diff --git a/litellm/constants.py b/litellm/constants.py index 36e578bd3232..711c8a413d35 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1157,6 +1157,7 @@ "openai.gpt-oss-120b-1:0", "anthropic.claude-haiku-4-5-20251001-v1:0", "anthropic.claude-sonnet-4-5-20250929-v1:0", + "anthropic.claude-fable-5", "anthropic.claude-opus-4-8", "anthropic.claude-opus-4-7", "anthropic.claude-opus-4-6-v1:0", diff --git a/litellm/integrations/datadog/datadog.py b/litellm/integrations/datadog/datadog.py index 79a9219a39cf..a79a852ff299 100644 --- a/litellm/integrations/datadog/datadog.py +++ b/litellm/integrations/datadog/datadog.py @@ -92,12 +92,22 @@ class DataDogLogger( # Class variables or attributes def __init__( self, + dd_api_key: Optional[str] = None, + dd_site: Optional[str] = None, + dd_agent_host: Optional[str] = None, + dd_agent_port: Optional[str] = None, **kwargs, ): """ Initializes the datadog logger, checks if the correct env variables are set - Required environment variables (Direct API): + Args: + dd_api_key: Datadog API key. Falls back to DD_API_KEY env var. + dd_site: Datadog site (e.g. "us5.datadoghq.com"). Falls back to DD_SITE env var. + dd_agent_host: Hostname or IP of DataDog agent. Falls back to LITELLM_DD_AGENT_HOST env var. + dd_agent_port: Port of DataDog agent (default: 10518). Falls back to LITELLM_DD_AGENT_PORT env var. + + Required environment variables (Direct API) when kwargs not provided: `DD_API_KEY` - your datadog api key `DD_SITE` - your datadog site, example = `"us5.datadoghq.com"` @@ -130,12 +140,19 @@ def __init__( ) # Configure DataDog endpoint (Agent or Direct API) - # Use LITELLM_DD_AGENT_HOST to avoid conflicts with ddtrace's DD_AGENT_HOST - dd_agent_host = os.getenv("LITELLM_DD_AGENT_HOST") - if dd_agent_host: - self._configure_dd_agent(dd_agent_host=dd_agent_host) + # Prefer explicit kwargs, then fall back to env vars + resolved_agent_host = dd_agent_host or os.getenv("LITELLM_DD_AGENT_HOST") + if resolved_agent_host: + self._configure_dd_agent( + dd_agent_host=resolved_agent_host, + dd_agent_port=dd_agent_port, + dd_api_key=dd_api_key, + ) else: - self._configure_dd_direct_api() + self._configure_dd_direct_api( + dd_api_key=dd_api_key, + dd_site=dd_site, + ) # Optional override for testing dd_base_url = get_datadog_base_url_from_env() @@ -172,34 +189,54 @@ def _get_datadog_params(self) -> Dict: ).model_dump() return dict_datadog_params - def _configure_dd_agent(self, dd_agent_host: str) -> None: + def _configure_dd_agent( + self, + dd_agent_host: str, + dd_agent_port: Optional[str] = None, + dd_api_key: Optional[str] = None, + ) -> None: """ Configure DataDog Agent for log forwarding Args: dd_agent_host: Hostname or IP of DataDog agent + dd_agent_port: Port of DataDog agent. Falls back to LITELLM_DD_AGENT_PORT env var (default: 10518). + dd_api_key: Datadog API key. Falls back to DD_API_KEY env var. Optional when using agent. """ - dd_agent_port = os.getenv( + resolved_port = dd_agent_port or os.getenv( "LITELLM_DD_AGENT_PORT", "10518" ) # default port for logs - self.intake_url = f"http://{dd_agent_host}:{dd_agent_port}/api/v2/logs" - self.DD_API_KEY = os.getenv("DD_API_KEY") # Optional when using agent + self.intake_url = f"http://{dd_agent_host}:{resolved_port}/api/v2/logs" + self.DD_API_KEY = dd_api_key or os.getenv( + "DD_API_KEY" + ) # Optional when using agent verbose_logger.debug(f"Datadog: Using DD Agent at {self.intake_url}") - def _configure_dd_direct_api(self) -> None: + def _configure_dd_direct_api( + self, + dd_api_key: Optional[str] = None, + dd_site: Optional[str] = None, + ) -> None: """ Configure direct DataDog API connection + Args: + dd_api_key: Datadog API key. Falls back to DD_API_KEY env var. + dd_site: Datadog site. Falls back to DD_SITE env var. + Raises: - Exception: If required environment variables are not set + Exception: If required credentials are not provided via args or env vars """ - if os.getenv("DD_API_KEY", None) is None: + resolved_api_key = dd_api_key or os.getenv("DD_API_KEY") + resolved_site = dd_site or os.getenv("DD_SITE") + + if resolved_api_key is None: raise Exception("DD_API_KEY is not set, set 'DD_API_KEY=<>") - if os.getenv("DD_SITE", None) is None: + if resolved_site is None: raise Exception("DD_SITE is not set in .env, set 'DD_SITE=<>") - self.DD_API_KEY = os.getenv("DD_API_KEY") - self.intake_url = f"https://http-intake.logs.{os.getenv('DD_SITE')}/api/v2/logs" + self.DD_API_KEY = resolved_api_key + self.intake_url = f"https://http-intake.logs.{resolved_site}/api/v2/logs" async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): """ diff --git a/litellm/integrations/datadog/datadog_team_handler.py b/litellm/integrations/datadog/datadog_team_handler.py new file mode 100644 index 000000000000..dda24c8e6a20 --- /dev/null +++ b/litellm/integrations/datadog/datadog_team_handler.py @@ -0,0 +1,117 @@ +""" +DataDog Team Handler + +Used to get the DataDogLogger for a given request. +Handles Key/Team Based Datadog Logging, following the same pattern as LangFuseHandler. +""" + +from typing import TYPE_CHECKING, Any, Dict, Optional, TypedDict + +from litellm._logging import verbose_logger +from litellm.litellm_core_utils.litellm_logging import StandardCallbackDynamicParams + +from .datadog import DataDogLogger + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import DynamicLoggingCache +else: + DynamicLoggingCache = Any + + +class DatadogLoggingConfig(TypedDict): + dd_api_key: Optional[str] + dd_site: Optional[str] + dd_agent_host: Optional[str] + dd_agent_port: Optional[str] + + +class DataDogHandler: + @staticmethod + def get_datadog_logger_for_request( + standard_callback_dynamic_params: StandardCallbackDynamicParams, + in_memory_dynamic_logger_cache: DynamicLoggingCache, + ) -> DataDogLogger: + """ + Get a team-scoped DataDogLogger for a given request. + + Resolves and caches per-team DataDogLogger instances using DynamicLoggingCache, + keyed by the team's DD credentials. Each unique set of credentials gets its own + logger instance with its own batch/flush loop. + + Note: This handler is only called when team-scoped DD credentials are present. + The global (env-var based) DataDogLogger is managed separately by + _init_custom_logger_compatible_class via _in_memory_loggers. + """ + _credentials = DataDogHandler.get_dynamic_datadog_logging_config( + standard_callback_dynamic_params=standard_callback_dynamic_params, + ) + credentials_dict = dict(_credentials) + + # check if datadog logger is already cached + temp_datadog_logger = in_memory_dynamic_logger_cache.get_cache( + credentials=credentials_dict, service_name="datadog" + ) + + # if not cached, create a new datadog logger and cache it + if temp_datadog_logger is None: + temp_datadog_logger = ( + DataDogHandler._create_datadog_logger_from_credentials( + credentials=credentials_dict, + in_memory_dynamic_logger_cache=in_memory_dynamic_logger_cache, + ) + ) + + return temp_datadog_logger + + @staticmethod + def _create_datadog_logger_from_credentials( + credentials: Dict, + in_memory_dynamic_logger_cache: DynamicLoggingCache, + ) -> DataDogLogger: + """ + Create a DataDogLogger from the credentials and cache it. + """ + datadog_logger = DataDogLogger( + dd_api_key=credentials.get("dd_api_key"), + dd_site=credentials.get("dd_site"), + dd_agent_host=credentials.get("dd_agent_host"), + dd_agent_port=credentials.get("dd_agent_port"), + ) + in_memory_dynamic_logger_cache.set_cache( + credentials=credentials, + service_name="datadog", + logging_obj=datadog_logger, + ) + verbose_logger.debug( + "Datadog: Created and cached new DataDogLogger for team-scoped credentials" + ) + return datadog_logger + + @staticmethod + def get_dynamic_datadog_logging_config( + standard_callback_dynamic_params: StandardCallbackDynamicParams, + ) -> DatadogLoggingConfig: + """ + Get the Datadog logging config for a given request from dynamic params. + """ + return DatadogLoggingConfig( + dd_api_key=standard_callback_dynamic_params.get("dd_api_key"), + dd_site=standard_callback_dynamic_params.get("dd_site"), + dd_agent_host=standard_callback_dynamic_params.get("dd_agent_host"), + dd_agent_port=standard_callback_dynamic_params.get("dd_agent_port"), + ) + + @staticmethod + def _dynamic_datadog_credentials_are_passed( + standard_callback_dynamic_params: StandardCallbackDynamicParams, + ) -> bool: + """ + Check if dynamic Datadog credentials are passed in standard_callback_dynamic_params. + """ + if ( + standard_callback_dynamic_params.get("dd_api_key") is not None + or standard_callback_dynamic_params.get("dd_site") is not None + or standard_callback_dynamic_params.get("dd_agent_host") is not None + ): + return True + return False diff --git a/litellm/litellm_core_utils/initialize_dynamic_callback_params.py b/litellm/litellm_core_utils/initialize_dynamic_callback_params.py index a89dae523168..949076aabf36 100644 --- a/litellm/litellm_core_utils/initialize_dynamic_callback_params.py +++ b/litellm/litellm_core_utils/initialize_dynamic_callback_params.py @@ -53,11 +53,19 @@ def validate_no_callback_env_reference( "braintrust_host", "slack_webhook_url", "lunary_public_key", + "dd_api_key", + "dd_site", + "dd_agent_host", + "dd_agent_port", ] _request_blocked_callback_params = { "gcs_bucket_name", "gcs_path_service_account", + "dd_api_key", + "dd_site", + "dd_agent_host", + "dd_agent_port", } diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index f20b66790c44..a76ca9546701 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -376,13 +376,14 @@ def __init__( List[Union[str, Callable, CustomLogger]] ] = dynamic_async_failure_callbacks - # Process dynamic callbacks - self.process_dynamic_callbacks() - ## DYNAMIC LANGFUSE / GCS / logging callback KEYS ## self.standard_callback_dynamic_params: StandardCallbackDynamicParams = ( self.initialize_standard_callback_dynamic_params(kwargs) ) + + # Process dynamic callbacks (after standard_callback_dynamic_params is initialized, + # so team-scoped credentials are available for callback initialization) + self.process_dynamic_callbacks() self.standard_built_in_tools_params: StandardBuiltInToolsParams = ( self.initialize_standard_built_in_tools_params(kwargs) ) @@ -477,8 +478,21 @@ def _process_dynamic_callback_list( isinstance(callback, str) and callback in litellm._known_custom_logger_compatible_callbacks ): + # For callbacks that support team-scoped credentials (e.g. datadog), + # pass only the relevant dynamic params as custom_logger_init_args. + _custom_logger_init_args: Optional[dict] = None + if callback == "datadog": + _custom_logger_init_args = { + k: v + for k, v in self.standard_callback_dynamic_params.items() + if k.startswith("dd_") + } + callback_class = _init_custom_logger_compatible_class( - callback, internal_usage_cache=None, llm_router=None # type: ignore + callback, # type: ignore[arg-type] + internal_usage_cache=None, + llm_router=None, # type: ignore + custom_logger_init_args=_custom_logger_init_args, ) if callback_class is not None: processed_list.append(callback_class) @@ -3890,6 +3904,24 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 _in_memory_loggers.append(_prometheus_logger) return _prometheus_logger # type: ignore elif logging_integration == "datadog": + # Check if team-scoped credentials are provided + _dd_api_key = custom_logger_init_args.get("dd_api_key") + _dd_site = custom_logger_init_args.get("dd_site") + _dd_agent_host = custom_logger_init_args.get("dd_agent_host") + _dd_agent_port = custom_logger_init_args.get("dd_agent_port") + + if _dd_api_key or _dd_site or _dd_agent_host: + # Team-scoped credentials: use DynamicLoggingCache for per-credential isolation + from litellm.integrations.datadog.datadog_team_handler import ( + DataDogHandler, + ) + + return DataDogHandler.get_datadog_logger_for_request( + standard_callback_dynamic_params=custom_logger_init_args, # type: ignore + in_memory_dynamic_logger_cache=in_memory_dynamic_logger_cache, + ) + + # Global (env-var based): reuse cached instance for callback in _in_memory_loggers: if isinstance(callback, DataDogLogger): return callback # type: ignore diff --git a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py index 8da66d4600d5..413ddb71bf80 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py +++ b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py @@ -6,6 +6,7 @@ import litellm from litellm.constants import OPENAI_FILE_SEARCH_COST_PER_1K_CALLS +from litellm.litellm_core_utils.llm_cost_calc.utils import _get_web_search_requests from litellm.types.llms.openai import ( FileSearchTool, ResponsesAPIResponse, @@ -339,8 +340,7 @@ def response_object_includes_web_search_call( # and _handle_web_search_cost() is never called. if ( hasattr(usage, "server_tool_use") - and usage.server_tool_use is not None - and usage.server_tool_use.web_search_requests is not None + and _get_web_search_requests(usage.server_tool_use) is not None ): return True return False @@ -352,8 +352,7 @@ def response_object_includes_web_search_call( elif usage is not None: if ( hasattr(usage, "server_tool_use") - and usage.server_tool_use is not None - and usage.server_tool_use.web_search_requests is not None + and _get_web_search_requests(usage.server_tool_use) is not None ): return True elif ( diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index f39c942f90f0..93049adf75aa 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -1,7 +1,7 @@ # What is this? ## Helper utilities for cost_per_token() -from typing import Literal, Optional, Tuple, TypedDict, cast +from typing import Any, Literal, Optional, Tuple, TypedDict, cast import litellm from litellm._logging import verbose_logger @@ -42,6 +42,26 @@ def _get_token_detail_value(details: object, key: str) -> Optional[int]: return value if isinstance(value, int) else None +def _get_web_search_requests(server_tool_use: Any) -> Optional[int]: + """ + Tolerantly read ``web_search_requests`` from a ``server_tool_use`` value + that may be ``None``, a ``dict``, a ``ServerToolUse`` pydantic instance, + or any other object supporting attribute access. + + Returns ``None`` when the value cannot be resolved — callers can + distinguish "absent" from "zero" using ``is None``. + + See https://github.com/BerriAI/litellm/issues/26153 — ``stream_chunk_builder`` + historically left this as a plain ``dict``, which broke direct attribute + access in cost calculation. + """ + if server_tool_use is None: + return None + if isinstance(server_tool_use, dict): + return server_tool_use.get("web_search_requests") + return getattr(server_tool_use, "web_search_requests", None) + + def _is_above_128k(tokens: float) -> bool: if tokens > 128000: return True diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index fe7c62c38425..c2a17ae8dcc1 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -588,7 +588,18 @@ def _calculate_usage_per_chunk( hasattr(usage_chunk, "server_tool_use") and usage_chunk.server_tool_use is not None ): - server_tool_use = usage_chunk.server_tool_use + # Coerce dict to ServerToolUse so downstream cost-calc code + # (which accesses .web_search_requests as an attribute) + # doesn't raise AttributeError. Some providers / streaming + # paths leave server_tool_use as a plain dict on the chunk. + if isinstance(usage_chunk.server_tool_use, dict): + server_tool_use = ServerToolUse(**usage_chunk.server_tool_use) + elif isinstance(usage_chunk.server_tool_use, ServerToolUse): + server_tool_use = usage_chunk.server_tool_use + else: + server_tool_use = ServerToolUse.model_validate( + usage_chunk.server_tool_use + ) if ( usage_chunk_dict["prompt_tokens_details"] is not None and getattr( diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 7949e150c23b..cc30db0ebad0 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -1455,10 +1455,15 @@ def map_openai_params( # noqa: PLR0915 _value = self._map_stop_sequences(value) if _value is not None: optional_params["stop_sequences"] = _value - elif param == "temperature": - optional_params["temperature"] = value - elif param == "top_p": - optional_params["top_p"] = value + elif param == "temperature" or param == "top_p": + AnthropicConfig._apply_sampling_param( + optional_params=optional_params, + model=model, + param=param, + value=value, + drop_params=drop_params, + output_key=param, + ) elif param == "response_format" and isinstance(value, dict): if any( substring in model @@ -1967,6 +1972,20 @@ def transform_request( optional_params.pop("is_vertex_request", None) optional_params.pop("client_metadata", None) + # ``top_k`` is a provider-specific kwarg that bypasses + # ``map_openai_params``; gate it here, the single boundary shared by + # the direct Anthropic, Bedrock invoke, Vertex, and Azure paths. + top_k = optional_params.pop("top_k", None) + if top_k is not None: + AnthropicConfig._apply_sampling_param( + optional_params=optional_params, + model=model, + param="top_k", + value=top_k, + drop_params=litellm_params.get("drop_params") is True, + output_key="top_k", + ) + data = { "model": model, "messages": anthropic_messages, diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 3f002d73cbca..5741513903c4 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -272,23 +272,68 @@ def _is_claude_4_7_model(model: str) -> bool: ) @staticmethod - def _supports_model_capability(model: str, key: str) -> bool: - """Check a boolean capability ``key`` in the model map. + def _supports_sampling_params(model: str) -> bool: + """Claude 4.7+ (Opus 4.7/4.8, Fable 5) removed sampling params: the API + rejects ``top_p``, ``top_k``, and any ``temperature`` other than 1 with + a 400 ("`temperature` is deprecated for this model"). + + Driven by the ``supports_sampling_params`` flag in the model map; the + name check remains only as a fallback for provider-routed ids whose + map entries predate the flag.""" + flag = AnthropicModelInfo._get_model_capability( + model, "supports_sampling_params" + ) + if flag is not None: + return flag + model_lower = model.lower() + return not any( + v in model_lower + for v in ( + "fable", + "opus-4-7", + "opus_4_7", + "opus-4.7", + "opus_4.7", + "opus-4-8", + "opus_4_8", + "opus-4.8", + "opus_4.8", + ) + ) - Strips bedrock/vertex prefixes so a provider-routed Claude still - resolves to the Anthropic model-map entry. - """ - from litellm.utils import _supports_factory + @staticmethod + def _apply_sampling_param( + optional_params: dict, + model: str, + param: str, + value: Any, + drop_params: bool, + output_key: str, + ) -> None: + """Forward ``temperature``/``top_p``/``top_k`` to + ``optional_params[output_key]`` unless the model removed sampling + params, in which case drop the param (with drop_params) or raise a + clean client-side 400.""" + if AnthropicModelInfo._supports_sampling_params(model) or ( + param == "temperature" and value == 1 + ): + optional_params[output_key] = value + elif not (litellm.drop_params or drop_params): + supported_hint = ( + "Only temperature=1 is supported. " if param == "temperature" else "" + ) + raise litellm.utils.UnsupportedParamsError( + message=( + f"{model} does not support {param}={value}. {supported_hint}" + "To drop unsupported params, set `litellm.drop_params = True`." + ), + status_code=400, + ) - try: - if _supports_factory( - model=model, - custom_llm_provider="anthropic", - key=key, - ): - return True - except Exception: - pass + @staticmethod + def _model_map_lookup_candidates(model: str) -> List[str]: + """Model-map keys to try for ``model``, stripping bedrock/vertex + prefixes so a provider-routed Claude still resolves to its entry.""" candidates = [model] for prefix in ( "bedrock/converse/", @@ -307,15 +352,40 @@ def _supports_model_capability(model: str, key: str) -> bool: candidates.append(f"bedrock/{base}") except Exception: pass + return candidates + + @staticmethod + def _get_model_capability(model: str, key: str) -> Optional[bool]: + """Read boolean capability ``key`` from the model map, or None when + no entry declares it.""" try: - for cand in candidates: - if cand in litellm.model_cost and ( - litellm.model_cost[cand].get(key) is True - ): - return True + for cand in AnthropicModelInfo._model_map_lookup_candidates(model): + value = litellm.model_cost.get(cand, {}).get(key) + if isinstance(value, bool): + return value except Exception: pass - return False + return None + + @staticmethod + def _supports_model_capability(model: str, key: str) -> bool: + """Check a boolean capability ``key`` in the model map. + + Strips bedrock/vertex prefixes so a provider-routed Claude still + resolves to the Anthropic model-map entry. + """ + from litellm.utils import _supports_factory + + try: + if _supports_factory( + model=model, + custom_llm_provider="anthropic", + key=key, + ): + return True + except Exception: + pass + return AnthropicModelInfo._get_model_capability(model, key) is True @staticmethod def _is_adaptive_thinking_model(model: str) -> bool: diff --git a/litellm/llms/anthropic/cost_calculation.py b/litellm/llms/anthropic/cost_calculation.py index 3882d8f978c2..6a031498dae2 100644 --- a/litellm/llms/anthropic/cost_calculation.py +++ b/litellm/llms/anthropic/cost_calculation.py @@ -7,6 +7,7 @@ from litellm.litellm_core_utils.llm_cost_calc.utils import ( _get_token_base_cost, + _get_web_search_requests, _parse_prompt_tokens_details, calculate_cache_writing_cost, generic_cost_per_token, @@ -110,11 +111,12 @@ def get_cost_for_anthropic_web_search( if model_info is None: return 0.0 - if ( - usage is None - or usage.server_tool_use is None - or usage.server_tool_use.web_search_requests is None - ): + if usage is None: + return 0.0 + web_search_requests = _get_web_search_requests( + getattr(usage, "server_tool_use", None) + ) + if web_search_requests is None: return 0.0 ## Get the cost per web search request @@ -128,5 +130,5 @@ def get_cost_for_anthropic_web_search( return 0.0 ## Calculate the total cost - total_cost = cost_per_web_search_request * usage.server_tool_use.web_search_requests + total_cost = cost_per_web_search_request * web_search_requests return total_cost diff --git a/litellm/llms/base_llm/responses/transformation.py b/litellm/llms/base_llm/responses/transformation.py index 853eb2827580..407d5ad8146a 100644 --- a/litellm/llms/base_llm/responses/transformation.py +++ b/litellm/llms/base_llm/responses/transformation.py @@ -62,6 +62,26 @@ def supports_native_file_search(self) -> bool: """ return False + def sign_request( + self, + headers: dict, + optional_params: dict, + request_data: dict, + api_base: str, + api_key: Optional[str] = None, + model: Optional[str] = None, + stream: Optional[bool] = None, + fake_stream: Optional[bool] = None, + ) -> Tuple[dict, Optional[bytes]]: + """Sign the request after the body is finalized. + + Default is a no-op (returns headers unchanged, no signed body). Providers + whose endpoint requires request signing (e.g. Bedrock Mantle SigV4) + override this and return the signed body bytes so the handler sends those + exact bytes. + """ + return headers, None + @abstractmethod def get_supported_openai_params(self, model: str) -> list: pass diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 90dfa13e9385..e56cd7c617b9 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -920,10 +920,15 @@ def map_openai_params( continue value = [value] optional_params["stopSequences"] = value - if param == "temperature": - optional_params["temperature"] = value - if param == "top_p": - optional_params["topP"] = value + if param == "temperature" or param == "top_p": + AnthropicConfig._apply_sampling_param( + optional_params=optional_params, + model=model, + param=param, + value=value, + drop_params=drop_params, + output_key="topP" if param == "top_p" else param, + ) if param == "tools" and isinstance(value, list): self._apply_tool_call_transformation( tools=cast(List[OpenAIChatCompletionToolParam], value), @@ -1221,7 +1226,9 @@ def _transform_inference_params(self, inference_params: dict) -> InferenceConfig inference_params["topK"] = inference_params.pop("top_k") return InferenceConfig(**inference_params) - def _handle_top_k_value(self, model: str, inference_params: dict) -> dict: + def _handle_top_k_value( + self, model: str, inference_params: dict, drop_params: bool = False + ) -> dict: base_model = BedrockModelInfo.get_base_model(model) val_top_k = None @@ -1230,16 +1237,25 @@ def _handle_top_k_value(self, model: str, inference_params: dict) -> dict: elif "top_k" in inference_params: val_top_k = inference_params.pop("top_k") - if val_top_k: + if val_top_k is not None: if base_model.startswith("anthropic"): - return {"top_k": val_top_k} + top_k_params: dict = {} + AnthropicConfig._apply_sampling_param( + optional_params=top_k_params, + model=model, + param="top_k", + value=val_top_k, + drop_params=drop_params, + output_key="top_k", + ) + return top_k_params if base_model.startswith("amazon.nova"): return {"inferenceConfig": {"topK": val_top_k}} return {} def _prepare_request_params( - self, optional_params: dict, model: str + self, optional_params: dict, model: str, drop_params: bool = False ) -> Tuple[dict, dict, dict, Optional[OutputConfigBlock]]: """Prepare and separate request parameters.""" # Consume the internal ``_output_config_normalized`` marker set by @@ -1338,7 +1354,7 @@ def _prepare_request_params( # Only set the topK value in for models that support it additional_request_params.update( - self._handle_top_k_value(model, inference_params) + self._handle_top_k_value(model, inference_params, drop_params) ) # Filter out internal/MCP-related parameters that shouldn't be sent to the API @@ -1572,6 +1588,7 @@ def _transform_request_helper( optional_params: dict, messages: Optional[List[AllMessageValues]] = None, headers: Optional[dict] = None, + drop_params: bool = False, ) -> CommonRequestObject: ## VALIDATE REQUEST """ @@ -1618,7 +1635,7 @@ def _transform_request_helper( additional_request_params, request_metadata, output_config, - ) = self._prepare_request_params(optional_params, model) + ) = self._prepare_request_params(optional_params, model, drop_params) original_tools = inference_params.pop("tools", []) @@ -1699,6 +1716,7 @@ async def _async_transform_request( optional_params=optional_params, messages=messages, headers=headers, + drop_params=litellm_params.get("drop_params") is True, ) bedrock_messages = ( @@ -1756,6 +1774,7 @@ def _transform_request( optional_params=optional_params, messages=messages, headers=headers, + drop_params=litellm_params.get("drop_params") is True, ) ## TRANSFORMATION ## diff --git a/litellm/llms/bedrock_mantle/responses/transformation.py b/litellm/llms/bedrock_mantle/responses/transformation.py index b63fd0ecdb1e..df219091074c 100644 --- a/litellm/llms/bedrock_mantle/responses/transformation.py +++ b/litellm/llms/bedrock_mantle/responses/transformation.py @@ -4,14 +4,26 @@ gpt-5.5 / gpt-5.4 on Mantle are exposed ONLY on the `/openai/v1/responses` path (not the standard `/v1/responses`). Payloads and SSE follow the OpenAI Responses spec, so this config inherits OpenAIResponsesAPIConfig and overrides -only the endpoint URL and Bearer authentication. +only the endpoint URL and authentication. -Auth: AWS Bedrock API key as Bearer token (BEDROCK_MANTLE_API_KEY or the -standard AWS_BEARER_TOKEN_BEDROCK), NOT SigV4. +Auth: Bearer token (BEDROCK_MANTLE_API_KEY or the standard +AWS_BEARER_TOKEN_BEDROCK, or litellm_params.api_key) when present; otherwise +AWS SigV4 (service name "bedrock") using the standard credential chain (IAM +role / access key / profile / web identity), signed via the shared +BaseAWSLLM._sign_request after the request body is finalized. """ -from typing import Optional +import re +from typing import Optional, Tuple +from botocore.exceptions import ( + CredentialRetrievalError, + NoCredentialsError, + PartialCredentialsError, + ProfileNotFound, +) + +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.secret_managers.main import get_secret_str from litellm.types.router import GenericLiteLLMParams @@ -29,22 +41,44 @@ "/v1", ) +# Standard Mantle host: https://bedrock-mantle..api.aws (group 1 = region). +_MANTLE_HOST_RE = re.compile( + r"^https?://bedrock-mantle\.([^/.]+)\.api\.aws", re.IGNORECASE +) + class BedrockMantleResponsesAPIConfig(OpenAIResponsesAPIConfig): + def __init__(self, aws_signer: Optional[BaseAWSLLM] = None): + super().__init__() + self._aws_signer = aws_signer or BaseAWSLLM() + @property def custom_llm_provider(self) -> LlmProviders: return LlmProviders.BEDROCK_MANTLE + @staticmethod + def _resolve_region(params: dict) -> str: + region = params.get("aws_region_name") + if region: + return region + base = params.get("api_base") or get_secret_str("BEDROCK_MANTLE_API_BASE") + if base: + match = _MANTLE_HOST_RE.match(base.rstrip("/")) + if match: + return match.group(1) + return ( + get_secret_str("BEDROCK_MANTLE_REGION") + or get_secret_str("AWS_REGION_NAME") + or get_secret_str("AWS_REGION") + or BEDROCK_MANTLE_DEFAULT_REGION + ) + def get_complete_url( self, api_base: Optional[str], litellm_params: dict, ) -> str: - region = ( - get_secret_str("BEDROCK_MANTLE_REGION") - or get_secret_str("AWS_REGION") - or BEDROCK_MANTLE_DEFAULT_REGION - ) + region = self._resolve_region({**litellm_params, "api_base": api_base}) base = ( api_base or get_secret_str("BEDROCK_MANTLE_API_BASE") @@ -55,6 +89,11 @@ def get_complete_url( if base.endswith(suffix): base = base[: -len(suffix)] break + # For the standard Mantle host (including the default-region base that + # responses/main.py auto-injects into litellm_params.api_base), pin to the + # single resolved region so aws_region_name wins; preserve custom proxy hosts. + if _MANTLE_HOST_RE.match(base): + base = f"https://bedrock-mantle.{region}.api.aws" return f"{base}/openai/v1/responses" def validate_environment( @@ -66,12 +105,8 @@ def validate_environment( or get_secret_str("BEDROCK_MANTLE_API_KEY") or get_secret_str("AWS_BEARER_TOKEN_BEDROCK") ) - if not api_key: - raise ValueError( - "Bedrock Mantle API key is required. Set BEDROCK_MANTLE_API_KEY " - "(or AWS_BEARER_TOKEN_BEDROCK) or pass api_key." - ) - headers["Authorization"] = f"Bearer {api_key}" + if api_key: + headers["Authorization"] = f"Bearer {api_key}" return headers def supports_native_file_search(self) -> bool: @@ -79,3 +114,58 @@ def supports_native_file_search(self) -> bool: def supports_native_websocket(self) -> bool: return False + + def sign_request( + self, + headers: dict, + optional_params: dict, + request_data: dict, + api_base: str, + api_key: Optional[str] = None, + model: Optional[str] = None, + stream: Optional[bool] = None, + fake_stream: Optional[bool] = None, + ) -> Tuple[dict, Optional[bytes]]: + bearer = ( + api_key + or get_secret_str("BEDROCK_MANTLE_API_KEY") + or get_secret_str("AWS_BEARER_TOKEN_BEDROCK") + ) + if not bearer: + # SigV4 path. Pin the credential-scope region to the region of the actual + # signing URL (api_base, already region-resolved by get_complete_url) so the + # SigV4 scope and the URL host can never disagree. Resolve from api_base first, + # then fall back to the regular precedence. Also drop any caller Authorization + # so _sign_request's restore-original-Authorization step cannot override the + # SigV4 header. + optional_params = { + **optional_params, + "aws_region_name": self._resolve_region( + {**optional_params, "api_base": api_base} + ), + } + headers = {k: v for k, v in headers.items() if k.lower() != "authorization"} + try: + return self._aws_signer._sign_request( + service_name="bedrock", + headers=headers, + optional_params=optional_params, + request_data=request_data, + api_base=api_base, + api_key=bearer, + model=model, + stream=stream, + fake_stream=fake_stream, + ) + except ( + NoCredentialsError, + PartialCredentialsError, + ProfileNotFound, + CredentialRetrievalError, + ) as e: + raise ValueError( + "Bedrock Mantle auth failed: no Bearer token and no usable AWS " + "credentials. Set BEDROCK_MANTLE_API_KEY (or AWS_BEARER_TOKEN_BEDROCK) " + "or pass api_key for Bearer auth, or provide AWS credentials " + "(IAM role / access key / profile / web identity) for SigV4." + ) from e diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 31c772510bad..25424feaeb4c 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -2318,6 +2318,31 @@ def response_api_handler( # but never included in the outbound provider payload. request_context["litellm_params"] = dict(litellm_params) + is_stream_request = bool(stream) + if is_stream_request and fake_stream is True: + stream, data = self._prepare_fake_stream_request( + stream=stream, + data=data, + fake_stream=fake_stream, + ) + + # Sign after the body is final (post-transform/normalize/extra_body and post + # fake-stream prep) so signed bytes match what we send. No-op for providers + # that inherit the default sign_request. + headers, signed_body = responses_api_provider_config.sign_request( + headers=headers, + optional_params=dict(litellm_params), + request_data=data, + api_base=api_base, + api_key=litellm_params.api_key, + model=model, + stream=stream, + fake_stream=fake_stream, + ) + body_kwargs: Dict[str, Any] = ( + {"data": signed_body} if signed_body is not None else {"json": data} + ) + ## LOGGING logging_obj.pre_call( input=input, @@ -2330,22 +2355,14 @@ def response_api_handler( ) try: - if stream: - # For streaming, use stream=True in the request - if fake_stream is True: - stream, data = self._prepare_fake_stream_request( - stream=stream, - data=data, - fake_stream=fake_stream, - ) - + if is_stream_request: response = sync_httpx_client.post( url=api_base, headers=headers, - json=data, timeout=timeout or float(response_api_optional_request_params.get("timeout", 0)), stream=stream, + **body_kwargs, ) if fake_stream is True: return MockResponsesAPIStreamingIterator( @@ -2370,13 +2387,12 @@ def response_api_handler( call_type=CallTypes.responses.value, ) else: - # For non-streaming requests response = sync_httpx_client.post( url=api_base, headers=headers, - json=data, timeout=timeout or float(response_api_optional_request_params.get("timeout", 0)), + **body_kwargs, ) except Exception as e: raise self._handle_error( @@ -2464,6 +2480,28 @@ async def async_response_api_handler( # but never included in the outbound provider payload. request_context["litellm_params"] = dict(litellm_params) + is_stream_request = bool(stream) + if is_stream_request and fake_stream is True: + stream, data = self._prepare_fake_stream_request( + stream=stream, + data=data, + fake_stream=fake_stream, + ) + + headers, signed_body = responses_api_provider_config.sign_request( + headers=headers, + optional_params=dict(litellm_params), + request_data=data, + api_base=api_base, + api_key=litellm_params.api_key, + model=model, + stream=stream, + fake_stream=fake_stream, + ) + body_kwargs: Dict[str, Any] = ( + {"data": signed_body} if signed_body is not None else {"json": data} + ) + ## LOGGING logging_obj.pre_call( input=input, @@ -2476,22 +2514,14 @@ async def async_response_api_handler( ) try: - if stream: - # For streaming, we need to use stream=True in the request - if fake_stream is True: - stream, data = self._prepare_fake_stream_request( - stream=stream, - data=data, - fake_stream=fake_stream, - ) - + if is_stream_request: response = await async_httpx_client.post( url=api_base, headers=headers, - json=data, timeout=timeout or float(response_api_optional_request_params.get("timeout", 0)), stream=stream, + **body_kwargs, ) if fake_stream is True: @@ -2518,13 +2548,12 @@ async def async_response_api_handler( call_type=CallTypes.responses.value, ) else: - # For non-streaming, proceed as before response = await async_httpx_client.post( url=api_base, headers=headers, - json=data, timeout=timeout or float(response_api_optional_request_params.get("timeout", 0)), + **body_kwargs, ) except Exception as e: @@ -4005,6 +4034,18 @@ def compact_response_api_handler( ) data = BaseResponsesAPIConfig.normalize_responses_api_request_dict(data) + headers, signed_body = responses_api_provider_config.sign_request( + headers=headers, + optional_params=dict(litellm_params), + request_data=data, + api_base=url, + api_key=litellm_params.api_key, + model=model, + ) + body_kwargs: Dict[str, Any] = ( + {"data": signed_body} if signed_body is not None else {"json": data} + ) + ## LOGGING logging_obj.pre_call( input=input, @@ -4018,7 +4059,7 @@ def compact_response_api_handler( try: response = sync_httpx_client.post( - url=url, headers=headers, json=data, timeout=timeout + url=url, headers=headers, timeout=timeout, **body_kwargs ) except Exception as e: @@ -4088,6 +4129,18 @@ async def async_compact_response_api_handler( ) data = BaseResponsesAPIConfig.normalize_responses_api_request_dict(data) + headers, signed_body = responses_api_provider_config.sign_request( + headers=headers, + optional_params=dict(litellm_params), + request_data=data, + api_base=url, + api_key=litellm_params.api_key, + model=model, + ) + body_kwargs: Dict[str, Any] = ( + {"data": signed_body} if signed_body is not None else {"json": data} + ) + ## LOGGING logging_obj.pre_call( input=input, @@ -4101,7 +4154,7 @@ async def async_compact_response_api_handler( try: response = await async_httpx_client.post( - url=url, headers=headers, json=data, timeout=timeout + url=url, headers=headers, timeout=timeout, **body_kwargs ) except Exception as e: diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 397f96fdb1e8..757aacf1cafa 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -1156,6 +1156,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1202,6 +1203,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1233,6 +1235,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1264,6 +1267,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1295,6 +1299,139 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh" + }, + "anthropic.claude-fable-5": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 1e-06, + "input_cost_per_token": 1e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh" + }, + "global.anthropic.claude-fable-5": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 1e-06, + "input_cost_per_token": 1e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh" + }, + "us.anthropic.claude-fable-5": { + "cache_creation_input_token_cost": 1.375e-05, + "cache_creation_input_token_cost_above_1hr": 2.2e-05, + "cache_read_input_token_cost": 1.1e-06, + "input_cost_per_token": 1.1e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh" + }, + "eu.anthropic.claude-fable-5": { + "cache_creation_input_token_cost": 1.375e-05, + "cache_creation_input_token_cost_above_1hr": 2.2e-05, + "cache_read_input_token_cost": 1.1e-06, + "input_cost_per_token": 1.1e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1327,6 +1464,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1359,6 +1497,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1391,6 +1530,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1423,6 +1563,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1455,6 +1596,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1485,6 +1627,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -2208,6 +2351,37 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true + }, + "azure_ai/claude-fable-5": { + "input_cost_per_token": 1e-05, + "output_cost_per_token": 5e-05, + "litellm_provider": "azure_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 1e-06, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -2237,6 +2411,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -10133,6 +10308,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -10167,6 +10343,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -10177,6 +10354,40 @@ }, "supports_output_config": true }, + "claude-fable-5": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 1e-06, + "input_cost_per_token": 1e-05, + "litellm_provider": "anthropic", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true + }, "claude-opus-4-8": { "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, @@ -10201,6 +10412,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -33967,6 +34179,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -33995,6 +34208,67 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true + }, + "vertex_ai/claude-fable-5": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 1e-06, + "input_cost_per_token": 1e-05, + "litellm_provider": "vertex_ai-anthropic_models", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true + }, + "vertex_ai/claude-fable-5@default": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 1e-06, + "input_cost_per_token": 1e-05, + "litellm_provider": "vertex_ai-anthropic_models", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -34024,6 +34298,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -34053,6 +34328,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 863e6acd41e6..914362d9b77c 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -63,9 +63,10 @@ def _is_mcp_passthrough_cold_start( spec-compliant WWW-Authenticate challenge instead of surfacing a generic admission error. - Uses "all" semantics (mirrors :meth:`MCPRequestHandler._target_servers_use_oauth2`): - one non-passthrough target in a co-targeted set must not flip the bypass - open for the others. Fails closed when any target cannot be resolved.""" + Uses "all" semantics (mirrors + :meth:`MCPRequestHandler._target_servers_delegate_auth_to_upstream`): one + non-passthrough target in a co-targeted set must not flip the bypass open + for the others. Fails closed when any target cannot be resolved.""" if not mcp_servers: return False from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( @@ -210,101 +211,64 @@ async def mock_body(): # Only OAuth metadata routes registered under /.well-known/ are public. if request_route.startswith("/.well-known/"): validated_user_api_key_auth = UserAPIKeyAuth() - elif ( - not litellm_api_key - and MCPRequestHandler._target_servers_delegate_auth_to_upstream( # noqa: E501 - path=request_route, - mcp_servers=mcp_servers, - client_ip=IPAddressUtils.get_mcp_client_ip(request), - ) - ): - # Operator opted this oauth2 server into upstream-delegated auth - # (PKCE passthrough): skip LiteLLM API-key/SSO entirely so the - # client authenticates directly with the upstream MCP server. - # Fires ONLY when neither x-litellm-api-key nor Authorization is - # present. If any LiteLLM key is supplied (primary or secondary - # header), we fall through so user_id is resolved, spend/rate - # limiting apply, and any stored OAuth token can be retrieved - # and forwarded upstream. Gated by - # _target_servers_delegate_auth_to_upstream, which only returns - # True when EVERY target is auth_type=oauth2 AND has the - # delegate_auth_to_upstream flag set — fails closed otherwise. - validated_user_api_key_auth = UserAPIKeyAuth() elif has_explicit_litellm_key: - # Explicit x-litellm-api-key provided - always validate normally + # An explicit x-litellm-api-key is always a LiteLLM credential, even + # for a delegated server, so validate it: identity / spend / rate + # limits resolve and any stored upstream token can be forwarded. validated_user_api_key_auth = await user_api_key_auth( api_key=litellm_api_key, request=request ) + elif MCPRequestHandler._target_servers_delegate_auth_to_upstream( + path=request_route, + mcp_servers=mcp_servers, + client_ip=IPAddressUtils.get_mcp_client_ip(request), + ): + # Operator opted this oauth2 server into upstream-delegated auth: the + # client authenticates directly with the upstream MCP server, so any + # Authorization bearer is an upstream token, never a LiteLLM key. Skip + # LiteLLM validation entirely — covering both the no-credential + # discovery request and the authenticated call carrying the upstream + # bearer — so a tool call that succeeds never carries a phantom 401 + # auth span; the bearer is forwarded upstream unchanged. Gated by + # _target_servers_delegate_auth_to_upstream, which returns True only + # when EVERY target is auth_type=oauth2 with delegate_auth_to_upstream + # set; fails closed otherwise. + validated_user_api_key_auth = UserAPIKeyAuth() elif oauth2_headers: - # No x-litellm-api-key, but Authorization header present. - # Could be a LiteLLM key (backward compat) OR an opaque OAuth2 token - # the operator wants forwarded to an upstream OAuth2-mode MCP server. - # Try LiteLLM auth first; on auth failure, only fall back to anonymous - # passthrough when the request actually targets a server whose operator - # configured ``auth_type=oauth2``. For any other server (api_key, - # bearer_token, basic, etc.), a failed LiteLLM auth is a real failure - # and must propagate — otherwise an attacker can exchange any garbage - # bearer for an anonymous session. + # Authorization on a non-delegated server: the bearer must be a real + # LiteLLM credential, so a failed validation is a genuine 401/403 and + # propagates. The sole anonymous fallback is the auth_type=none + # pass-through cold-start (RFC 9728 discovery return), gated on a 401 + # so a recognized-but-forbidden key still fails closed. + client_ip = IPAddressUtils.get_mcp_client_ip(request) try: validated_user_api_key_auth = await user_api_key_auth( api_key=litellm_api_key, request=request ) except (HTTPException, ProxyException) as e: - # HTTPException.status_code is int; ProxyException.code is - # normalized to str in its __init__ but can be ``"None"`` or any - # non-numeric string when the caller didn't supply a numeric - # code, so we compare against both int and str forms rather - # than coercing (``int("None")`` would raise ValueError and - # rewrite the auth error as a 500). + # ProxyException.code is normalized to str (possibly "None"), so + # compare both int and str forms rather than coercing. status = e.status_code if isinstance(e, HTTPException) else e.code - is_auth_error = status in (401, 403, "401", "403") is_unauthenticated = status in (401, "401") - client_ip = IPAddressUtils.get_mcp_client_ip(request) - if is_auth_error and MCPRequestHandler._target_servers_use_oauth2( - path=request_route, - mcp_servers=mcp_servers, - client_ip=client_ip, + mcp_servers_from_path = _parse_mcp_server_names_from_path( + request_route, mcp_servers + ) + if ( + is_unauthenticated + and mcp_servers_from_path is not None + and not _has_client_supplied_mcp_auth( + mcp_auth_header, + mcp_server_auth_headers, + ) + and _is_mcp_passthrough_cold_start( + mcp_servers_from_path, client_ip=client_ip + ) ): verbose_logger.debug( - "MCP OAuth2: target server is OAuth2-mode, treating " - "Authorization as upstream OAuth2 token passthrough" + "MCP pass-through return: forwarding Authorization as " + "upstream OAuth token for delegated auth" ) validated_user_api_key_auth = UserAPIKeyAuth() - elif is_unauthenticated: - # Pass-through cold-start return: per RFC 9728 / MCP - # Authorization spec the client completes upstream OAuth - # discovery and returns with ``Authorization: Bearer - # ``. For ``auth_type=none`` passthrough - # servers that bearer is not a LiteLLM key (auth above - # failed) but is meant to be forwarded upstream - # unchanged. Fall back to anonymous admission so the - # caller is not rejected for following the discovery - # flow without also setting ``x-litellm-api-key``. - # Only trigger on 401 (token unrecognized); a 403 means - # the key WAS recognized but is forbidden (e.g. over - # budget / rate limited) and must propagate so those - # controls are not bypassed via anonymous admission. - mcp_servers_from_path = _parse_mcp_server_names_from_path( - request_route, mcp_servers - ) - if ( - mcp_servers_from_path is not None - and not _has_client_supplied_mcp_auth( - mcp_auth_header, - mcp_server_auth_headers, - ) - and _is_mcp_passthrough_cold_start( - mcp_servers_from_path, client_ip=client_ip - ) - ): - verbose_logger.debug( - "MCP pass-through return: target server is " - "passthrough, treating Authorization as " - "upstream OAuth token for delegated auth" - ) - validated_user_api_key_auth = UserAPIKeyAuth() - else: - raise else: raise else: @@ -408,45 +372,6 @@ def _extract_target_server_names_from_path(path: str) -> List[str]: return [single_server_match.group(1)] return [servers_and_path] - @staticmethod - def _target_servers_use_oauth2( - path: str, mcp_servers: Optional[List[str]], client_ip: Optional[str] - ) -> bool: - """ - True only when EVERY MCP server the request targets is configured for - ``auth_type == oauth2``. If any target is non-OAuth2 — or if the target - cannot be resolved at all — return False so the caller fails closed. - - Used to gate the "treat Authorization as opaque OAuth2 token" fallback - in :meth:`process_mcp_request` so a failed LiteLLM-auth cannot be - exchanged for an anonymous session against a non-OAuth2 server. - """ - # Inline imports avoid a circular dependency: mcp_server_manager imports - # from this module. - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from litellm.types.mcp import MCPAuth - - # Resolve the same target list downstream routing will use. For - # ``/mcp/...`` routes, ``extract_mcp_auth_context`` overrides the - # ``x-mcp-servers`` header with path-derived names, so we must mirror - # that here — otherwise a caller could set the header to a permissive - # server while the path targets a stricter one (header/path TOCTOU). - target_names = MCPRequestHandler._resolve_target_server_names( - path=path, mcp_servers_header=mcp_servers - ) - if not target_names: - return False - - for name in target_names: - server = global_mcp_server_manager.get_mcp_server_by_name( - name, client_ip=client_ip - ) - if server is None or server.auth_type != MCPAuth.oauth2: - return False - return True - @staticmethod def _target_servers_delegate_auth_to_upstream( path: str, mcp_servers: Optional[List[str]], client_ip: Optional[str] @@ -468,8 +393,8 @@ def _target_servers_delegate_auth_to_upstream( ) from litellm.types.mcp import MCPAuth - # See _target_servers_use_oauth2: must mirror the downstream - # header-vs-path override or an attacker could set + # Must mirror the downstream header-vs-path override + # (``extract_mcp_auth_context``) or an attacker could set # ``x-mcp-servers`` to a delegate-enabled server while the URL path # targets a non-delegate server, skipping LiteLLM auth for it. target_names = MCPRequestHandler._resolve_target_server_names( diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 0ba0181200f5..449664828305 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -557,10 +557,12 @@ async def delete_mcp_server( """ Delete the mcp server from the db by server_id - The server-row delete is the commit point. Per-user env var rows have no FK - cascade, so they are cleaned up afterwards on a best-effort basis: a transient - failure there leaves only orphaned rows pointing at a now-missing server and - must not turn a successful delete into a caller-visible error. + The server-row delete is the commit point. Per-user credential and env var + rows have no FK cascade, so they are cleaned up afterwards on a best-effort + basis: a transient failure there leaves only orphaned rows pointing at a + now-missing server and must not turn a successful delete into a + caller-visible error. Each table is cleaned independently so a failure on one + still attempts the other. Returns the deleted mcp server record if it exists, otherwise None """ @@ -570,17 +572,20 @@ async def delete_mcp_server( }, ) if deleted_server is not None: - try: - await prisma_client.db.litellm_mcpuserenvvars.delete_many( - where={"server_id": server_id} - ) - except Exception as e: - verbose_proxy_logger.warning( - "MCP server %s deleted but per-user env var cleanup failed; " - "orphaned rows can be removed on a later delete: %s", - server_id, - e, - ) + for model, label in ( + (prisma_client.db.litellm_mcpusercredentials, "credential"), + (prisma_client.db.litellm_mcpuserenvvars, "env var"), + ): + try: + await model.delete_many(where={"server_id": server_id}) + except Exception as e: + verbose_proxy_logger.warning( + "MCP server %s deleted but per-user %s cleanup failed; " + "orphaned rows can be removed on a later delete: %s", + server_id, + label, + e, + ) return deleted_server diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index ed374635fead..3beddd2c4354 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -512,12 +512,13 @@ async def exchange_token_with_server( result = { "access_token": access_token, "token_type": token_response.get("token_type", "Bearer"), - "expires_in": token_response.get("expires_in", 3600), } - if "refresh_token" in token_response and token_response["refresh_token"]: + if token_response.get("expires_in") is not None: + result["expires_in"] = token_response["expires_in"] + if token_response.get("refresh_token"): result["refresh_token"] = token_response["refresh_token"] - if "scope" in token_response and token_response["scope"]: + if token_response.get("scope"): result["scope"] = token_response["scope"] # RFC 6749 §5.1: token responses must not be cached. diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 725f7a335bc4..2149f079a3d2 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -386,8 +386,15 @@ async def _get_tools_for_single_server( raw_headers: Optional[Dict[str, str]] = None, user_api_key_auth: Optional[UserAPIKeyAuth] = None, extra_headers: Optional[Dict[str, str]] = None, + apply_tool_filters: bool = True, ): - """Helper function to get tools for a single server.""" + """Helper function to get tools for a single server. + + When ``apply_tool_filters`` is False the raw server catalog is returned + without the allowed_tools/disallowed_tools gate or the per-key tool + permissions. This is the admin-only configuration view; every runtime + path keeps the default True so callable tools stay filtered. + """ tools = await global_mcp_server_manager._get_tools_from_server( server=server, mcp_auth_header=server_auth_header, @@ -397,6 +404,9 @@ async def _get_tools_for_single_server( user_api_key_auth=user_api_key_auth, ) + if not apply_tool_filters: + return _create_tool_response_objects(tools, server.mcp_info) + # Always apply allowed_tools/disallowed_tools so the blacklist is # enforced even when no allowlist is set (matches the SSE/HTTP path). tools = filter_tools_by_allowed_tools(tools, server) @@ -463,6 +473,7 @@ async def _list_tools_for_single_server( mcp_auth_header: Optional[str], raw_headers_from_request: dict, user_api_key_dict: UserAPIKeyAuth, + apply_tool_filters: bool = True, ) -> dict: """Handle tool listing for a single server_id request.""" # Resolve a server name to its UUID if needed @@ -527,6 +538,7 @@ async def _list_tools_for_single_server( raw_headers_from_request, user_api_key_dict, extra_headers=user_oauth_extra_headers, + apply_tool_filters=apply_tool_filters, ) except MCPUpstreamAuthError: # Surface the upstream 401/403 to the caller so it can emit the @@ -552,6 +564,14 @@ async def list_tool_rest_api( server_id: Optional[str] = Query( None, description="The server id to list tools for" ), + include_disabled_tools: bool = Query( + False, + description=( + "Admin only. Return the full server tool catalog without the " + "allowed_tools filter or per-key tool permissions, so the MCP " + "settings UI can configure the allowlist. Ignored for non-admins." + ), + ), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ) -> dict: """ @@ -579,6 +599,13 @@ async def list_tool_rest_api( ) try: + # The full catalog (allowlist filter skipped) is admin-only so the + # REST endpoint can't be used to enumerate deliberately-disabled tools. + apply_tool_filters = not ( + include_disabled_tools + and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN + ) + # Extract auth headers from request headers = request.headers raw_headers_from_request = dict(headers) @@ -620,6 +647,7 @@ async def list_tool_rest_api( mcp_auth_header=mcp_auth_header, raw_headers_from_request=raw_headers_from_request, user_api_key_dict=user_api_key_dict, + apply_tool_filters=apply_tool_filters, ) else: if not allowed_server_ids: @@ -677,6 +705,7 @@ async def list_tool_rest_api( raw_headers_from_request, user_api_key_dict, extra_headers=user_oauth_extra_headers, + apply_tool_filters=apply_tool_filters, ) list_tools_result.extend(tools_result) except Exception as e: diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index e5d709330636..c9fe3a0c88e9 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2473,6 +2473,17 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): "`statement_cache_size`). Keys here override any default LiteLLM sets." ), ) + database_disable_prepared_statements: Optional[bool] = Field( + None, + description=( + "Disable server-side prepared statements by setting Prisma's " + "`pgbouncer=true` URL param. Use this for pgbouncer transaction-pooling " + "deployments, or to prevent the 'cached plan must not change result " + "type' error that pooled connections hit during rolling schema " + "migrations. An explicit `pgbouncer` in `database_extra_connection_params` " + "takes precedence." + ), + ) database_type: Optional[Literal["dynamo_db"]] = Field( None, description="to use dynamodb instead of postgres db" ) @@ -2609,6 +2620,24 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): None, description="List of MCP server fields that must be filled in for a submission to pass standards checks (e.g. ['description', 'source_url', 'alias']).", ) + disable_budget_reservation: Optional[bool] = Field( + None, + description=( + "If True, disables the optimistic per-request budget reservation " + "introduced in v1.84.0. " + "WARNING: This weakens hard budget enforcement. Without the reservation, " + "a burst of concurrent requests from a single key can each pass the " + "read-time spend check before any of them is charged, allowing a " + "configured budget to be exceeded under high concurrency. " + "Budgets are still evaluated on every request at read time, so " + "an already-exhausted budget is still rejected. " + "Enable only if your deployment is experiencing phantom " + "BudgetExceededError responses caused by leaked reservations " + "(see GitHub issue #27639). " + "A proxy-level WARNING is logged on every request while this flag " + "is active as a reminder that hard enforcement is relaxed." + ), + ) class ConfigYAML(LiteLLMPydanticObjectBase): diff --git a/litellm/proxy/auth/auth_exception_handler.py b/litellm/proxy/auth/auth_exception_handler.py index e06ac760237d..a8c10d38e9ec 100644 --- a/litellm/proxy/auth/auth_exception_handler.py +++ b/litellm/proxy/auth/auth_exception_handler.py @@ -154,6 +154,16 @@ async def _handle_authentication_error( ) elif isinstance(e, ProxyException): raise e + if PrismaDBExceptionHandler.is_database_service_unavailable_error(e): + raise ProxyException( + message=( + "Service Unavailable, the authentication database is " + "temporarily unreachable. Please retry shortly." + ), + type=ProxyErrorTypes.no_db_connection, + param="None", + code=status.HTTP_503_SERVICE_UNAVAILABLE, + ) raise ProxyException( message="Authentication Error, " + str(e), type=ProxyErrorTypes.auth_error, diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index a5501fefa4eb..a62e74003a5f 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -2422,6 +2422,7 @@ async def _run_centralized_common_checks( # noqa: PLR0915 user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, skip_budget_checks=skip_budget_checks, + general_settings=general_settings, ) @@ -2442,12 +2443,23 @@ async def _reserve_budget_after_common_checks( user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, skip_budget_checks: bool, + general_settings: dict, end_user_id: Optional[str] = None, end_user_object: Optional[LiteLLM_EndUserTable] = None, ) -> None: user_api_key_auth_obj.budget_reservation = None if skip_budget_checks: return + if general_settings.get("disable_budget_reservation") is True: + verbose_proxy_logger.warning( + "disable_budget_reservation is enabled: skipping optimistic budget " + "reservation. Budget enforcement is read-time only — concurrent " + "requests can each pass the spend check before their cost is recorded, " + "so a configured budget may be briefly exceeded under high concurrency. " + "Set disable_budget_reservation to False or remove it to restore " + "hard per-request budget enforcement." + ) + return from litellm.proxy.spend_tracking.budget_reservation import ( reserve_budget_for_request, diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index ab9d341aa510..c500e727595e 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -109,6 +109,92 @@ def is_database_transport_error(e: Exception) -> bool: return True return False + @staticmethod + def is_prisma_engine_internal_error(e: Exception) -> bool: + """True iff ``e`` is a non-``PrismaError`` exception raised from inside + prisma-client-py's query-engine layer. + + During the instant a DB connection is torn down, the query engine can + return a malformed error payload (``user_facing_error.meta`` is + ``null``). prisma-client-py's ``handle_response_errors`` then crashes + with ``AttributeError: 'NoneType' object has no attribute 'get'`` + before it can raise the proper P1001 "can't reach database server" + error. That AttributeError carries no connection keyword, so it can't + be matched by message; identify it by its ``prisma.engine`` origin + instead. + + Recognized ``PrismaError`` subclasses are excluded: connectivity ones + are already classified by type/keyword above, and data-layer ones + (the DB IS reachable) must stay 401. + """ + import prisma + + if isinstance(e, prisma.errors.PrismaError): + return False + tb = getattr(e, "__traceback__", None) + while tb is not None: + if tb.tb_frame.f_globals.get("__name__", "").startswith("prisma.engine"): + return True + tb = tb.tb_next + return False + + @staticmethod + def is_database_service_unavailable_error(e: Exception) -> bool: + """True iff the exception means the database could not answer at the + infrastructure level (connection refused, socket/interface failure, + timeout) rather than a genuine auth failure (key not found) or a + data-layer error (the DB IS reachable and rejected the data). + + Auth must answer 401 only for a key the DB confirms is invalid. When + the DB itself is unreachable, the request has to surface as 503 so + callers retry instead of treating valid keys as invalid during an + outage. + + Note: prisma-client-py mislabels the P1001 "can't reach database + server" connectivity failure as a ``DataError`` (a data-layer type), + so a type-only check misses real outages. ``is_database_transport_error`` + keyword-matches the connection message and catches that masquerade, + while genuine data errors (no connection keyword) correctly stay 401. + + The Postgres "cached plan must not change result type" error is matched + here, not in ``is_database_transport_error``: it is a transient stale-DB- + state condition (not an invalid key), but the connection is healthy so it + must not trigger a reconnect. + + A non-``PrismaError`` raised from inside the prisma query engine (e.g. + the ``AttributeError`` from ``handle_response_errors`` when the engine + returns a malformed error payload mid-tear-down) is also treated as + unavailable; see ``is_prisma_engine_internal_error``. + """ + import asyncio + + if PrismaDBExceptionHandler.is_database_connection_error(e): + return True + if PrismaDBExceptionHandler.is_database_transport_error(e): + return True + if PrismaDBExceptionHandler.is_prisma_engine_internal_error(e): + return True + if "cached plan must not change result type" in str(e).lower(): + return True + + # OSError already covers ConnectionError and (Py3.3+) TimeoutError. + # asyncio.TimeoutError is a distinct class before Py3.11. + if isinstance(e, (OSError, asyncio.TimeoutError)): + return True + + try: + import asyncpg + except ImportError: + return False + + return isinstance( + e, + ( + asyncpg.exceptions.PostgresConnectionError, + asyncpg.exceptions.InterfaceError, + ), + ) + @staticmethod def handle_db_exception(e: Exception): """ diff --git a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py b/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py index 14d950ecdf4c..16d1dfa62efd 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py +++ b/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py @@ -105,6 +105,16 @@ def _extract_text_from_content(content: object) -> str: return "" +def _merge_metadata_bags(request_data: Mapping[str, Any]) -> Optional[dict[str, Any]]: + merged: dict[str, Any] = {} + present = False + for bag in (request_data.get("metadata"), request_data.get("litellm_metadata")): + if isinstance(bag, Mapping): + present = True + merged.update(bag) + return merged if present else None + + class CrowdStrikeAIDRHandler(CustomGuardrail): """ CrowdStrike AIDR AI Guardrail handler to interact with the CrowdStrike AIDR @@ -317,6 +327,22 @@ async def apply_guardrail( "event_type": event_type, } + model = inputs.get("model") + if model: + ai_guard_payload["model"] = model + + metadata = _merge_metadata_bags(request_data) + if metadata is not None: + user_id = metadata.get("user_api_key_user_id") + if user_id: + ai_guard_payload["user_id"] = user_id + + extra_info: dict[str, str] = {} + user_email = metadata.get("user_api_key_user_email") + if user_email: + extra_info["user_name"] = user_email + ai_guard_payload["extra_info"] = extra_info + ai_guard_response = await self._call_crowdstrike_aidr_guard( ai_guard_payload, hook_name ) diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index 8473b5e77de6..c5715872373f 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -518,11 +518,17 @@ async def count_input_file_usage( # Check if this is a managed file (base64 encoded unified file ID) from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, + get_models_from_unified_file_id, ) # Managed files require bypassing the HTTP endpoint (which runs access-check hooks) # and calling the managed files hook directly with the user's credentials. is_managed_file = _is_base64_encoded_unified_file_id(file_id) + target_model_names = ( + get_models_from_unified_file_id(is_managed_file) + if is_managed_file + else [] + ) if is_managed_file and user_api_key_dict is not None: file_content = await self._fetch_managed_file_content( file_id=file_id, @@ -560,6 +566,7 @@ async def count_input_file_usage( await self._enforce_batch_file_model_access( user_api_key_dict=user_api_key_dict, file_content_as_dict=file_content_as_dict, + target_model_names=target_model_names or None, ) input_file_usage = _get_batch_job_input_file_usage( @@ -595,9 +602,13 @@ async def _enforce_batch_file_model_access( self, user_api_key_dict: UserAPIKeyAuth, file_content_as_dict: List[dict], + target_model_names: Optional[List[str]] = None, ) -> None: - """Reject the batch if the caller is not authorized for every - ``body.model`` named inside the JSONL. + """Reject the batch if the caller is not authorized for the upload target. + + For managed files, ``target_model_names`` (from the unified file id) is + the proxy alias the file was uploaded for and is used directly for auth. + For legacy/non-managed files, falls back to ``body.model`` values in the JSONL. Reuses standard auth helpers so the same model access rules the proxy enforces on `/chat/completions` apply here. @@ -614,9 +625,12 @@ async def _enforce_batch_file_model_access( from litellm.proxy.proxy_server import proxy_logging_obj from litellm.proxy.proxy_server import user_api_key_cache - models = _get_models_from_batch_input_file_content(file_content_as_dict) - if not models: - return + if target_model_names: + models = target_model_names + else: + models = _get_models_from_batch_input_file_content(file_content_as_dict) + if not models: + return team_object = None if ( @@ -647,12 +661,7 @@ async def _enforce_batch_file_model_access( llm_model_list = llm_router.model_list if llm_router is not None else None for model in models: - # body.model may be the provider id after replace_model_in_jsonl; map to proxy model_name for auth. model_to_check = model - if llm_router is not None: - proxy_model_name = llm_router.resolve_model_name_from_model_id(model) - if proxy_model_name is not None: - model_to_check = proxy_model_name try: if team_object is not None: try: diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index f1edcc9c7b2f..46f5c1f7f824 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -21,7 +21,7 @@ import os from dataclasses import dataclass from datetime import datetime, timedelta, timezone -from typing import Any, Dict, Iterable, List, Literal, Optional +from typing import Any, Dict, Iterable, List, Literal, Optional, Set from fastapi import ( APIRouter, @@ -1714,11 +1714,13 @@ async def _get_cached_temporary_mcp_server_or_404( status_code=status.HTTP_403_FORBIDDEN, detail={"error": f"Access denied to MCP server {server_id}"}, ) - allowed_server_ids = ( - await global_mcp_server_manager.get_allowed_mcp_servers( - user_api_key_dict + allowed_server_ids: Set[str] = set() + for auth_context in await build_effective_auth_contexts(user_api_key_dict): + allowed_server_ids.update( + await global_mcp_server_manager.get_allowed_mcp_servers( + auth_context + ) ) - ) if server.server_id not in allowed_server_ids: raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index ae7da0d29f26..f9351c1bcb68 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -4684,15 +4684,36 @@ async def team_model_add( detail={"error": "Only proxy admin or team admin can modify team models"}, ) - updated_models = add_new_models_to_team(team_obj=team_obj, new_models=data.models) - # Update team. `include` mirrors the relations the auth path consumes - # off the cached team object so that `_refresh_cached_team` doesn't - # null them out — see object_permission_utils.validate_key_search_tools_against_team - # and the MCP/agent authz paths, which treat a missing object_permission - # as "no team-level restriction". + # Atomic array append with dedup at the database level so concurrent + # BYOK model creates don't overwrite each other's team.models entries. + # When the team currently has models=[] (unrestricted access), the + # CASE expression inserts the 'all-proxy-models' sentinel first. + models_to_add = list(data.models) + await prisma_client.db.execute_raw( + 'UPDATE "LiteLLM_TeamTable" ' + "SET models = (" + " SELECT ARRAY(SELECT DISTINCT unnest(" + " CASE WHEN cardinality(COALESCE(models, ARRAY[]::text[])) = 0 " + " THEN ARRAY['all-proxy-models']::text[] " + " ELSE models " + " END || $1::text[]" + " ))" + ") " + "WHERE team_id = $2", + models_to_add, + data.team_id, + ) + # Re-fetch via update (write-routed) instead of find_unique (read-routed) + # to avoid returning stale data from a read replica. The models column was + # already set by execute_raw above; this bumps updated_at. `include` mirrors + # the relations the auth path consumes off the cached team object so that + # `_refresh_cached_team` doesn't null them out — see + # object_permission_utils.validate_key_search_tools_against_team and the + # MCP/agent authz paths, which treat a missing object_permission as + # "no team-level restriction". updated_team = await prisma_client.db.litellm_teamtable.update( where={"team_id": data.team_id}, - data={"models": updated_models}, + data={"updated_at": datetime.now(timezone.utc)}, include={"object_permission": True}, # type: ignore ) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py index a94672f94874..a912a88a993b 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py @@ -100,6 +100,42 @@ def _get_user_from_metadata( return get_end_user_id_from_request_body(request_body) return None + @staticmethod + def _resolve_costing_model(model: str, logging_obj: LiteLLMLoggingObj) -> str: + if model and model != "unknown": + return model + litellm_params = (getattr(logging_obj, "model_call_details", {}) or {}).get( + "litellm_params", {} + ) or {} + deployment_model = litellm_params.get("model") + if deployment_model and deployment_model != "unknown": + return deployment_model + model_group = (litellm_params.get("metadata", {}) or {}).get("model_group") + if model_group: + return model_group.removeprefix("passthrough/") + return model + + @staticmethod + def _extract_model_from_anthropic_chunks( + all_chunks: Sequence[Union[str, bytes]], + ) -> Optional[str]: + for raw in all_chunks: + text = raw.decode("utf-8") if isinstance(raw, bytes) else raw + for line in text.splitlines(): + if not line.startswith("data:"): + continue + try: + data = json.loads(line[len("data:") :].strip()) + except (json.JSONDecodeError, ValueError): + continue + if not isinstance(data, dict): + continue + if data.get("type") == "message_start": + model = (data.get("message") or {}).get("model") + if model: + return model + return None + @staticmethod def _create_anthropic_response_logging_payload( litellm_model_response: Union[ModelResponse, TextCompletionResponse], @@ -127,6 +163,10 @@ def _create_anthropic_response_logging_payload( "custom_llm_provider" ) + model = AnthropicPassthroughLoggingHandler._resolve_costing_model( + model, logging_obj + ) + # Prepend custom_llm_provider to model if not already present model_for_cost = model if custom_llm_provider and not model.startswith(f"{custom_llm_provider}/"): @@ -213,6 +253,15 @@ def _handle_logging_anthropic_collected_chunks( ): model = cast(str, litellm_logging_obj.model_call_details.get("model")) + if not model or model == "unknown": + chunk_model = ( + AnthropicPassthroughLoggingHandler._extract_model_from_anthropic_chunks( + all_chunks + ) + ) + if chunk_model: + model = chunk_model + complete_streaming_response = ( AnthropicPassthroughLoggingHandler._build_complete_streaming_response( all_chunks=all_chunks, @@ -468,6 +517,13 @@ def _build_complete_streaming_response_legacy( # Process each individual event for event_str in individual_events: try: + # Skip OpenAI-style [DONE] sentinels some Anthropic-compatible + # providers emit. Match the whole SSE line so a valid chunk whose + # text payload happens to contain "[DONE]" is not dropped. + if any( + line.strip() == "data: [DONE]" for line in event_str.split("\n") + ): + continue transformed_openai_chunk = anthropic_model_response_iterator.convert_str_chunk_to_generic_chunk( chunk=event_str ) @@ -476,6 +532,14 @@ def _build_complete_streaming_response_legacy( except (StopIteration, StopAsyncIteration): break + except json.JSONDecodeError: + # Some upstreams emit non-JSON SSE lines; skip them so the + # logging pipeline is not broken by a single bad frame. + verbose_proxy_logger.debug( + "Skipping non-JSON SSE event: %s", + event_str[:200], + ) + continue complete_streaming_response = litellm.stream_chunk_builder( chunks=all_openai_chunks, diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index e4567b9f494e..c5ca9c50a9b8 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -44,15 +44,19 @@ def _build_db_connection_url_params( pool_timeout: Optional[Union[int, float]], connect_timeout: Optional[Union[int, float]] = None, socket_timeout: Optional[Union[int, float]] = None, + disable_prepared_statements: bool = False, extra_params: Optional[dict] = None, ) -> dict: """Build the Prisma DATABASE_URL query params controlling connection pool behavior. `connect_timeout` / `socket_timeout` map to the Prisma URL params of the same name (https://www.prisma.io/docs/orm/overview/databases/postgresql) and are - omitted when None so Prisma's defaults apply. `extra_params` is an - untyped passthrough — keys it provides win over the named arguments above, - so it can be used to override any default we set here. + omitted when None so Prisma's defaults apply. `disable_prepared_statements` + sets `pgbouncer=true`, which makes Prisma stop using server-side prepared + statements (pgbouncer transaction-pool compatible; also sidesteps the + "cached plan must not change result type" error during rolling migrations). + `extra_params` is an untyped passthrough — keys it provides win over the + named arguments above, so it can be used to override any default we set here. """ params: dict = { "connection_limit": connection_limit, @@ -63,6 +67,8 @@ def _build_db_connection_url_params( params["connect_timeout"] = connect_timeout if socket_timeout is not None: params["socket_timeout"] = socket_timeout + if disable_prepared_statements: + params["pgbouncer"] = "true" if extra_params: params.update(extra_params) return params @@ -947,6 +953,7 @@ def run_server( # noqa: PLR0915 db_connection_timeout: Optional[Union[int, float]] = 60 db_connect_timeout: Optional[Union[int, float]] = None db_socket_timeout: Optional[Union[int, float]] = None + db_disable_prepared_statements: bool = False db_extra_connection_params: Optional[dict] = None general_settings = {} ### GET DB TOKEN FOR IAM AUTH ### @@ -1067,6 +1074,17 @@ def run_server( # noqa: PLR0915 ) db_connect_timeout = general_settings.get("database_connect_timeout") db_socket_timeout = general_settings.get("database_socket_timeout") + _disable_prepared_statements = general_settings.get( + "database_disable_prepared_statements", False + ) + if isinstance(_disable_prepared_statements, str): + from litellm.secret_managers.main import str_to_bool + + db_disable_prepared_statements = ( + str_to_bool(_disable_prepared_statements) is True + ) + else: + db_disable_prepared_statements = bool(_disable_prepared_statements) db_extra_connection_params = general_settings.get( "database_extra_connection_params" ) @@ -1114,6 +1132,7 @@ def run_server( # noqa: PLR0915 pool_timeout=db_connection_timeout, connect_timeout=db_connect_timeout, socket_timeout=db_socket_timeout, + disable_prepared_statements=db_disable_prepared_statements, extra_params=db_extra_connection_params, ) if os.getenv("DATABASE_URL", None) is not None: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 72423b2a7968..a6413138bd57 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -10920,16 +10920,26 @@ def get_direct_access_models( return direct_access_models -async def get_all_team_and_direct_access_models( +def _filter_models_to_user_accessible(all_models: List[Dict]) -> List[Dict]: + """Keep only deployments the caller can use via direct access or team membership.""" + return [ + _model + for _model in all_models + if _model.get("model_info", {}).get("direct_access", False) + or _model.get("model_info", {}).get("access_via_team_ids", []) + ] + + +async def _populate_team_access_on_models( user_api_key_dict: UserAPIKeyAuth, prisma_client: PrismaClient, llm_router: Router, all_models: List[Dict], ) -> List[Dict]: """ - Get all models across all teams user is in. + Populate `model_info.access_via_team_ids` and `model_info.direct_access` + without filtering the model list. """ - user_teams: Optional[Union[List[str], Literal["*"]]] = None direct_access_models: List[str] = [] if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: @@ -10948,7 +10958,6 @@ async def get_all_team_and_direct_access_models( user_db_object=user_object, llm_router=llm_router, ) - ## ADD ACCESS_VIA_TEAM_IDS TO ALL MODELS if user_teams is not None: team_models = await get_all_team_models( user_teams=user_teams, @@ -10971,23 +10980,33 @@ async def get_all_team_and_direct_access_models( model_id, [] ) - ## ADD DIRECT_ACCESS TO RELEVANT MODELS - + direct_access_model_ids = set(direct_access_models) for _model in all_models: model_id = _model.get("model_info", {}).get("id", None) - if model_id is not None and model_id in direct_access_models: - _model["model_info"]["direct_access"] = True + if model_id is not None: + _model["model_info"]["direct_access"] = model_id in direct_access_model_ids - ## FILTER OUT MODELS THAT ARE NOT IN DIRECT_ACCESS_MODELS OR ACCESS_VIA_TEAM_IDS - only show user models they can call - all_models = [ - _model - for _model in all_models - if _model.get("model_info", {}).get("direct_access", False) - or _model.get("model_info", {}).get("access_via_team_ids", []) - ] return all_models +async def get_all_team_and_direct_access_models( + user_api_key_dict: UserAPIKeyAuth, + prisma_client: PrismaClient, + llm_router: Router, + all_models: List[Dict], +) -> List[Dict]: + """ + Get all models across all teams user is in. + """ + all_models = await _populate_team_access_on_models( + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + llm_router=llm_router, + all_models=all_models, + ) + return _filter_models_to_user_accessible(all_models) + + def _enrich_model_info_with_litellm_data( model: Dict[str, Any], debug: bool = False, llm_router: Optional[Router] = None ) -> Dict[str, Any]: @@ -11096,6 +11115,22 @@ async def _get_caller_byok_team_scope( return set(user_row.teams or []) +def _byok_row_outside_caller_teams( + model_info_dict: Dict[str, Any], allowed_team_ids: Optional[Set[str]] +) -> bool: + """Whether a team BYOK row belongs to a team the caller is not a member of. + + `team_id` is only set on team BYOK rows; non-team rows fall through + unaffected. `allowed_team_ids is None` means no scoping (e.g. admins). + """ + if allowed_team_ids is None: + return False + team_id = model_info_dict.get("team_id") + if team_id is None: + return False + return team_id not in allowed_team_ids + + # Hard cap on rows the DB-side BYOK search may pull when results need to be # sorted across the full match set. Without this, an authenticated caller # can hit `/v2/model/info?search=&sortBy=` and force the @@ -11217,15 +11252,7 @@ async def _apply_search_filter_to_models( ) def _is_byok_outside_caller_teams(model_info_dict: Dict[str, Any]) -> bool: - # `team_id` is only set on team BYOK rows. Non-team rows fall - # through unaffected — they are gated by other paths (router - # membership, direct_access, include_team_models). - if allowed_team_ids is None: - return False - team_id = model_info_dict.get("team_id") - if team_id is None: - return False - return team_id not in allowed_team_ids + return _byok_row_outside_caller_teams(model_info_dict, allowed_team_ids) def _model_matches_search(m: Dict[str, Any]) -> bool: # Team BYOK models persist an internal `model_name` @@ -11723,10 +11750,8 @@ async def _find_model_by_id( @router.get( "/v2/model/info", - description="v2 - returns models available to the user based on their API key permissions. Shows model info from config.yaml (except api key and api base). Filter to just user-added models with ?user_models_only=true", tags=["model management"], dependencies=[Depends(user_api_key_auth)], - include_in_schema=False, ) async def model_info_v2( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), @@ -11762,7 +11787,49 @@ async def model_info_v2( ), ): """ - BETA ENDPOINT. Might change unexpectedly. Use `/v1/model/info` for now. + Paginated model metadata for proxy deployments (pricing, provider, team access). + + Returns configured router deployments with enriched `model_info` (costs, provider, + context window, etc.). Sensitive fields such as API keys and api_base are omitted. + + Query parameters: + model: Filter to a single public `model_name`. + user_models_only: When true, only return models created by the calling user. + include_team_models: When true, populate `access_via_team_ids` and `direct_access` + on each model and filter to deployments the caller can use. + page / size: Pagination controls (defaults: page=1, size=50). + search: Case-insensitive partial match on model name or team public name. + modelId: Return a single deployment by LiteLLM model id. + teamId: Filter to models with direct access or team membership for this team id. + sortBy / sortOrder: Sort by model_name, created_at, updated_at, costs, or status. + + Example request: + ``` + curl -X GET 'http://localhost:4000/v2/model/info?include_team_models=true&page=1&size=50' \\ + --header 'Authorization: Bearer sk-1234' + ``` + + Example response: + ```json + { + "data": [ + { + "model_name": "gpt-4", + "litellm_params": {"model": "openai/gpt-4.1"}, + "model_info": { + "id": "abc123", + "litellm_provider": "openai", + "access_via_team_ids": ["team-1"], + "direct_access": true + } + } + ], + "total_count": 1, + "current_page": 1, + "total_pages": 1, + "size": 50 + } + ``` """ global llm_model_list, general_settings, user_config_file_path, proxy_config, llm_router @@ -12325,6 +12392,72 @@ async def model_metrics_exceptions( return {"data": response, "exception_types": list(exception_types)} +def _deployment_matches_allowed_model_names( + model: Dict[str, Any], allowed_model_names: Set[str] +) -> bool: + """Match a router deployment against allowed public model names. + + Team-scoped rows store an internal routing key in ``model_name``; callers + with key/team restrictions still refer to the public name in + ``model_info.team_public_model_name``. + """ + if model.get("model_name") in allowed_model_names: + return True + model_info = model.get("model_info") + if not isinstance(model_info, dict): + return False + team_public_model_name = model_info.get("team_public_model_name") + return ( + isinstance(team_public_model_name, str) + and team_public_model_name in allowed_model_names + ) + + +def _get_v1_model_info_allowed_model_names( + user_api_key_dict: UserAPIKeyAuth, + llm_router: Router, +) -> Optional[Set[str]]: + """Return key/team allowlisted public model names, or None if unrestricted.""" + model_access_groups = llm_router.get_model_access_groups() + proxy_model_list = llm_router.get_model_names() + key_models = get_key_models( + user_api_key_dict=user_api_key_dict, + proxy_model_list=proxy_model_list, + model_access_groups=model_access_groups, + ) + team_models = get_team_models( + team_models=user_api_key_dict.team_models, + proxy_model_list=proxy_model_list, + model_access_groups=model_access_groups, + ) + if not key_models and not team_models: + return None + return set( + get_complete_model_list( + key_models=key_models, + team_models=team_models, + proxy_model_list=proxy_model_list, + user_model=user_model, + infer_model_from_keys=general_settings.get("infer_model_from_keys", False), + llm_router=llm_router, + return_wildcard_routes=False, + ) + ) + + +def _filter_v1_model_info_deployments( + all_models: List[dict], + allowed_model_names: Optional[Set[str]], +) -> List[dict]: + if allowed_model_names is None: + return all_models + return [ + model + for model in all_models + if _deployment_matches_allowed_model_names(model, allowed_model_names) + ] + + def _translate_model_name_for_response(model: dict) -> dict: """For team-scoped DB rows, replace `model_name` with the public name in `model_info.team_public_model_name` before returning. The DB column @@ -12408,6 +12541,14 @@ def _get_proxy_model_info(model: dict) -> dict: async def model_info_v1( # noqa: PLR0915 user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), litellm_model_id: Optional[str] = None, + include_team_models: Optional[bool] = fastapi.Query( + False, + description="When true, filter to deployments the caller can use via direct access or team membership.", + ), + teamId: Optional[str] = fastapi.Query( + None, + description="Filter models by team ID. Returns models with direct_access=True or teamId in access_via_team_ids", + ), ): """ Provides more info about each model in /models, including config.yaml descriptions (except api key and api base) @@ -12417,6 +12558,11 @@ async def model_info_v1( # noqa: PLR0915 - When litellm_model_id is passed, it will return the info for that specific model - When litellm_model_id is not passed, it will return the info for all models + - include_team_models: When true, filter to deployments the caller can use (same as /v2/model/info). + - teamId: Filter to models accessible by the given team. + + Each model in the list response includes `model_info.access_via_team_ids` and + `model_info.direct_access` when the proxy database is connected. Returns: Returns a dictionary containing information about each model. @@ -12443,6 +12589,12 @@ async def model_info_v1( # noqa: PLR0915 """ global llm_model_list, general_settings, user_config_file_path, proxy_config, llm_router, user_model + # Unit tests call this handler directly; FastAPI normally resolves Query defaults. + if not isinstance(include_team_models, bool): + include_team_models = False + if not isinstance(teamId, str): + teamId = None + if user_model is not None: # user is trying to get specific model from litellm router try: @@ -12479,6 +12631,14 @@ async def model_info_v1( # noqa: PLR0915 }, ) + if prisma_client is None and ( + include_team_models or (teamId is not None and teamId.strip()) + ): + raise HTTPException( + status_code=500, + detail={"error": CommonProxyErrors.db_not_connected_error.value}, + ) + if litellm_model_id is not None: # user is trying to get specific model from litellm router deployment_info = llm_router.get_deployment(model_id=litellm_model_id) @@ -12492,51 +12652,82 @@ async def model_info_v1( # noqa: PLR0915 _deployment_info_dict = _get_proxy_model_info( model=deployment_info.model_dump(exclude_none=True) ) - return {"data": [_deployment_info_dict]} + single_model_list: List[dict] = [_deployment_info_dict] + if prisma_client is not None: + single_model_list = await _populate_team_access_on_models( + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + llm_router=llm_router, + all_models=single_model_list, + ) + if include_team_models: + single_model_list = _filter_models_to_user_accessible(single_model_list) + if teamId is not None and teamId.strip(): + single_model_list = await _filter_models_by_team_id( + all_models=single_model_list, + team_id=teamId.strip(), + prisma_client=prisma_client, + llm_router=llm_router, + user_api_key_dict=user_api_key_dict, + ) + return {"data": single_model_list} - all_models: List[dict] = [] - model_access_groups: Dict[str, List[str]] = defaultdict(list) - ## CHECK IF MODEL RESTRICTIONS ARE SET AT KEY/TEAM LEVEL ## - if llm_router is None: - proxy_model_list = [] - else: - proxy_model_list = llm_router.get_model_names() - model_access_groups = llm_router.get_model_access_groups() - key_models = get_key_models( + # Return router deployments (same source as /v2/model/info), not wildcard- + # expanded model names from get_complete_model_list(). Team-scoped rows + # use internal routing keys (model_name_{team_id}_{uuid}) and were omitted + # when v1 resolved models only via public model_name strings. + all_models: List[dict] = copy.deepcopy(llm_router.model_list) + allowed_model_names = _get_v1_model_info_allowed_model_names( user_api_key_dict=user_api_key_dict, - proxy_model_list=proxy_model_list, - model_access_groups=model_access_groups, + llm_router=llm_router, ) - team_models = get_team_models( - team_models=user_api_key_dict.team_models, - proxy_model_list=proxy_model_list, - model_access_groups=model_access_groups, + + all_models = _filter_v1_model_info_deployments( + all_models=all_models, + allowed_model_names=allowed_model_names, ) - all_models_str = get_complete_model_list( - key_models=key_models, - team_models=team_models, - proxy_model_list=proxy_model_list, - user_model=user_model, - infer_model_from_keys=general_settings.get("infer_model_from_keys", False), - llm_router=llm_router, + + # Team BYOK deployments carry an internal routing key and other teams' + # public name/team_id/api_base; drop the ones the caller cannot access so + # listing the full router model_list does not leak cross-team metadata. + allowed_team_ids = await _get_caller_byok_team_scope( + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, ) + all_models = [ + model + for model in all_models + if not _byok_row_outside_caller_teams( + model.get("model_info") or {}, allowed_team_ids + ) + ] - if len(all_models_str) > 0: - _relevant_models = [] - for model in all_models_str: - router_models = llm_router.get_model_list(model_name=model) - if router_models is not None: - _relevant_models.extend(router_models) - if llm_model_list is not None: - all_models = copy.deepcopy(_relevant_models) # type: ignore - else: - all_models = [] + if prisma_client is not None: + all_models = await _populate_team_access_on_models( + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + llm_router=llm_router, + all_models=all_models, + ) + + if include_team_models: + all_models = _filter_models_to_user_accessible(all_models) - # Reassign each entry: _get_proxy_model_info returns a (possibly new) - # dict via _translate_model_name_for_response, which does NOT mutate in - # place. Binding only the loop variable would drop the public-name swap - # for team-scoped rows and leak the internal routing key (#28382). - all_models = [_get_proxy_model_info(model=model) for model in all_models] + all_models = [ + _translate_model_name_for_response( + _enrich_model_info_with_litellm_data(model=model, llm_router=llm_router) + ) + for model in all_models + ] + + if teamId is not None and teamId.strip(): + all_models = await _filter_models_by_team_id( + all_models=all_models, + team_id=teamId.strip(), + prisma_client=cast(PrismaClient, prisma_client), + llm_router=llm_router, + user_api_key_dict=user_api_key_dict, + ) verbose_proxy_logger.debug("all_models: %s", all_models) return {"data": all_models} diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index e77e24c9e712..c15c37f6ad8a 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -3239,40 +3239,49 @@ async def _query_first_with_cached_plan_fallback( self, sql_query: str, *args ) -> Optional[dict]: """ - Execute a query with automatic fallback for PostgreSQL cached plan errors. - - This handles the "cached plan must not change result type" error that occurs - during rolling deployments when schema changes are applied while old pods - still have cached query plans expecting the old schema. - - Args: - sql_query: SQL query string to execute - - Returns: - Query result or None - - Raises: - Original exception if not a cached plan error + Execute a query, recovering once from PostgreSQL's "cached plan must not + change result type" error. + + That error surfaces during rolling deployments when a schema change + invalidates the prepared-statement plans that pooled connections still + hold. Clearing only the server-side plans with DEALLOCATE ALL makes + things worse: Prisma's query engine keeps a per-connection client-side + cache of prepared-statement names, so once the server drops a plan the + engine re-sends a name PostgreSQL no longer recognizes and the + connection breaks with `prepared statement "sN" does not exist`. With a + small pool that connection stays poisoned and every auth lookup fails. + + Recreating the Prisma client kills the engine subprocess and drops the + server-side plans and the engine's client-side name cache together, so + the retried query is prepared fresh. We reconnect through + `attempt_db_reconnect`, which is singleflight: when a schema change + poisons every pooled connection at once, the first cached-plan error + recreates the client and the concurrent waiters reuse that single + recreate instead of racing to kill each other's fresh engine. We then + retry the identical query exactly once. + + The retry reuses the original query byte-for-byte. Mutating the SQL + (e.g. injecting a unique comment) would defeat PostgreSQL's plan cache, + forcing a fresh plan on every request and pegging the database CPU. + + If the reconnect is skipped because a recent reconnect is still within + its cooldown, the retry runs against the same connection and may fail + again; the get_data backoff decorator re-runs the lookup and a later + attempt reconnects once the cooldown elapses. """ try: return await self.db.query_first(sql_query, *args) except Exception as e: - error_str = str(e) - if "cached plan must not change result type" in error_str: - # Force PostgreSQL to re-plan by invalidating the cache - # Add a unique comment to make the query different - sql_query_retry = sql_query.replace( - "SELECT", - f"SELECT /* cache_invalidated_{int(time.time() * 1000)} */", - ) - verbose_proxy_logger.warning( - "PostgreSQL cached plan error detected for token lookup, " - "retrying with fresh plan. This may occur during rolling deployments " - "when schema changes are applied." - ) - return await self.db.query_first(sql_query_retry, *args) - else: + if "cached plan must not change result type" not in str(e): raise + verbose_proxy_logger.warning( + "PostgreSQL cached plan error detected for token lookup; " + "recreating the database connection and retrying with the same " + "query. This may occur during rolling deployments when schema " + "changes are applied." + ) + await self.attempt_db_reconnect(reason="postgres_cached_plan_error") + return await self.db.query_first(sql_query, *args) @backoff.on_exception( backoff.expo, @@ -3628,7 +3637,10 @@ async def get_data( # noqa: PLR0915 db=self.db, hashed_token=hashed_token ) if active_token_id: - response = await self.get_data( + # The recursive call returns a finished + # LiteLLM_VerificationTokenView; the dict + # normalization below would crash subscripting it. + deprecated_response = await self.get_data( token=active_token_id, table_name="combined_view", query_type="find_unique", @@ -3636,10 +3648,11 @@ async def get_data( # noqa: PLR0915 proxy_logging_obj=proxy_logging_obj, check_deprecated=False, ) - if response is not None: + if deprecated_response is not None: verbose_proxy_logger.debug( "Deprecated key used during grace period" ) + return deprecated_response if response is not None: if response["team_models"] is None: diff --git a/litellm/setup_wizard.py b/litellm/setup_wizard.py index 862ca13e7ba4..2f0cb1233ae3 100644 --- a/litellm/setup_wizard.py +++ b/litellm/setup_wizard.py @@ -52,11 +52,12 @@ { "id": "anthropic", "name": "Anthropic", - "description": "Claude Opus 4.8, Opus 4.7, Opus 4.6, Sonnet 4.6, Haiku 4.5", + "description": "Claude Fable 5, Opus 4.8, Opus 4.7, Opus 4.6, Sonnet 4.6, Haiku 4.5", "env_key": "ANTHROPIC_API_KEY", "key_hint": "sk-ant-...", "test_model": "claude-haiku-4-5-20251001", "models": [ + "claude-fable-5", "claude-opus-4-8", "claude-opus-4-7", "claude-opus-4-6", diff --git a/litellm/types/utils.py b/litellm/types/utils.py index a7a0b0f62386..0200d1a0831d 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1672,6 +1672,9 @@ def __init__( # noqa: PLR0915 prompt_tokens_details=_prompt_tokens_details or None, ) + if isinstance(server_tool_use, dict): + server_tool_use = ServerToolUse(**server_tool_use) + if server_tool_use is not None: self.server_tool_use = server_tool_use else: # maintain openai compatibility in usage object if possible @@ -3026,6 +3029,12 @@ class StandardCallbackDynamicParams(TypedDict, total=False): wandb_api_key: Optional[str] weave_project_id: Optional[str] + # Datadog dynamic params + dd_api_key: Optional[str] + dd_site: Optional[str] + dd_agent_host: Optional[str] + dd_agent_port: Optional[str] + # Logging settings turn_off_message_logging: Optional[bool] # when true will not log messages litellm_disabled_callbacks: Optional[List[str]] diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index b2836a096b71..ddd7d51d76b3 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -1156,6 +1156,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1202,6 +1203,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1233,6 +1235,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1264,6 +1267,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1295,6 +1299,139 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh" + }, + "anthropic.claude-fable-5": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 1e-06, + "input_cost_per_token": 1e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh" + }, + "global.anthropic.claude-fable-5": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 1e-06, + "input_cost_per_token": 1e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh" + }, + "us.anthropic.claude-fable-5": { + "cache_creation_input_token_cost": 1.375e-05, + "cache_creation_input_token_cost_above_1hr": 2.2e-05, + "cache_read_input_token_cost": 1.1e-06, + "input_cost_per_token": 1.1e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh" + }, + "eu.anthropic.claude-fable-5": { + "cache_creation_input_token_cost": 1.375e-05, + "cache_creation_input_token_cost_above_1hr": 2.2e-05, + "cache_read_input_token_cost": 1.1e-06, + "input_cost_per_token": 1.1e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1327,6 +1464,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1359,6 +1497,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1391,6 +1530,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1423,6 +1563,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1455,6 +1596,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1485,6 +1627,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -2208,6 +2351,37 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true + }, + "azure_ai/claude-fable-5": { + "input_cost_per_token": 1e-05, + "output_cost_per_token": 5e-05, + "litellm_provider": "azure_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 1e-06, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -2237,6 +2411,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -10133,6 +10308,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -10167,6 +10343,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -10177,6 +10354,40 @@ }, "supports_output_config": true }, + "claude-fable-5": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 1e-06, + "input_cost_per_token": 1e-05, + "litellm_provider": "anthropic", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true + }, "claude-opus-4-8": { "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, @@ -10201,6 +10412,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -34007,6 +34219,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -34035,6 +34248,67 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true + }, + "vertex_ai/claude-fable-5": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 1e-06, + "input_cost_per_token": 1e-05, + "litellm_provider": "vertex_ai-anthropic_models", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true + }, + "vertex_ai/claude-fable-5@default": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 1e-06, + "input_cost_per_token": 1e-05, + "litellm_provider": "vertex_ai-anthropic_models", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -34064,6 +34338,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -34093,6 +34368,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, diff --git a/pyproject.toml b/pyproject.toml index 577e800d79bd..4a456b2c9a7e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm" -version = "1.89.0" +version = "1.89.1" description = "Library to easily interface with LLM API providers" readme = "README.md" requires-python = ">=3.10, <3.14" @@ -125,7 +125,7 @@ proxy-runtime = [ "mangum>=0.17.0,<1.0", "azure-ai-contentsafety>=1.0.0,<2.0", "azure-storage-file-datalake>=12.20.0,<13.0", - "pypdf>=6.10.2,<7.0; python_version < '3.14'", + "pypdf>=6.12.0,<7.0; python_version < '3.14'", "llm-sandbox>=0.3.39,<1.0", "detect-secrets>=1.5.0,<2.0", ] @@ -231,6 +231,10 @@ requires = ["uv_build==0.11.8"] build-backend = "uv_build" [tool.uv] +constraint-dependencies = [ + "tornado>=6.5.6", + "aiohttp>=3.13.5,<3.14", +] default-groups = ["dev"] required-version = ">=0.10.9" exclude-newer = "3 days" @@ -260,7 +264,7 @@ source-exclude = [ profile = "black" [tool.commitizen] -version = "1.89.0" +version = "1.89.1" version_files = [ "pyproject.toml:^version", ] diff --git a/tests/llm_translation/reasoning_effort_grid/grid_spec.py b/tests/llm_translation/reasoning_effort_grid/grid_spec.py index a08013cd4391..83a2c286d649 100644 --- a/tests/llm_translation/reasoning_effort_grid/grid_spec.py +++ b/tests/llm_translation/reasoning_effort_grid/grid_spec.py @@ -1,7 +1,6 @@ from dataclasses import dataclass, field from typing import Dict, FrozenSet, List, Optional, Tuple - OMIT = object() @@ -136,6 +135,13 @@ def expected(model: ModelEntry, effort: str) -> CellExpectation: ANTHROPIC_DIRECT_MODELS: Tuple[ModelEntry, ...] = ( + ModelEntry( + alias="claude-fable-5", + model="anthropic/claude-fable-5", + mode="adaptive", + required_env=_ANTHROPIC_REQ, + caps=_CAPS_XHIGH_MAX, + ), ModelEntry( alias="claude-opus-4-8", model="anthropic/claude-opus-4-8", @@ -168,6 +174,19 @@ def expected(model: ModelEntry, effort: str) -> CellExpectation: AZURE_AI_MODELS: Tuple[ModelEntry, ...] = ( + ModelEntry( + alias="azure-claude-fable-5", + model="azure_ai/claude-fable-5", + mode="adaptive", + required_env=_AZURE_FOUNDRY_REQ, + caps=_CAPS_XHIGH_MAX, + fail_reason=( + "claude-fable-5 has no deployment on the CI Microsoft Foundry " + "resource yet; Foundry returns DeploymentNotFound until someone " + "creates the fable-5 deployment, so this cell stays loud in CI. " + "Remove this fail_reason once the deployment exists." + ), + ), ModelEntry( alias="azure-claude-opus-4-8", model="azure_ai/claude-opus-4-8", @@ -213,6 +232,20 @@ def expected(model: ModelEntry, effort: str) -> CellExpectation: VERTEX_AI_MODELS: Tuple[ModelEntry, ...] = ( + ModelEntry( + alias="vertex-claude-fable-5", + model="vertex_ai/claude-fable-5", + mode="adaptive", + extra_params=(("vertex_location", "global"),), + required_env=_VERTEX_REQ, + caps=_CAPS_XHIGH_MAX, + fail_reason=( + "claude-fable-5 availability on the CI Vertex project is not yet " + "confirmed for this brand-new release, so this cell stays loud in " + "CI until verified. Remove this fail_reason once the model is " + "confirmed available on the global Vertex endpoint." + ), + ), ModelEntry( alias="vertex-claude-opus-4-8", model="vertex_ai/claude-opus-4-8", @@ -263,6 +296,23 @@ def expected(model: ModelEntry, effort: str) -> CellExpectation: BEDROCK_CONVERSE_MODELS: Tuple[ModelEntry, ...] = ( + ModelEntry( + alias="bedrock-claude-fable-5", + model="bedrock/converse/us.anthropic.claude-fable-5", + mode="adaptive", + extra_params=(("aws_region_name", "us-east-1"),), + required_env=_BEDROCK_REQ, + caps=_CAPS_XHIGH_MAX, + bedrock_effort_ceiling="xhigh", + unavailable_error="is not available for this account", + fail_reason=( + "claude-fable-5 on Bedrock requires the account to opt in to " + "provider data sharing (data retention mode " + "'provider_data_sharing' via the Data Retention API); the CI " + "account has not opted in yet, so this cell stays loud in CI. " + "Remove this fail_reason once the opt-in is done." + ), + ), ModelEntry( alias="bedrock-claude-opus-4-8", model="bedrock/converse/us.anthropic.claude-opus-4-8", diff --git a/tests/llm_translation/reasoning_effort_grid/test_reasoning_effort_grid.py b/tests/llm_translation/reasoning_effort_grid/test_reasoning_effort_grid.py index 551ab8459d15..a5f16f928e51 100644 --- a/tests/llm_translation/reasoning_effort_grid/test_reasoning_effort_grid.py +++ b/tests/llm_translation/reasoning_effort_grid/test_reasoning_effort_grid.py @@ -15,7 +15,6 @@ all_cells, ) - _PROMPT_MESSAGES: List[Dict[str, str]] = [ {"role": "user", "content": "Step by step, calculate 47 * 53. Show your work."} ] @@ -201,8 +200,8 @@ async def test_reasoning_effort_grid( def test_grid_cell_count() -> None: - assert len(_PARAMS) == 25 * 11, ( - f"expected 275 cells (25 provider x model combos x 11 efforts), " + assert len(_PARAMS) == 29 * 11, ( + f"expected 319 cells (29 provider x model combos x 11 efforts), " f"got {len(_PARAMS)}" ) diff --git a/tests/test_litellm/integrations/datadog/test_datadog_team_handler.py b/tests/test_litellm/integrations/datadog/test_datadog_team_handler.py new file mode 100644 index 000000000000..1ee497476d1b --- /dev/null +++ b/tests/test_litellm/integrations/datadog/test_datadog_team_handler.py @@ -0,0 +1,194 @@ +""" +Tests for team-scoped Datadog callback support. + +Verifies that DataDogLogger can be instantiated with per-team credentials +(dd_api_key, dd_site) instead of relying solely on environment variables, +and that the DataDogHandler correctly resolves and caches per-team loggers. +""" + +from unittest.mock import patch + +import pytest + +from litellm.integrations.datadog.datadog import DataDogLogger +from litellm.integrations.datadog.datadog_team_handler import ( + DataDogHandler, + DatadogLoggingConfig, +) +from litellm.litellm_core_utils.specialty_caches.dynamic_logging_cache import ( + DynamicLoggingCache, +) +from litellm.types.utils import StandardCallbackDynamicParams + + +@pytest.fixture +def datadog_env(monkeypatch): + """Set global DD env vars for the default/global logger.""" + monkeypatch.setenv("DD_API_KEY", "global_api_key") + monkeypatch.setenv("DD_SITE", "us1.datadoghq.com") + + +class TestDataDogLoggerCredentialKwargs: + """Test that DataDogLogger accepts credentials as kwargs.""" + + def test_init_with_explicit_credentials(self): + """Logger should use explicit kwargs instead of env vars.""" + with patch("asyncio.create_task"): + logger = DataDogLogger( + dd_api_key="team_api_key", + dd_site="eu1.datadoghq.com", + ) + + assert logger.DD_API_KEY == "team_api_key" + assert "eu1.datadoghq.com" in logger.intake_url + + def test_init_falls_back_to_env_vars(self, datadog_env): + """Logger should fall back to env vars when no kwargs provided.""" + with patch("asyncio.create_task"): + logger = DataDogLogger() + + assert logger.DD_API_KEY == "global_api_key" + assert "us1.datadoghq.com" in logger.intake_url + + def test_init_kwargs_override_env_vars(self, datadog_env): + """Explicit kwargs should take precedence over env vars.""" + with patch("asyncio.create_task"): + logger = DataDogLogger( + dd_api_key="override_key", + dd_site="ap1.datadoghq.com", + ) + + assert logger.DD_API_KEY == "override_key" + assert "ap1.datadoghq.com" in logger.intake_url + + def test_init_with_agent_credentials(self): + """Logger should use agent mode when dd_agent_host is provided.""" + with patch("asyncio.create_task"): + logger = DataDogLogger( + dd_agent_host="dd-agent.local", + dd_agent_port="8125", + dd_api_key="agent_api_key", + ) + + assert "dd-agent.local:8125" in logger.intake_url + assert logger.DD_API_KEY == "agent_api_key" + + def test_init_raises_without_credentials(self, monkeypatch): + """Logger should raise if no credentials are available.""" + monkeypatch.delenv("DD_API_KEY", raising=False) + monkeypatch.delenv("DD_SITE", raising=False) + monkeypatch.delenv("LITELLM_DD_AGENT_HOST", raising=False) + + with pytest.raises(Exception, match="DD_API_KEY"): + with patch("asyncio.create_task"): + DataDogLogger() + + +class TestDataDogHandler: + """Test that DataDogHandler resolves the correct logger per team.""" + + def test_creates_team_logger_with_dynamic_credentials(self, datadog_env): + """Should create a new logger when team credentials are provided.""" + cache = DynamicLoggingCache() + params = StandardCallbackDynamicParams( + dd_api_key="team_a_key", + dd_site="eu1.datadoghq.com", + ) + + with patch("asyncio.create_task"): + result = DataDogHandler.get_datadog_logger_for_request( + standard_callback_dynamic_params=params, + in_memory_dynamic_logger_cache=cache, + ) + + assert result.DD_API_KEY == "team_a_key" + assert "eu1.datadoghq.com" in result.intake_url + + def test_caches_team_logger(self, datadog_env): + """Same team credentials should return the same cached logger instance.""" + cache = DynamicLoggingCache() + params = StandardCallbackDynamicParams( + dd_api_key="team_b_key", + dd_site="us5.datadoghq.com", + ) + + with patch("asyncio.create_task"): + result1 = DataDogHandler.get_datadog_logger_for_request( + standard_callback_dynamic_params=params, + in_memory_dynamic_logger_cache=cache, + ) + result2 = DataDogHandler.get_datadog_logger_for_request( + standard_callback_dynamic_params=params, + in_memory_dynamic_logger_cache=cache, + ) + + assert result1 is result2 + + def test_different_teams_get_different_loggers(self, datadog_env): + """Different team credentials should create separate logger instances.""" + cache = DynamicLoggingCache() + + params_a = StandardCallbackDynamicParams( + dd_api_key="team_a_key", + dd_site="us1.datadoghq.com", + ) + params_b = StandardCallbackDynamicParams( + dd_api_key="team_b_key", + dd_site="eu1.datadoghq.com", + ) + + with patch("asyncio.create_task"): + result_a = DataDogHandler.get_datadog_logger_for_request( + standard_callback_dynamic_params=params_a, + in_memory_dynamic_logger_cache=cache, + ) + result_b = DataDogHandler.get_datadog_logger_for_request( + standard_callback_dynamic_params=params_b, + in_memory_dynamic_logger_cache=cache, + ) + + assert result_a is not result_b + assert result_a.DD_API_KEY == "team_a_key" + assert result_b.DD_API_KEY == "team_b_key" + + def test_request_blocked_callback_params_includes_dd(self): + """DD params should be blocked from request-level metadata (security).""" + from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( + _request_blocked_callback_params, + ) + + assert "dd_api_key" in _request_blocked_callback_params + assert "dd_site" in _request_blocked_callback_params + assert "dd_agent_host" in _request_blocked_callback_params + assert "dd_agent_port" in _request_blocked_callback_params + + +class TestDynamicCredentialDetection: + """Test that _dynamic_datadog_credentials_are_passed works correctly.""" + + def test_no_credentials(self): + params = StandardCallbackDynamicParams() + assert DataDogHandler._dynamic_datadog_credentials_are_passed(params) is False + + def test_dd_api_key_only(self): + params = StandardCallbackDynamicParams(dd_api_key="key") + assert DataDogHandler._dynamic_datadog_credentials_are_passed(params) is True + + def test_dd_site_only(self): + params = StandardCallbackDynamicParams(dd_site="site") + assert DataDogHandler._dynamic_datadog_credentials_are_passed(params) is True + + def test_dd_agent_host_only(self): + params = StandardCallbackDynamicParams(dd_agent_host="host") + assert DataDogHandler._dynamic_datadog_credentials_are_passed(params) is True + + +class TestStandardCallbackDynamicParamsIncludesDatadog: + """Verify that Datadog params are in the allow-list.""" + + def test_dd_params_in_annotations(self): + annotations = StandardCallbackDynamicParams.__annotations__ + assert "dd_api_key" in annotations + assert "dd_site" in annotations + assert "dd_agent_host" in annotations + assert "dd_agent_port" in annotations diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking_dict_safety.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking_dict_safety.py new file mode 100644 index 000000000000..4eee6b59d34f --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking_dict_safety.py @@ -0,0 +1,88 @@ +""" +Tests that the cost-tracking call sites tolerate ``server_tool_use`` being +either a ``dict`` or a ``ServerToolUse`` pydantic instance. + +See https://github.com/BerriAI/litellm/issues/26153. +""" + +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( + StandardBuiltInToolCostTracking, + _get_web_search_requests, +) +from litellm.types.utils import ModelResponse, ServerToolUse, Usage + + +class _UsageWithDictServerToolUse: + """ + Tiny stand-in that mimics the broken streaming-rebuild shape: + ``server_tool_use`` is a plain dict. + """ + + def __init__(self, server_tool_use): + self.server_tool_use = server_tool_use + self.prompt_tokens_details = None + + +def test_get_web_search_requests_handles_none(): + assert _get_web_search_requests(None) is None + + +def test_get_web_search_requests_handles_dict(): + assert _get_web_search_requests({"web_search_requests": 5}) == 5 + + +def test_get_web_search_requests_handles_dict_missing_key(): + assert _get_web_search_requests({}) is None + + +def test_get_web_search_requests_handles_pydantic(): + stu = ServerToolUse(web_search_requests=7) + assert _get_web_search_requests(stu) == 7 + + +def test_get_web_search_requests_handles_pydantic_with_none_value(): + stu = ServerToolUse() + assert _get_web_search_requests(stu) is None + + +def test_response_object_includes_web_search_call_with_dict_server_tool_use(): + """ + The exact bug: ``usage.server_tool_use`` is a dict and the check in + ``response_object_includes_web_search_call`` used to crash with + ``AttributeError``. + """ + response = ModelResponse() + usage = _UsageWithDictServerToolUse({"web_search_requests": 2}) + + # Must not raise — and must correctly detect the web search call. + result = StandardBuiltInToolCostTracking.response_object_includes_web_search_call( + response_object=response, usage=usage # type: ignore[arg-type] + ) + assert result is True + + +def test_response_object_includes_web_search_call_with_pydantic_server_tool_use(): + response = ModelResponse() + usage = _UsageWithDictServerToolUse(ServerToolUse(web_search_requests=2)) + + result = StandardBuiltInToolCostTracking.response_object_includes_web_search_call( + response_object=response, usage=usage # type: ignore[arg-type] + ) + assert result is True + + +def test_response_object_includes_web_search_call_with_none_server_tool_use(): + response = ModelResponse() + usage = _UsageWithDictServerToolUse(None) + + result = StandardBuiltInToolCostTracking.response_object_includes_web_search_call( + response_object=response, usage=usage # type: ignore[arg-type] + ) + assert result is False diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_server_tool_use.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_server_tool_use.py new file mode 100644 index 000000000000..4e28d5ba7d2f --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_server_tool_use.py @@ -0,0 +1,130 @@ +""" +Regression tests for https://github.com/BerriAI/litellm/issues/26153 + +``stream_chunk_builder`` used to leave ``usage.server_tool_use`` as a plain +``dict`` when reconstructing a streaming response. Downstream cost-calculation +code (``StandardBuiltInToolCostTracking.response_object_includes_web_search_call`` +and ``get_cost_for_anthropic_web_search``) accesses +``usage.server_tool_use.web_search_requests`` as an attribute, which raised +``AttributeError: 'dict' object has no attribute 'web_search_requests'``. + +These tests reconstruct streaming chunks for an Anthropic-style web_search +response and assert: + +1. ``stream_chunk_builder`` returns ``ServerToolUse`` (not ``dict``) for + ``usage.server_tool_use``. +2. ``completion_cost`` runs end-to-end on the rebuilt response without + raising ``AttributeError``. +""" + +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm import completion_cost, stream_chunk_builder +from litellm.types.utils import ( + Delta, + ModelResponseStream, + ServerToolUse, + StreamingChoices, + Usage, +) + + +def _make_text_chunk(text: str) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-test-26153", + created=1700000000, + model="claude-3-haiku-20240307", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta(role="assistant", content=text), + ) + ], + ) + + +def _make_finish_chunk_with_usage_dict_server_tool_use() -> ModelResponseStream: + """Final chunk where server_tool_use is a *dict* — reproduces the bug shape.""" + return ModelResponseStream( + id="chatcmpl-test-26153", + created=1700000000, + model="claude-3-haiku-20240307", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(), + ) + ], + usage=Usage( + prompt_tokens=42, + completion_tokens=11, + total_tokens=53, + # NOTE: passed as a dict on purpose — this is the shape that + # historically slipped through stream_chunk_builder unchanged. + server_tool_use={"web_search_requests": 3}, + ), + ) + + +def test_stream_chunk_builder_coerces_server_tool_use_to_pydantic(): + """ + Regression: stream_chunk_builder must produce ServerToolUse, not dict. + """ + chunks = [ + _make_text_chunk("Otters "), + _make_text_chunk("are great."), + _make_finish_chunk_with_usage_dict_server_tool_use(), + ] + + rebuilt = stream_chunk_builder(chunks) + + assert rebuilt is not None + assert rebuilt.usage is not None # type: ignore[attr-defined] + server_tool_use = rebuilt.usage.server_tool_use # type: ignore[attr-defined] + + assert ( + server_tool_use is not None + ), "server_tool_use should be carried through from the final chunk" + assert isinstance(server_tool_use, ServerToolUse), ( + f"expected ServerToolUse, got {type(server_tool_use).__name__}: " + f"{server_tool_use!r}" + ) + # Attribute access must not raise (this is exactly what was broken). + assert server_tool_use.web_search_requests == 3 + + +def test_completion_cost_does_not_raise_on_streaming_web_search_response(): + """ + Regression: completion_cost(...) must not raise AttributeError when the + response was reconstructed by stream_chunk_builder from a streaming + Anthropic web_search call. + """ + chunks = [ + _make_text_chunk("hello"), + _make_finish_chunk_with_usage_dict_server_tool_use(), + ] + + rebuilt = stream_chunk_builder(chunks) + assert rebuilt is not None + + # The exact dollar amount depends on the model-pricing table; what matters + # for this regression is that it does NOT raise AttributeError on + # `dict has no attribute 'web_search_requests'`. + try: + cost = completion_cost(completion_response=rebuilt) + except AttributeError as e: # pragma: no cover - regression guard + pytest.fail( + "completion_cost raised AttributeError after stream_chunk_builder " + f"(issue #26153 regression): {e}" + ) + + assert isinstance(cost, (int, float)) diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py index e40a0817fd94..35aca525f6ce 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -520,7 +520,10 @@ def test_stream_chunk_builder_anthropic_web_search(): assert usage.prompt_tokens == 50 assert usage.completion_tokens == 27 assert usage.total_tokens == 77 - assert usage.server_tool_use["web_search_requests"] == 2 + # server_tool_use must be a ServerToolUse pydantic so downstream cost-calc + # (which uses attribute access) works. See issue #26153. + assert isinstance(usage.server_tool_use, ServerToolUse) + assert usage.server_tool_use.web_search_requests == 2 def test_sort_chunks_handles_dict_hidden_params_created_at(): diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py index 4c3303129305..c91c9c3fdf4d 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -5188,3 +5188,140 @@ def test_client_metadata_stripped_from_anthropic_request(): headers={}, ) assert "client_metadata" not in result + + +@pytest.mark.parametrize( + "model", + ["claude-fable-5", "claude-opus-4-7", "claude-opus-4-8-20260120"], +) +def test_sampling_params_dropped_for_models_that_removed_them(model): + """Fable 5 / Opus 4.7 / 4.8 reject temperature != 1 and any top_p with a + 400; with drop_params set they must be dropped, not forwarded (#30064).""" + config = AnthropicConfig() + + result = config.map_openai_params( + non_default_params={"temperature": 0.5, "top_p": 0.9}, + optional_params={}, + model=model, + drop_params=True, + ) + + assert "temperature" not in result + assert "top_p" not in result + + +@pytest.mark.parametrize("params", [{"temperature": 0.5}, {"top_p": 0.9}, {"top_p": 1}]) +def test_sampling_params_raise_clean_error_without_drop_params(params, monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + config = AnthropicConfig() + + with pytest.raises(litellm.utils.UnsupportedParamsError, match="drop_params"): + config.map_openai_params( + non_default_params=params, + optional_params={}, + model="claude-fable-5", + drop_params=False, + ) + + +def test_temperature_1_forwarded_on_models_that_removed_sampling_params(): + """temperature=1 (the API default) is still accepted and must pass through.""" + config = AnthropicConfig() + + result = config.map_openai_params( + non_default_params={"temperature": 1}, + optional_params={}, + model="claude-fable-5", + drop_params=False, + ) + + assert result["temperature"] == 1 + + +@pytest.mark.parametrize("model", ["claude-opus-4-6", "claude-sonnet-4-6"]) +def test_sampling_params_forwarded_on_models_that_accept_them(model): + config = AnthropicConfig() + + result = config.map_openai_params( + non_default_params={"temperature": 0.5, "top_p": 0.9}, + optional_params={}, + model=model, + drop_params=True, + ) + + assert result["temperature"] == 0.5 + assert result["top_p"] == 0.9 + + +def test_sampling_param_gating_driven_by_model_map_flag(monkeypatch): + """The drop/raise decision must come from ``supports_sampling_params`` in + the model map, not just name matching: a flagged entry gates a model whose + name says nothing, and an explicit ``true`` overrides the name fallback.""" + monkeypatch.setitem( + litellm.model_cost, "claude-zeta-9", {"supports_sampling_params": False} + ) + monkeypatch.setitem( + litellm.model_cost, "claude-fable-5-test", {"supports_sampling_params": True} + ) + config = AnthropicConfig() + + flagged_off = config.map_openai_params( + non_default_params={"top_p": 0.9}, + optional_params={}, + model="claude-zeta-9", + drop_params=True, + ) + assert "top_p" not in flagged_off + + flagged_on = config.map_openai_params( + non_default_params={"top_p": 0.9}, + optional_params={}, + model="claude-fable-5-test", + drop_params=True, + ) + assert flagged_on["top_p"] == 0.9 + + +def test_top_k_dropped_at_transform_for_models_that_removed_it(): + """``top_k`` is a provider-specific kwarg that bypasses + ``map_openai_params``, so it must be stripped at the transform_request + boundary shared by the direct, invoke, Vertex, and Azure paths (#30064).""" + config = AnthropicConfig() + + result = config.transform_request( + model="claude-fable-5", + messages=[{"role": "user", "content": "hello"}], + optional_params={"max_tokens": 10, "top_k": 40}, + litellm_params={"drop_params": True}, + headers={}, + ) + + assert "top_k" not in result + + +def test_top_k_raises_at_transform_without_drop_params(monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + config = AnthropicConfig() + + with pytest.raises(litellm.utils.UnsupportedParamsError, match="drop_params"): + config.transform_request( + model="claude-fable-5", + messages=[{"role": "user", "content": "hello"}], + optional_params={"max_tokens": 10, "top_k": 40}, + litellm_params={}, + headers={}, + ) + + +def test_top_k_forwarded_at_transform_on_models_that_accept_it(): + config = AnthropicConfig() + + result = config.transform_request( + model="claude-sonnet-4-6", + messages=[{"role": "user", "content": "hello"}], + optional_params={"max_tokens": 10, "top_k": 40}, + litellm_params={"drop_params": True}, + headers={}, + ) + + assert result["top_k"] == 40 diff --git a/tests/test_litellm/llms/anthropic/test_cost_calculation_dict_safety.py b/tests/test_litellm/llms/anthropic/test_cost_calculation_dict_safety.py new file mode 100644 index 000000000000..70fef0162e62 --- /dev/null +++ b/tests/test_litellm/llms/anthropic/test_cost_calculation_dict_safety.py @@ -0,0 +1,94 @@ +""" +Tests that ``get_cost_for_anthropic_web_search`` tolerates ``server_tool_use`` +being either a ``dict`` or a ``ServerToolUse`` pydantic instance. + +See https://github.com/BerriAI/litellm/issues/26153. +""" + +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.llms.anthropic.cost_calculation import ( + _get_web_search_requests, + get_cost_for_anthropic_web_search, +) +from litellm.types.utils import ModelInfo, ServerToolUse + + +class _UsageWithServerToolUse: + def __init__(self, server_tool_use): + self.server_tool_use = server_tool_use + + +def _make_model_info(cost_per_query: float = 0.01) -> ModelInfo: + info: ModelInfo = { # type: ignore[typeddict-item] + "search_context_cost_per_query": { + "search_context_size_low": cost_per_query, + "search_context_size_medium": cost_per_query, + "search_context_size_high": cost_per_query, + } + } + return info + + +def test_get_web_search_requests_handles_none(): + assert _get_web_search_requests(None) is None + + +def test_get_web_search_requests_handles_dict(): + assert _get_web_search_requests({"web_search_requests": 4}) == 4 + + +def test_get_web_search_requests_handles_dict_missing_key(): + assert _get_web_search_requests({}) is None + + +def test_get_web_search_requests_handles_pydantic(): + assert _get_web_search_requests(ServerToolUse(web_search_requests=2)) == 2 + + +def test_get_cost_for_anthropic_web_search_with_dict_server_tool_use(): + """ + Regression: ``server_tool_use`` was a dict from ``stream_chunk_builder`` and + direct attribute access on it raised ``AttributeError``. + """ + usage = _UsageWithServerToolUse({"web_search_requests": 3}) + info = _make_model_info(cost_per_query=0.01) + + cost = get_cost_for_anthropic_web_search( + model_info=info, usage=usage # type: ignore[arg-type] + ) + + assert cost == pytest.approx(0.03) + + +def test_get_cost_for_anthropic_web_search_with_pydantic_server_tool_use(): + usage = _UsageWithServerToolUse(ServerToolUse(web_search_requests=3)) + info = _make_model_info(cost_per_query=0.01) + + cost = get_cost_for_anthropic_web_search( + model_info=info, usage=usage # type: ignore[arg-type] + ) + + assert cost == pytest.approx(0.03) + + +def test_get_cost_for_anthropic_web_search_with_none_server_tool_use(): + usage = _UsageWithServerToolUse(None) + info = _make_model_info(cost_per_query=0.01) + + cost = get_cost_for_anthropic_web_search( + model_info=info, usage=usage # type: ignore[arg-type] + ) + + assert cost == 0.0 + + +def test_get_cost_for_anthropic_web_search_with_no_usage(): + info = _make_model_info(cost_per_query=0.01) + cost = get_cost_for_anthropic_web_search(model_info=info, usage=None) + assert cost == 0.0 diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index a6aa35ee6d12..fec215e5c440 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -5268,3 +5268,122 @@ def text(self): msg = str(exc_info.value) assert "secret content" not in msg assert "Error converting to valid response block" in msg + + +def test_converse_drops_sampling_params_for_models_that_removed_them(): + """Fable 5 / Opus 4.7 / 4.8 reject temperature != 1 and any top_p; with + drop_params set, converse must drop them instead of forwarding (#30064).""" + config = AmazonConverseConfig() + + result = config.map_openai_params( + non_default_params={"temperature": 0.5, "top_p": 0.9}, + optional_params={}, + model="us.anthropic.claude-fable-5", + drop_params=True, + ) + + assert "temperature" not in result + assert "topP" not in result + + +def test_converse_sampling_params_raise_without_drop_params(monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + config = AmazonConverseConfig() + + with pytest.raises(litellm.utils.UnsupportedParamsError, match="drop_params"): + config.map_openai_params( + non_default_params={"temperature": 0.5}, + optional_params={}, + model="global.anthropic.claude-opus-4-8-v1:0", + drop_params=False, + ) + + +def test_converse_sampling_params_forwarded_on_models_that_accept_them(): + config = AmazonConverseConfig() + + result = config.map_openai_params( + non_default_params={"temperature": 0.5, "top_p": 0.9}, + optional_params={}, + model="us.anthropic.claude-sonnet-4-6", + drop_params=True, + ) + + assert result["temperature"] == 0.5 + assert result["topP"] == 0.9 + + +def test_converse_top_k_dropped_for_models_that_removed_it(): + """``top_k`` reaches converse as a provider-specific kwarg destined for + ``additionalModelRequestFields``, bypassing ``map_openai_params``; the + transform must strip it for models that removed sampling params (#30064).""" + config = AmazonConverseConfig() + + result = config.transform_request( + model="us.anthropic.claude-fable-5", + messages=[{"role": "user", "content": "hello"}], + optional_params={"top_k": 40}, + litellm_params={"drop_params": True}, + headers={}, + ) + + assert "top_k" not in result.get("additionalModelRequestFields", {}) + + +def test_converse_top_k_raises_without_drop_params(monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + config = AmazonConverseConfig() + + with pytest.raises(litellm.utils.UnsupportedParamsError, match="drop_params"): + config.transform_request( + model="us.anthropic.claude-fable-5", + messages=[{"role": "user", "content": "hello"}], + optional_params={"top_k": 40}, + litellm_params={}, + headers={}, + ) + + +def test_converse_top_k_forwarded_on_models_that_accept_it(): + config = AmazonConverseConfig() + + result = config.transform_request( + model="us.anthropic.claude-sonnet-4-6", + messages=[{"role": "user", "content": "hello"}], + optional_params={"top_k": 40}, + litellm_params={"drop_params": True}, + headers={}, + ) + + assert result["additionalModelRequestFields"]["top_k"] == 40 + + +def test_converse_top_k_zero_raises_without_drop_params(monkeypatch): + """``top_k=0`` must hit the same gating as any other value; previously the + truthiness check let it silently disappear on models that removed sampling + params, diverging from the Anthropic boundary that treats ``0`` as present.""" + monkeypatch.setattr(litellm, "drop_params", False) + config = AmazonConverseConfig() + + with pytest.raises(litellm.utils.UnsupportedParamsError, match="drop_params"): + config.transform_request( + model="us.anthropic.claude-fable-5", + messages=[{"role": "user", "content": "hello"}], + optional_params={"top_k": 0}, + litellm_params={}, + headers={}, + ) + + +def test_converse_top_k_zero_forwarded_on_models_that_accept_it(): + config = AmazonConverseConfig() + + result = config.transform_request( + model="us.anthropic.claude-sonnet-4-6", + messages=[{"role": "user", "content": "hello"}], + optional_params={"top_k": 0}, + litellm_params={"drop_params": True}, + headers={}, + ) + + assert result["additionalModelRequestFields"]["top_k"] == 0 diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py index e2133d56f897..92b5ca7b10bc 100644 --- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py +++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py @@ -12,6 +12,11 @@ sys.path.insert(0, os.path.abspath("../../../../..")) import pytest +from botocore.exceptions import ( + ConnectTimeoutError, + PartialCredentialsError, + ProfileNotFound, +) import litellm from litellm.llms.bedrock_mantle.responses.transformation import ( @@ -114,16 +119,15 @@ def test_bedrock_bearer_token_fallback(self, monkeypatch): ) assert headers["Authorization"] == "Bearer bearer-key" - def test_missing_key_raises(self, monkeypatch): + def test_missing_bearer_does_not_raise_in_validate_environment(self, monkeypatch): + # SigV4 may still apply, so validate_environment must defer instead of raising. monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) cfg = BedrockMantleResponsesAPIConfig() - with pytest.raises(ValueError, match="Bedrock Mantle API key"): - cfg.validate_environment( - headers={}, - model="openai.gpt-5.5", - litellm_params=GenericLiteLLMParams(), - ) + headers = cfg.validate_environment( + headers={}, model="openai.gpt-5.5", litellm_params=GenericLiteLLMParams() + ) + assert "Authorization" not in headers def test_custom_llm_provider(self): cfg = BedrockMantleResponsesAPIConfig() @@ -261,6 +265,386 @@ def local_cost_map(monkeypatch): litellm.get_model_info.cache_clear() +class TestBedrockMantleResponsesSigV4: + def test_bearer_short_circuits_without_credentials(self, monkeypatch): + from unittest.mock import MagicMock + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + + signer = BaseAWSLLM() + signer.get_credentials = MagicMock( + side_effect=AssertionError("get_credentials must not run for bearer auth") + ) + cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer) + + headers, signed_body = cfg.sign_request( + headers={}, + optional_params={}, + request_data={"input": "hi"}, + api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses", + api_key="bearer-from-config", + ) + assert headers["Authorization"] == "Bearer bearer-from-config" + assert signed_body == b'{"input": "hi"}' + signer.get_credentials.assert_not_called() + + def test_bearer_resolved_from_mantle_env_key(self, monkeypatch): + from unittest.mock import MagicMock + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.setenv("BEDROCK_MANTLE_API_KEY", "env-bearer") + + signer = BaseAWSLLM() + signer.get_credentials = MagicMock( + side_effect=AssertionError("get_credentials must not run for bearer auth") + ) + cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer) + + headers, _ = cfg.sign_request( + headers={}, + optional_params={}, + request_data={"input": "hi"}, + api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses", + api_key=None, + ) + assert headers["Authorization"] == "Bearer env-bearer" + + def test_bearer_arg_takes_priority_over_mantle_env_key(self, monkeypatch): + # The passed api_key (e.g. litellm_params.api_key) must win over the env + # bearer; a reordered precedence chain would silently use the wrong token. + from unittest.mock import MagicMock + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.setenv("BEDROCK_MANTLE_API_KEY", "env-bearer") + + signer = BaseAWSLLM() + signer.get_credentials = MagicMock( + side_effect=AssertionError("get_credentials must not run for bearer auth") + ) + cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer) + + headers, _ = cfg.sign_request( + headers={}, + optional_params={}, + request_data={"input": "hi"}, + api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses", + api_key="arg-bearer", + ) + assert headers["Authorization"] == "Bearer arg-bearer" + signer.get_credentials.assert_not_called() + + def test_access_key_produces_sigv4_headers(self, monkeypatch): + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + + cfg = BedrockMantleResponsesAPIConfig(aws_signer=BaseAWSLLM()) + headers, signed_body = cfg.sign_request( + headers={}, + optional_params={ + "aws_access_key_id": "AKIAEXAMPLE", + "aws_secret_access_key": "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0", + "aws_session_token": "session-token-test", + "aws_region_name": "us-east-2", + }, + request_data={"input": "hi"}, + api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses", + api_key=None, + ) + assert headers["Authorization"].startswith("AWS4-HMAC-SHA256") + assert "Credential=AKIAEXAMPLE/" in headers["Authorization"] + assert "/us-east-2/bedrock/aws4_request" in headers["Authorization"] + assert "X-Amz-Date" in headers + assert headers["X-Amz-Security-Token"] == "session-token-test" + assert signed_body == b'{"input": "hi"}' + + def test_assume_role_path_produces_sigv4_headers(self, monkeypatch): + from unittest.mock import MagicMock + from botocore.credentials import Credentials + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + + signer = BaseAWSLLM() + signer.get_credentials = MagicMock( + return_value=Credentials( + access_key="ASIAEXAMPLE", + secret_key="YXNzdW1lZC1yb2xlLXNlY3JldC1hc3N1bWVk", + token="assumed-session-token", + ) + ) + cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer) + + headers, _ = cfg.sign_request( + headers={}, + optional_params={ + "aws_role_name": "arn:aws:iam::000000000000:role/test-role", + "aws_session_name": "litellm-test", + "aws_region_name": "us-east-2", + }, + request_data={"input": "hi"}, + api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses", + api_key=None, + ) + signer.get_credentials.assert_called_once() + call = signer.get_credentials.call_args.kwargs + assert call["aws_role_name"] == "arn:aws:iam::000000000000:role/test-role" + assert call["aws_session_name"] == "litellm-test" + assert headers["Authorization"].startswith("AWS4-HMAC-SHA256") + assert "/us-east-2/bedrock/aws4_request" in headers["Authorization"] + + def test_signed_body_matches_final_data_after_normalize(self, monkeypatch): + """Core regression: the signed bytes must equal the bytes actually sent. + + Sign the *final* data dict and assert the returned signed_body decodes to + exactly that dict, so a later change to the data would break the SigV4 hash. + """ + import json + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + + final_data = {"model": "openai.gpt-5.5", "input": "hi", "max_output_tokens": 16} + cfg = BedrockMantleResponsesAPIConfig(aws_signer=BaseAWSLLM()) + _, signed_body = cfg.sign_request( + headers={}, + optional_params={ + "aws_access_key_id": "AKIAEXAMPLE", + "aws_secret_access_key": "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0", + "aws_region_name": "us-east-2", + }, + request_data=final_data, + api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses", + api_key=None, + ) + assert signed_body is not None + assert json.loads(signed_body) == final_data + + def test_region_comes_from_optional_params(self, monkeypatch): + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + monkeypatch.delenv("AWS_REGION", raising=False) + monkeypatch.delenv("AWS_REGION_NAME", raising=False) + + cfg = BedrockMantleResponsesAPIConfig(aws_signer=BaseAWSLLM()) + headers, _ = cfg.sign_request( + headers={}, + optional_params={ + "aws_access_key_id": "AKIAEXAMPLE", + "aws_secret_access_key": "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0", + "aws_region_name": "eu-west-1", + }, + request_data={"input": "hi"}, + api_base="https://bedrock-mantle.eu-west-1.api.aws/openai/v1/responses", + api_key=None, + ) + assert "/eu-west-1/bedrock/aws4_request" in headers["Authorization"] + + def test_url_region_and_sigv4_region_agree_from_litellm_params(self, monkeypatch): + """Adversarial-review regression: a caller-supplied aws_region_name (no region + env set) must shape BOTH the URL host and the SigV4 credential scope, or the + request is signed for one region and sent to another -> 401. + """ + monkeypatch.delenv("BEDROCK_MANTLE_REGION", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False) + monkeypatch.delenv("AWS_REGION", raising=False) + monkeypatch.delenv("AWS_REGION_NAME", raising=False) + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + params = { + "aws_region_name": "ap-southeast-2", + "aws_access_key_id": "AKIAEXAMPLE", + "aws_secret_access_key": "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0", + } + cfg = BedrockMantleResponsesAPIConfig(aws_signer=BaseAWSLLM()) + url = cfg.get_complete_url(api_base=None, litellm_params=params) + assert ( + url == "https://bedrock-mantle.ap-southeast-2.api.aws/openai/v1/responses" + ) + + headers, _ = cfg.sign_request( + headers={}, + optional_params=params, + request_data={"input": "hi"}, + api_base=url, + api_key=None, + ) + assert "/ap-southeast-2/bedrock/aws4_request" in headers["Authorization"] + + def test_injected_default_region_base_does_not_override_aws_region_name( + self, monkeypatch + ): + """2nd-round adversarial regression: responses/main.py auto-injects + litellm_params.api_base = https://bedrock-mantle..api.aws/v1 (default + region, ignoring aws_region_name). The config must still pin BOTH the URL host + and the SigV4 scope to aws_region_name, or the IAM deployment 401s. A naive + 'resolve region only when api_base is None' fix would fail this test. + """ + monkeypatch.delenv("BEDROCK_MANTLE_REGION", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False) + monkeypatch.delenv("AWS_REGION", raising=False) + monkeypatch.delenv("AWS_REGION_NAME", raising=False) + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + injected_base = "https://bedrock-mantle.us-east-1.api.aws/v1" # default region + params = { + "aws_region_name": "us-east-2", # what the caller actually wants + "api_base": injected_base, + "aws_access_key_id": "AKIAEXAMPLE", + "aws_secret_access_key": "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0", + } + cfg = BedrockMantleResponsesAPIConfig(aws_signer=BaseAWSLLM()) + url = cfg.get_complete_url(api_base=injected_base, litellm_params=params) + assert url == "https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses" + + headers, _ = cfg.sign_request( + headers={}, + optional_params=params, + request_data={"input": "hi"}, + api_base=url, + api_key=None, + ) + assert "/us-east-2/bedrock/aws4_request" in headers["Authorization"] + assert "us-east-1" not in headers["Authorization"] + + def test_custom_proxy_host_is_preserved(self, monkeypatch): + """A genuinely custom (non-Mantle) api_base host must be preserved, not rewritten + to a bedrock-mantle host. Only standard Mantle hosts are region-pinned. + """ + monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False) + cfg = BedrockMantleResponsesAPIConfig() + url = cfg.get_complete_url( + api_base="https://mantle-proxy.internal.example/openai/v1", + litellm_params={"aws_region_name": "us-east-2"}, + ) + assert url == "https://mantle-proxy.internal.example/openai/v1/responses" + + def test_caller_authorization_does_not_override_sigv4(self, monkeypatch): + """Adversarial-review regression: a caller-supplied Authorization header (e.g. + from extra_headers, surviving the relaxed validate_environment) must not clobber + the SigV4 Authorization that _sign_request would otherwise restore. + """ + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + cfg = BedrockMantleResponsesAPIConfig(aws_signer=BaseAWSLLM()) + headers, _ = cfg.sign_request( + headers={"Authorization": "Bearer stale-caller-token"}, + optional_params={ + "aws_access_key_id": "AKIAEXAMPLE", + "aws_secret_access_key": "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0", + "aws_region_name": "us-east-2", + }, + request_data={"input": "hi"}, + api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses", + api_key=None, + ) + assert headers["Authorization"].startswith("AWS4-HMAC-SHA256") + assert "Bearer stale-caller-token" not in headers["Authorization"] + + def test_no_bearer_and_no_credentials_raises_both_paths(self, monkeypatch): + from unittest.mock import MagicMock + from botocore.exceptions import NoCredentialsError + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + + signer = BaseAWSLLM() + signer.get_credentials = MagicMock(side_effect=NoCredentialsError()) + cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer) + + with pytest.raises(ValueError) as exc: + cfg.sign_request( + headers={}, + optional_params={"aws_region_name": "us-east-2"}, + request_data={"input": "hi"}, + api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses", + api_key=None, + ) + msg = str(exc.value) + assert "Bearer" in msg + assert "SigV4" in msg or "IAM" in msg + + @pytest.mark.parametrize( + "cred_error", + [ + PartialCredentialsError(provider="env", cred_var="aws_secret_access_key"), + ProfileNotFound(profile="missing-profile"), + ], + ) + def test_partial_credentials_raises_both_paths(self, monkeypatch, cred_error): + from unittest.mock import MagicMock + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + + signer = BaseAWSLLM() + signer.get_credentials = MagicMock(side_effect=cred_error) + cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer) + + with pytest.raises(ValueError) as exc: + cfg.sign_request( + headers={}, + optional_params={"aws_region_name": "us-east-2"}, + request_data={"input": "hi"}, + api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses", + api_key=None, + ) + msg = str(exc.value) + assert "Bearer" in msg + assert "SigV4" in msg or "IAM" in msg + + def test_sts_transport_error_is_not_masked_as_credentials(self, monkeypatch): + # An AssumeRole / web-identity flow hits STS over the network, so a transient + # connection error must surface as itself, not be rewritten into the + # "no usable AWS credentials" message that would send the user to fix the + # wrong thing. + from unittest.mock import MagicMock + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + + signer = BaseAWSLLM() + signer.get_credentials = MagicMock( + side_effect=ConnectTimeoutError( + endpoint_url="https://sts.us-east-2.amazonaws.com" + ) + ) + cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer) + + with pytest.raises(ConnectTimeoutError): + cfg.sign_request( + headers={}, + optional_params={ + "aws_role_name": "arn:aws:iam::000000000000:role/test-role", + "aws_region_name": "us-east-2", + }, + request_data={"input": "hi"}, + api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses", + api_key=None, + ) + + class TestBedrockMantleResponsesPricing: def test_gpt_5_5_pricing_and_mode(self, local_cost_map): info = litellm.get_model_info("bedrock_mantle/openai.gpt-5.5") diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index 279e9730e69f..7321abcee461 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -742,3 +742,241 @@ def _mutate(e, request_data): assert first_sent == prebuilt # attempt 0 used prebuilt assert second_sent == _json.dumps(request_body) # attempt 1 re-serialized assert "MUTATED" in second_sent # ... the mutated body + + +def test_base_responses_config_sign_request_is_noop_by_default(): + """Default responses sign_request must be a no-op: unchanged headers, no signed body. + + Guards the 15 existing responses providers from accidental signing when the + handler starts calling sign_request. + """ + from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig + + cfg = OpenAIResponsesAPIConfig() + headers = {"Authorization": "Bearer sk-existing"} + out_headers, signed_body = cfg.sign_request( + headers=headers, + optional_params={}, + request_data={"input": "hi"}, + api_base="https://api.openai.com/v1/responses", + ) + assert out_headers == {"Authorization": "Bearer sk-existing"} + assert signed_body is None + + +def _make_responses_handler_call(signed_body): + """Drive BaseLLMHTTPHandler.response_api_handler with a fully mocked provider + config + sync client, returning the kwargs the client.post was called with. + + signed_body=None simulates a no-op (non-signing) provider; bytes simulates a + signing provider (e.g. Bedrock Mantle). + """ + from unittest.mock import MagicMock + from litellm.llms.custom_httpx.http_handler import HTTPHandler + from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler + from litellm.types.router import GenericLiteLLMParams + + provider_config = MagicMock() + provider_config.validate_environment.return_value = {} + provider_config.get_complete_url.return_value = ( + "https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses" + ) + provider_config.transform_responses_api_request.return_value = {"input": "hi"} + provider_config.should_fake_stream.return_value = False + provider_config.sign_request.return_value = ({"X-Signed": "1"}, signed_body) + + mock_client = MagicMock(spec=HTTPHandler) + mock_client.post.return_value = MagicMock() + + handler = BaseLLMHTTPHandler() + handler.response_api_handler( + model="openai.gpt-5.5", + input="hi", + responses_api_provider_config=provider_config, + response_api_optional_request_params={}, + custom_llm_provider="bedrock_mantle", + litellm_params=GenericLiteLLMParams(aws_region_name="us-east-2"), + logging_obj=MagicMock(), + client=mock_client, + _is_async=False, + ) + return mock_client.post.call_args.kwargs + + +def test_responses_handler_sends_json_when_not_signed(): + """No-op provider (signed_body is None) -> handler posts json=data, no data= bytes.""" + kwargs = _make_responses_handler_call(signed_body=None) + assert kwargs.get("json") == {"input": "hi"} + assert "data" not in kwargs + + +def test_responses_handler_sends_signed_bytes_when_signed(): + """Signing provider -> handler posts the exact signed bytes via data=, not json=.""" + kwargs = _make_responses_handler_call(signed_body=b'{"input": "hi"}') + assert kwargs.get("data") == b'{"input": "hi"}' + assert "json" not in kwargs + assert kwargs["headers"] == {"X-Signed": "1"} + + +def test_responses_handler_signs_after_fake_stream_prep_strips_stream(): + """Fake-stream signing-order invariant: the bytes SIGNED must equal the bytes SENT. + + In the streaming + fake-stream path the handler first runs + _prepare_fake_stream_request, which pops "stream" out of the body, and only + then calls sign_request. If signing ran before that pop, the signed body + would still carry "stream" while the body sent over the wire would not, + producing a SigV4 payload-hash mismatch (401) for a real Mantle deployment. + We snapshot request_data at sign time and assert "stream" is already gone. + """ + from unittest.mock import MagicMock + from litellm.llms.custom_httpx.http_handler import HTTPHandler + from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler + from litellm.types.llms.openai import ResponsesAPIResponse + from litellm.types.router import GenericLiteLLMParams + + provider_config = MagicMock() + provider_config.validate_environment.return_value = {} + provider_config.get_complete_url.return_value = ( + "https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses" + ) + provider_config.transform_responses_api_request.return_value = { + "input": "hi", + "stream": True, + } + provider_config.should_fake_stream.return_value = True + provider_config.transform_response_api_response.return_value = ResponsesAPIResponse( + id="resp_1", + created_at=0, + output=[], + status="completed", + model="openai.gpt-5.5", + ) + + captured = {} + + def _capture_sign(**kwargs): + captured["request_data"] = dict(kwargs["request_data"]) + return ({"X-Signed": "1"}, b'{"input": "hi"}') + + provider_config.sign_request.side_effect = _capture_sign + + mock_client = MagicMock(spec=HTTPHandler) + mock_client.post.return_value = MagicMock() + + handler = BaseLLMHTTPHandler() + handler.response_api_handler( + model="openai.gpt-5.5", + input="hi", + responses_api_provider_config=provider_config, + response_api_optional_request_params={"stream": True}, + custom_llm_provider="bedrock_mantle", + litellm_params=GenericLiteLLMParams(aws_region_name="us-east-2"), + logging_obj=MagicMock(), + client=mock_client, + _is_async=False, + fake_stream=True, + ) + + assert "stream" not in captured["request_data"] + assert "input" in captured["request_data"] + + post_kwargs = mock_client.post.call_args.kwargs + assert post_kwargs.get("data") == b'{"input": "hi"}' + assert "json" not in post_kwargs + assert "stream" in post_kwargs + + +def _make_compact_handler_call(signed_body, is_async): + """Drive (async_)compact_response_api_handler with a fully mocked provider config + + client, returning the kwargs the client.post was called with. + + signed_body=None simulates a no-op (non-signing) provider; bytes simulates a + signing provider (e.g. Bedrock Mantle SigV4 / bearer). + """ + from unittest.mock import MagicMock + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler + from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler + from litellm.types.router import GenericLiteLLMParams + + compact_url = "https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses/compact" + provider_config = MagicMock() + provider_config.validate_environment.return_value = {} + provider_config.get_complete_url.return_value = ( + "https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses" + ) + provider_config.transform_compact_response_api_request.return_value = ( + compact_url, + {"model": "openai.gpt-5.5", "input": "hi"}, + ) + provider_config.sign_request.return_value = ({"X-Signed": "1"}, signed_body) + provider_config.transform_compact_response_api_response.return_value = "ok" + + spec = AsyncHTTPHandler if is_async else HTTPHandler + mock_client = MagicMock(spec=spec) + if is_async: + mock_client.post = AsyncMock(return_value=MagicMock()) + else: + mock_client.post.return_value = MagicMock() + + handler = BaseLLMHTTPHandler() + result = handler.compact_response_api_handler( + model="openai.gpt-5.5", + input="hi", + responses_api_provider_config=provider_config, + response_api_optional_request_params={}, + custom_llm_provider="bedrock_mantle", + litellm_params=GenericLiteLLMParams(aws_region_name="us-east-2"), + logging_obj=MagicMock(), + client=mock_client, + _is_async=is_async, + ) + if is_async: + asyncio.run(result) + return provider_config, mock_client.post.call_args.kwargs + + +def test_compact_handler_sends_json_when_not_signed(): + """No-op provider on compact (signed_body is None) -> posts json=data, no data= bytes.""" + provider_config, kwargs = _make_compact_handler_call( + signed_body=None, is_async=False + ) + provider_config.sign_request.assert_called_once() + assert kwargs.get("json") == {"model": "openai.gpt-5.5", "input": "hi"} + assert "data" not in kwargs + + +def test_compact_handler_sends_signed_bytes_when_signed(): + """Signing provider on compact -> posts the signed bytes via data=, not json=. + + Regression for the adversarial-review finding that /responses/compact bypassed + the SigV4 signing hook, so IAM-only Mantle callers sent unsigned bodies. + """ + provider_config, kwargs = _make_compact_handler_call( + signed_body=b'{"model": "openai.gpt-5.5", "input": "hi"}', is_async=False + ) + assert kwargs.get("data") == b'{"model": "openai.gpt-5.5", "input": "hi"}' + assert "json" not in kwargs + assert kwargs["headers"] == {"X-Signed": "1"} + # signing must use the compact endpoint as api_base, not the create URL + assert provider_config.sign_request.call_args.kwargs["api_base"].endswith( + "/openai/v1/responses/compact" + ) + + +def test_async_compact_handler_sends_signed_bytes_when_signed(): + """Async compact must sign identically to sync (same omission in the async twin).""" + provider_config, kwargs = _make_compact_handler_call( + signed_body=b'{"model": "openai.gpt-5.5", "input": "hi"}', is_async=True + ) + assert kwargs.get("data") == b'{"model": "openai.gpt-5.5", "input": "hi"}' + assert "json" not in kwargs + assert kwargs["headers"] == {"X-Signed": "1"} + + +def test_async_compact_handler_sends_json_when_not_signed(): + """Async no-op provider on compact -> posts json=data, no data= bytes.""" + _provider_config, kwargs = _make_compact_handler_call( + signed_body=None, is_async=True + ) + assert kwargs.get("json") == {"model": "openai.gpt-5.5", "input": "hi"} + assert "data" not in kwargs diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 7753378ab4f0..ab42ee1e979f 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -658,12 +658,11 @@ class TestMCPOAuth2AuthFlow: async def test_oauth2_token_in_authorization_header_fallback(self): """ - When only Authorization header is present with a non-LiteLLM OAuth2 token - AND the target server is operator-configured for ``auth_type=oauth2``, - auth should fall back to permissive mode (OAuth2 passthrough). + When only the Authorization header is present with a non-LiteLLM OAuth2 + token AND the target server delegates auth to upstream, LiteLLM skips its + own validation entirely (so the upstream token is never mistaken for a + virtual key) and forwards the bearer upstream. """ - from fastapi import HTTPException - from litellm.types.mcp import MCPAuth scope = { @@ -675,17 +674,16 @@ async def test_oauth2_token_in_authorization_header_fallback(self): ], } - async def mock_user_api_key_auth_fails(api_key, request): - raise HTTPException(status_code=401, detail="Invalid API key") - oauth2_server = MagicMock() oauth2_server.auth_type = MCPAuth.oauth2 + oauth2_server.delegate_auth_to_upstream = True + oauth2_server.has_client_credentials = False with ( patch( "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", - side_effect=mock_user_api_key_auth_fails, - ), + new_callable=AsyncMock, + ) as mock_auth, patch( "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager" ) as mock_mgr, @@ -700,10 +698,10 @@ async def mock_user_api_key_auth_fails(api_key, request): raw_headers, ) = await MCPRequestHandler.process_mcp_request(scope) - # Should succeed with default UserAPIKeyAuth (OAuth2 fallback) - assert auth_result is not None assert isinstance(auth_result, UserAPIKeyAuth) - # OAuth2 headers should contain the token for upstream forwarding + # The upstream token is never validated as a LiteLLM key ... + mock_auth.assert_not_called() + # ... and is preserved for upstream forwarding. assert ( oauth2_headers.get("Authorization") == "Bearer atlassian-oauth2-access-token-xyz" @@ -813,11 +811,12 @@ async def mock_user_api_key_auth_server_error(api_key, request): await MCPRequestHandler.process_mcp_request(scope) assert exc_info.value.status_code == 500 - async def test_proxy_exception_oauth2_fallback(self): + async def test_proxy_exception_non_delegate_oauth2_propagates(self): """ - user_api_key_auth raises ProxyException (not HTTPException) in production. - The OAuth2 fallback must catch ProxyException with code 401/403 too, - but only when the target server is operator-configured for ``auth_type=oauth2``. + Production raises ProxyException (not HTTPException) on auth failure. For + a non-delegate oauth2 server the bearer is treated as a LiteLLM credential + and a 401 must propagate as a real auth error, not be exchanged for an + anonymous upstream-passthrough session. """ from litellm.proxy._types import ProxyException from litellm.types.mcp import MCPAuth @@ -841,6 +840,8 @@ async def mock_user_api_key_auth_proxy_exception(api_key, request): oauth2_server = MagicMock() oauth2_server.auth_type = MCPAuth.oauth2 + oauth2_server.delegate_auth_to_upstream = False + oauth2_server.is_oauth_passthrough = False with ( patch( @@ -852,22 +853,9 @@ async def mock_user_api_key_auth_proxy_exception(api_key, request): ) as mock_mgr, ): mock_mgr.get_mcp_server_by_name.return_value = oauth2_server - ( - auth_result, - mcp_auth_header, - mcp_servers, - mcp_server_auth_headers, - oauth2_headers, - raw_headers, - ) = await MCPRequestHandler.process_mcp_request(scope) - - # Should succeed with default UserAPIKeyAuth (OAuth2 fallback) - assert auth_result is not None - assert isinstance(auth_result, UserAPIKeyAuth) - assert ( - oauth2_headers.get("Authorization") - == "Bearer atlassian-oauth2-access-token-xyz" - ) + with pytest.raises(ProxyException) as exc_info: + await MCPRequestHandler.process_mcp_request(scope) + assert str(exc_info.value.code) == "401" async def test_proxy_exception_non_auth_still_raises(self): """ @@ -1355,11 +1343,15 @@ async def mock_user_api_key_auth_fails(api_key, request): await MCPRequestHandler.process_mcp_request(scope) assert exc_info.value.status_code == 401 - async def test_fallback_allowed_when_target_is_oauth2_mode(self): + async def test_non_delegate_oauth2_does_not_fall_back_to_anonymous(self): """ - Operator-configured OAuth2 passthrough still works: target server has - ``auth_type=oauth2`` → failed LiteLLM auth falls back to anonymous so - the bearer can be forwarded to upstream. + An ``auth_type=oauth2`` server that has NOT opted into + ``delegate_auth_to_upstream`` must not exchange a failed LiteLLM auth for + an anonymous session: forwarding an arbitrary bearer upstream is only + allowed once the operator explicitly delegates auth. A failed validation + here is a genuine 401 and propagates (which is also what keeps the + success-path trace free of a phantom 401, since no doomed validation runs + for a delegated server). """ from fastapi import HTTPException @@ -1389,8 +1381,9 @@ async def mock_user_api_key_auth_fails(api_key, request): mock_mgr.get_mcp_server_by_name.return_value = ( TestMCPOAuth2FallbackTargetGating._make_server(MCPAuth.oauth2) ) - auth_result, *_rest = await MCPRequestHandler.process_mcp_request(scope) - assert isinstance(auth_result, UserAPIKeyAuth) + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request(scope) + assert exc_info.value.status_code == 401 async def test_fallback_allowed_when_target_is_passthrough(self): """ @@ -1668,19 +1661,16 @@ async def test_delegate_skips_litellm_auth_with_no_authorization(self): assert isinstance(auth_result, UserAPIKeyAuth) mock_auth.assert_not_called() - async def test_delegate_with_upstream_token_in_authorization_falls_back_to_anonymous( + async def test_delegate_with_upstream_token_in_authorization_skips_litellm_auth( self, ): """ oauth2 + delegate_auth_to_upstream=True with an upstream OAuth token in - ``Authorization`` (not a LiteLLM key): LiteLLM auth is attempted first - (and fails), then the existing oauth2 fallback returns anonymous so the - bearer is forwarded upstream untouched. The delegate branch itself does - not fire when Authorization is present — that is what protects spend - tracking for callers using Authorization-style LiteLLM keys. + ``Authorization``: the delegate gate fires before any LiteLLM validation, + so ``user_api_key_auth`` is never called and the bearer is forwarded + upstream untouched. Skipping the doomed validation is what keeps a tool + call that actually succeeds from carrying a phantom 401 auth span. """ - from fastapi import HTTPException - from litellm.types.mcp import MCPAuth scope = { @@ -1690,14 +1680,11 @@ async def test_delegate_with_upstream_token_in_authorization_falls_back_to_anony "headers": [(b"authorization", b"Bearer upstream-pkce-token")], } - async def mock_user_api_key_auth_fails(api_key, request): - raise HTTPException(status_code=401, detail="Invalid API key") - with ( patch( "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", - side_effect=mock_user_api_key_auth_fails, - ), + new_callable=AsyncMock, + ) as mock_auth, patch( "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager" ) as mock_mgr, @@ -1718,6 +1705,7 @@ async def mock_user_api_key_auth_fails(api_key, request): ) = await MCPRequestHandler.process_mcp_request(scope) assert isinstance(auth_result, UserAPIKeyAuth) assert oauth2_headers.get("Authorization") == "Bearer upstream-pkce-token" + mock_auth.assert_not_called() async def test_delegate_off_still_requires_litellm_auth(self): """ @@ -1912,12 +1900,15 @@ async def test_explicit_litellm_key_takes_precedence_over_delegate(self): assert auth_result.user_id == "real-user" mock_auth.assert_called_once() - async def test_litellm_key_via_authorization_header_not_bypassed(self): + async def test_authorization_bearer_on_delegate_server_treated_as_upstream(self): """ - Regression: a LiteLLM key sent via the secondary ``Authorization`` header - (e.g. ``Authorization: Bearer sk-...``) must still trigger normal auth - and not be silently swallowed by the delegate bypass — otherwise spend - tracking and rate limiting are skipped for those callers. + On a delegate server the ``Authorization`` header is, by contract, an + upstream token rather than a LiteLLM key — even when it is sk-shaped. It + is forwarded upstream without LiteLLM validation, so ``user_api_key_auth`` + is not called and no LiteLLM identity is resolved. Callers who need + LiteLLM identity / spend tracking on a delegate server must supply + ``x-litellm-api-key`` (see + test_explicit_litellm_key_takes_precedence_over_delegate). """ from litellm.types.mcp import MCPAuth @@ -1944,10 +1935,18 @@ async def test_litellm_key_via_authorization_header_not_bypassed(self): delegate_auth_to_upstream=True, ) ) - auth_result, *_rest = await MCPRequestHandler.process_mcp_request(scope) + ( + auth_result, + _, + _, + _, + oauth2_headers, + _, + ) = await MCPRequestHandler.process_mcp_request(scope) assert isinstance(auth_result, UserAPIKeyAuth) - assert auth_result.user_id == "real-user" - mock_auth.assert_called_once() + assert auth_result.user_id is None + assert oauth2_headers.get("Authorization") == "Bearer sk-1234" + mock_auth.assert_not_called() async def test_delegate_ignored_for_client_credentials_server(self): """ diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index da66d60aed8a..6fd935e3364c 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -1,5 +1,6 @@ """Tests for MCP OAuth discoverable endpoints""" +import json from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -2661,3 +2662,74 @@ async def test_token_endpoint_sets_no_store_cache_control(): assert response.headers["cache-control"] == "no-store" assert response.headers["pragma"] == "no-cache" + + +async def _exchange_with_upstream_token_response(upstream_body): + from fastapi import Request + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + exchange_token_with_server, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp import MCPAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="t", + name="t", + server_name="t", + alias="t", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id="cid", + client_secret="cs", + authorization_url="https://provider.com/oauth/authorize", + token_url="https://provider.com/oauth/token", + ) + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://litellm.example.com/" + mock_request.headers = {} + + fake_http_response = MagicMock() + fake_http_response.json.return_value = upstream_body + fake_http_response.raise_for_status = MagicMock() + fake_http_client = MagicMock() + fake_http_client.post = AsyncMock(return_value=fake_http_response) + + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=fake_http_client, + ): + response = await exchange_token_with_server( + request=mock_request, + mcp_server=server, + grant_type="authorization_code", + code="c", + redirect_uri="http://127.0.0.1:3000/cb", + client_id="cid", + client_secret=None, + code_verifier=None, + ) + return json.loads(response.body) + + +@pytest.mark.asyncio +async def test_token_exchange_omits_expires_in_when_upstream_omits_it(): + """A provider that issues a non-expiring token (e.g. Slack without token + rotation) returns no ``expires_in``. The exchange must mirror that and omit + ``expires_in`` rather than fabricate a 1-hour TTL, so the stored credential + is treated as non-expiring instead of dying after an hour.""" + body = await _exchange_with_upstream_token_response( + {"access_token": "tok", "token_type": "Bearer"} + ) + assert "expires_in" not in body + + +@pytest.mark.asyncio +async def test_token_exchange_passes_through_upstream_expires_in(): + """When the provider does send ``expires_in`` (e.g. Slack with token + rotation), the exchange forwards the real value unchanged.""" + body = await _exchange_with_upstream_token_response( + {"access_token": "tok", "token_type": "Bearer", "expires_in": 43200} + ) + assert body["expires_in"] == 43200 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_env_vars.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_env_vars.py index 19065ff816ba..a846ca247392 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_env_vars.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_env_vars.py @@ -872,6 +872,7 @@ def _mock_env_vars_prisma(row=None): prisma.db.litellm_mcpuserenvvars.find_many = AsyncMock(return_value=[]) prisma.db.litellm_mcpuserenvvars.upsert = AsyncMock() prisma.db.litellm_mcpuserenvvars.delete_many = AsyncMock() + prisma.db.litellm_mcpusercredentials.delete_many = AsyncMock() return prisma @@ -1252,6 +1253,64 @@ async def test_delete_mcp_server_succeeds_when_orphan_cleanup_fails(): prisma.db.litellm_mcpuserenvvars.delete_many.assert_awaited_once() +@pytest.mark.asyncio +async def test_delete_mcp_server_removes_orphaned_user_credentials(): + """Deleting a server must also drop every user's stored BYOK/OAuth credential + rows for it; there is no FK cascade, so skipping this leaves encrypted secrets + pointing at a now-missing server.""" + from unittest.mock import AsyncMock + + from litellm.proxy._experimental.mcp_server.db import delete_mcp_server + + prisma = _mock_env_vars_prisma() + prisma.db.litellm_mcpservertable.delete = AsyncMock(return_value=object()) + + await delete_mcp_server(prisma, "srv-1") + + prisma.db.litellm_mcpusercredentials.delete_many.assert_awaited_once() + call = prisma.db.litellm_mcpusercredentials.delete_many.call_args + assert call.kwargs["where"] == {"server_id": "srv-1"} + + +@pytest.mark.asyncio +async def test_delete_mcp_server_skips_credential_cleanup_when_server_missing(): + """A no-op delete (server not found) must not touch the credential table.""" + from unittest.mock import AsyncMock + + from litellm.proxy._experimental.mcp_server.db import delete_mcp_server + + prisma = _mock_env_vars_prisma() + prisma.db.litellm_mcpservertable.delete = AsyncMock(return_value=None) + + result = await delete_mcp_server(prisma, "srv-1") + + assert result is None + prisma.db.litellm_mcpusercredentials.delete_many.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_delete_mcp_server_credential_cleanup_failure_still_cleans_env_vars(): + """Each per-user table is cleaned independently: a failure dropping credential + rows must not skip the env var cleanup (or vice versa), and the delete must + still succeed for the caller.""" + from unittest.mock import AsyncMock + + from litellm.proxy._experimental.mcp_server.db import delete_mcp_server + + deleted = object() + prisma = _mock_env_vars_prisma() + prisma.db.litellm_mcpservertable.delete = AsyncMock(return_value=deleted) + prisma.db.litellm_mcpusercredentials.delete_many = AsyncMock( + side_effect=Exception("connection pool exhausted") + ) + + result = await delete_mcp_server(prisma, "srv-1") + + assert result is deleted + prisma.db.litellm_mcpusercredentials.delete_many.assert_awaited_once() + prisma.db.litellm_mcpuserenvvars.delete_many.assert_awaited_once() + + # ── DB helpers: global env vars encrypted at rest ───────────────────────── 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 caff9ea2d289..47c9396f1210 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 @@ -501,6 +501,7 @@ async def fake_get_tools( raw_headers=None, user_api_key_auth=None, extra_headers=None, + apply_tool_filters=True, ): captured["called"] = True captured["server"] = server @@ -545,6 +546,78 @@ async def fake_get_tools( assert result["error"] is None assert result["message"] == "Successfully retrieved tools" + 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 + to enumerate deliberately-disabled tools.""" + from litellm.proxy._types import LitellmUserRoles + + 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"] + + class StubServer: + alias = "server-1" + server_name = "server-1" + name = "stub" + allowed_tools = ["tool1"] + mcp_info = {"server_name": "stub"} + available_on_public_internet = True + + stub_server = StubServer() + captured = {} + + async def fake_get_tools( + server, server_auth_header, *args, apply_tool_filters=True, **kwargs + ): + captured["apply_tool_filters"] = apply_tool_filters + return ["tool-1"] + + 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, + "get_mcp_server_by_id", + lambda server_id: stub_server if server_id == "server-1" else None, + raising=False, + ) + monkeypatch.setattr( + rest_endpoints, + "_get_tools_for_single_server", + fake_get_tools, + raising=False, + ) + + request = _build_request(path="/mcp-rest/tools/list", method="GET") + + await rest_endpoints.list_tool_rest_api( + request, + server_id="server-1", + include_disabled_tools=True, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + assert captured["apply_tool_filters"] is False + + await rest_endpoints.list_tool_rest_api( + request, + server_id="server-1", + include_disabled_tools=True, + user_api_key_dict=UserAPIKeyAuth(), + ) + assert captured["apply_tool_filters"] is True + @pytest.mark.parametrize("upstream_status", [401, 403]) async def test_upstream_auth_failure_surfaces_status_and_challenge( self, monkeypatch, upstream_status @@ -649,6 +722,7 @@ async def fake_get_tools( raw_headers=None, user_api_key_auth=None, extra_headers=None, + apply_tool_filters=True, ): captured["called"] = True captured["server_arg"] = server @@ -792,6 +866,7 @@ async def fake_get_tools( raw_headers=None, user_api_key_auth=None, extra_headers=None, + apply_tool_filters=True, ): captured["server"] = server captured["auth_header"] = server_auth_header @@ -1284,6 +1359,56 @@ async def fake_get_tools_from_server(**kwargs): assert "tool1" not in tool_names assert "tool4" not in tool_names + async def test_apply_tool_filters_false_returns_full_catalog(self, monkeypatch): + """apply_tool_filters=False returns the raw catalog without the server + allowed_tools gate, so the config UI can render disabled tools as off.""" + from litellm.proxy._experimental.mcp_server.server import MCPServer + from litellm.types.mcp import MCPTransport + + class MockTool: + def __init__(self, name): + self.name = name + self.description = name + self.inputSchema = {} + + mock_tools = [MockTool("tool1"), MockTool("tool2"), MockTool("tool3")] + + async def fake_get_tools_from_server(**kwargs): + return mock_tools + + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "_get_tools_from_server", + fake_get_tools_from_server, + raising=False, + ) + + # Server enforces an allowlist of just tool1. + server = MCPServer( + server_id="test-server-id", + name="test-server", + transport=MCPTransport.sse, + allowed_tools=["tool1"], + ) + user_api_key_dict = UserAPIKeyAuth(api_key="test-key", object_permission=None) + + # Runtime default: only the allowed tool comes back. + filtered = await rest_endpoints._get_tools_for_single_server( + server=server, + server_auth_header=None, + user_api_key_auth=user_api_key_dict, + ) + assert [t.name for t in filtered] == ["tool1"] + + # Config view: full catalog, including the disabled tools. + full = await rest_endpoints._get_tools_for_single_server( + server=server, + server_auth_header=None, + user_api_key_auth=user_api_key_dict, + apply_tool_filters=False, + ) + assert {t.name for t in full} == {"tool1", "tool2", "tool3"} + class TestStdioCommandAllowlist: """Tests for MCP stdio command allowlist validation.""" diff --git a/tests/test_litellm/proxy/auth/test_auth_exception_handler.py b/tests/test_litellm/proxy/auth/test_auth_exception_handler.py index 27f6015e6f49..11e6f483e358 100644 --- a/tests/test_litellm/proxy/auth/test_auth_exception_handler.py +++ b/tests/test_litellm/proxy/auth/test_auth_exception_handler.py @@ -112,6 +112,166 @@ async def test_handle_authentication_error_data_layer_errors_do_not_fall_back( ) +@pytest.mark.asyncio +@pytest.mark.parametrize( + "db_error", + [ + ConnectionError("connection refused"), + TimeoutError("timed out"), + asyncio.TimeoutError(), + OSError("network is unreachable"), + HTTPClientClosedError(), + PrismaError("can't reach database server"), + RawQueryError( + data={ + "user_facing_error": { + "message": "cached plan must not change result type", + "meta": {"table": "t"}, + } + } + ), + ], +) +async def test_handle_authentication_error_db_infra_error_returns_503(db_error): + """Regression for the outage where valid keys got 401 for 4 hours: an + infrastructure-level DB failure during auth must surface as 503 (the DB + could not confirm the key), never as 401 ("Invalid API key").""" + handler = UserAPIKeyAuthExceptionHandler() + + with ( + patch( + "litellm.proxy.proxy_server.proxy_logging_obj.post_call_failure_hook", + new_callable=AsyncMock, + return_value=None, + ), + patch( + "litellm.proxy.auth.auth_exception_handler.seed_request_identity", + ), + patch( + "litellm.proxy.proxy_server.general_settings", + {"allow_requests_on_db_unavailable": False}, + ), + ): + with pytest.raises(ProxyException) as exc_info: + await handler._handle_authentication_error( + db_error, + MagicMock(), + {}, + "/v1/chat/completions", + None, + "sk-valid-but-db-down", + ) + + assert int(exc_info.value.code) == status.HTTP_503_SERVICE_UNAVAILABLE + assert exc_info.value.type == ProxyErrorTypes.no_db_connection + assert "Invalid API key" not in str(exc_info.value.message) + + +@pytest.mark.asyncio +async def test_handle_authentication_error_prisma_engine_teardown_returns_503(): + """Regression for the first-request-of-an-outage edge case: at the instant + the DB socket drops, the prisma query engine returns a malformed error + payload and prisma-client-py crashes with a bare + ``AttributeError: 'NoneType' object has no attribute 'get'`` before it can + raise P1001. That AttributeError reached auth and fell through to 401. It + must surface as 503 like every other infra failure during the outage.""" + from prisma.engine import utils as prisma_engine_utils + + malformed_payload = [ + { + "error": "Can't reach database server", + "user_facing_error": { + "error_code": "P1001", + "message": "Can't reach database server at `localhost`:`5503`", + "meta": None, + }, + } + ] + try: + prisma_engine_utils.handle_response_errors(None, malformed_payload) + raise AssertionError("expected prisma to raise AttributeError") + except AttributeError as e: + teardown_error = e + + handler = UserAPIKeyAuthExceptionHandler() + + with ( + patch( + "litellm.proxy.proxy_server.proxy_logging_obj.post_call_failure_hook", + new_callable=AsyncMock, + return_value=None, + ), + patch( + "litellm.proxy.auth.auth_exception_handler.seed_request_identity", + ), + patch( + "litellm.proxy.proxy_server.general_settings", + {"allow_requests_on_db_unavailable": False}, + ), + ): + with pytest.raises(ProxyException) as exc_info: + await handler._handle_authentication_error( + teardown_error, + MagicMock(), + {}, + "/v1/chat/completions", + None, + "sk-valid-but-db-down", + ) + + assert int(exc_info.value.code) == status.HTTP_503_SERVICE_UNAVAILABLE + assert exc_info.value.type == ProxyErrorTypes.no_db_connection + assert "Invalid API key" not in str(exc_info.value.message) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "auth_error", + [ + # DB returned no row -> get_key_object raises this exact 401. + ProxyException( + message="Authentication Error, Invalid proxy server token passed.", + type=ProxyErrorTypes.token_not_found_in_db, + param="key", + code=status.HTTP_401_UNAUTHORIZED, + ), + # A bare auth failure raised as a plain Exception (e.g. master-key-only + # route) must keep returning 401, not get reclassified as 503. + Exception("Invalid proxy server token passed"), + ], +) +async def test_handle_authentication_error_genuine_auth_failure_stays_401(auth_error): + """Guard against the 503 conversion being too broad: a genuine auth + failure (missing key / wrong key) must still be 401.""" + handler = UserAPIKeyAuthExceptionHandler() + + with ( + patch( + "litellm.proxy.proxy_server.proxy_logging_obj.post_call_failure_hook", + new_callable=AsyncMock, + return_value=None, + ), + patch( + "litellm.proxy.auth.auth_exception_handler.seed_request_identity", + ), + patch( + "litellm.proxy.proxy_server.general_settings", + {"allow_requests_on_db_unavailable": False}, + ), + ): + with pytest.raises(ProxyException) as exc_info: + await handler._handle_authentication_error( + auth_error, + MagicMock(), + {}, + "/v1/chat/completions", + None, + "sk-bad-key", + ) + + assert int(exc_info.value.code) == status.HTTP_401_UNAUTHORIZED + + @pytest.mark.asyncio async def test_handle_authentication_error_budget_exceeded(): handler = UserAPIKeyAuthExceptionHandler() diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index fa6cc8bed1b7..80f12d4459f4 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -112,11 +112,71 @@ async def test_should_clear_stale_budget_reservation_when_budget_checks_skip(): user_api_key_cache=MagicMock(), proxy_logging_obj=MagicMock(), skip_budget_checks=True, + general_settings={}, ) assert user_api_key_auth_obj.budget_reservation is None +@pytest.mark.asyncio +async def test_disable_budget_reservation_skips_reservation(): + """#27639: general_settings.disable_budget_reservation turns off the optimistic Redis + reservation so operators hit by phantom BudgetExceededError can opt out of it.""" + user_api_key_auth_obj = UserAPIKeyAuth(token="test_token") + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.reserve_budget_for_request", + new=AsyncMock(return_value={"reserved_cost": 0.5, "entries": []}), + ) as mock_reserve: + await _reserve_budget_after_common_checks( + user_api_key_auth_obj=user_api_key_auth_obj, + request_data={"model": "gpt-4o"}, + route="/v1/chat/completions", + llm_router=None, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + skip_budget_checks=False, + general_settings={"disable_budget_reservation": True}, + ) + + mock_reserve.assert_not_called() + assert user_api_key_auth_obj.budget_reservation is None + + +@pytest.mark.asyncio +async def test_budget_reservation_runs_when_not_disabled(): + """Control for #27639: with the flag absent, the reservation still runs and is stored.""" + user_api_key_auth_obj = UserAPIKeyAuth(token="test_token") + reservation = { + "reserved_cost": 0.5, + "entries": [{"counter_key": "spend:key:test_token"}], + } + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.reserve_budget_for_request", + new=AsyncMock(return_value=reservation), + ) as mock_reserve: + await _reserve_budget_after_common_checks( + user_api_key_auth_obj=user_api_key_auth_obj, + request_data={"model": "gpt-4o"}, + route="/v1/chat/completions", + llm_router=None, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + skip_budget_checks=False, + general_settings={}, + ) + + mock_reserve.assert_awaited_once() + assert user_api_key_auth_obj.budget_reservation == reservation + + @pytest.mark.asyncio async def test_should_not_reuse_cached_key_object_for_request_state(): key_cache = DualCache() @@ -1636,7 +1696,9 @@ async def test_auto_register_passes_validated_org_context_to_generated_key(self) assert mock_auto_register.call_args.kwargs["team_id"] == "validated-team" assert mock_auto_register.call_args.kwargs["user_id"] == "validated-user" assert mock_auto_register.call_args.kwargs["org_id"] == "validated-org" - assert mock_auto_register.call_args.kwargs["end_user_id"] == "validated-end-user" + assert ( + mock_auto_register.call_args.kwargs["end_user_id"] == "validated-end-user" + ) assert result.org_id == "validated-org" @pytest.mark.asyncio @@ -3548,3 +3610,118 @@ async def test_user_api_key_auth_does_not_overwrite_end_user_id_set_by_builder() finally: for k, v in originals.items(): setattr(_proxy_server_mod, k, v) + + +def _proxy_attrs_for_db_lookup(): + """Minimal proxy_server attributes for driving the real + ``_user_api_key_auth_builder`` down to the DB key lookup.""" + proxy_logging_obj = MagicMock() + proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) + return { + "prisma_client": MagicMock(), + "user_api_key_cache": DualCache(), + "proxy_logging_obj": proxy_logging_obj, + "master_key": "sk-test-master", + "general_settings": {"allow_requests_on_db_unavailable": False}, + "llm_model_list": [], + "llm_router": None, + "open_telemetry_logger": None, + "model_max_budget_limiter": MagicMock(), + "user_custom_auth": None, + "jwt_handler": None, + "litellm_proxy_admin_name": "admin", + } + + +async def _run_builder_with_key_lookup(get_key_object_mock): + """Drive the real auth builder with ``get_key_object`` replaced by the + given mock. Returns the builder result. Patches ``seed_request_identity`` + so the failure path doesn't touch OTEL.""" + from fastapi import Request + from starlette.datastructures import URL + + import litellm.proxy.proxy_server as _proxy_server_mod + from litellm.proxy.auth.user_api_key_auth import _user_api_key_auth_builder + + attrs = _proxy_attrs_for_db_lookup() + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + with ( + patch( + "litellm.proxy.auth.user_api_key_auth.get_key_object", + get_key_object_mock, + ), + patch( + "litellm.proxy.auth.auth_exception_handler.seed_request_identity", + ), + ): + return await _user_api_key_auth_builder( + request=request, + api_key="Bearer sk-db-lookup-test", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={}, + ) + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + +@pytest.mark.asyncio +async def test_builder_returns_503_when_db_lookup_raises_infra_error(): + """End-to-end: a DB infrastructure failure during the key lookup must + propagate past the ``except ProxyException`` guard and surface as 503, + not the 401 that masked the 4-hour outage. Killing the new 503 branch + flips this to 401 and fails the test.""" + get_key_object = AsyncMock(side_effect=ConnectionError("connection refused")) + + with pytest.raises(ProxyException) as exc_info: + await _run_builder_with_key_lookup(get_key_object) + + assert int(exc_info.value.code) == status.HTTP_503_SERVICE_UNAVAILABLE + assert exc_info.value.type == ProxyErrorTypes.no_db_connection + assert "Invalid API key" not in str(exc_info.value.message) + + +@pytest.mark.asyncio +async def test_builder_returns_401_when_db_lookup_reports_missing_key(): + """Regression guard: a genuinely missing key (DB returned no row, which + ``get_key_object`` raises as a 401 ProxyException) must still be 401.""" + missing_key_error = ProxyException( + message="Authentication Error, Invalid proxy server token passed. key=..., not found in db.", + type=ProxyErrorTypes.token_not_found_in_db, + param="key", + code=status.HTTP_401_UNAUTHORIZED, + ) + get_key_object = AsyncMock(side_effect=missing_key_error) + + with pytest.raises(ProxyException) as exc_info: + await _run_builder_with_key_lookup(get_key_object) + + assert int(exc_info.value.code) == status.HTTP_401_UNAUTHORIZED + + +@pytest.mark.asyncio +async def test_builder_succeeds_when_db_lookup_returns_valid_token(): + """Regression guard: a valid key still authenticates. Proves the 503 + conversion only fires on the failure path and never intercepts success.""" + valid_token = UserAPIKeyAuth(api_key="sk-db-lookup-test", token="hashed-valid") + get_key_object = AsyncMock(return_value=valid_token) + + with patch( + "litellm.proxy.auth.user_api_key_auth._return_user_api_key_auth_obj", + new_callable=AsyncMock, + return_value=valid_token, + ) as mock_return: + result = await _run_builder_with_key_lookup(get_key_object) + + assert isinstance(result, UserAPIKeyAuth) + # Reaching the success-assembly return (never the exception handler) + # proves a valid key is unaffected by the 503 conversion. + mock_return.assert_awaited_once() diff --git a/tests/test_litellm/proxy/db/test_exception_handler.py b/tests/test_litellm/proxy/db/test_exception_handler.py index 9dcf5df4aebf..6021c2214267 100644 --- a/tests/test_litellm/proxy/db/test_exception_handler.py +++ b/tests/test_litellm/proxy/db/test_exception_handler.py @@ -107,6 +107,201 @@ def test_is_database_connection_generic_errors(): ) +@pytest.mark.parametrize( + "error", + [ + ConnectionError("connection refused"), + TimeoutError("timed out"), + OSError("network is unreachable"), + asyncio.TimeoutError(), + HTTPClientClosedError(), + ClientNotConnectedError(), + PrismaError("can't reach database server"), + PrismaError(), + ], +) +def test_is_database_service_unavailable_error_infra_failures(error): + """Infrastructure-level failures (socket/connection/timeout, prisma + transport, unknown PrismaError) mean the DB could not answer, so auth + must surface 503 instead of treating a valid key as invalid.""" + assert PrismaDBExceptionHandler.is_database_service_unavailable_error(error) is True + + +def test_is_database_service_unavailable_error_prisma_p1001_masquerades_as_dataerror(): + """Real-world regression: prisma-client-py raises the P1001 "can't reach + database server" connectivity failure as a DataError (a data-layer type). + A type-only check would miss it and return 401 during a genuine outage; + the message keyword must still classify it as service-unavailable -> 503.""" + p1001_as_dataerror = DataError( + data={ + "user_facing_error": { + "message": "Can't reach database server at `127.0.0.1`:`5499`", + "meta": {"table": "t"}, + } + } + ) + assert ( + PrismaDBExceptionHandler.is_database_service_unavailable_error( + p1001_as_dataerror + ) + is True + ) + + +def test_is_database_service_unavailable_error_cached_plan_escapes_as_503(): + """Composes with the cached-plan retry: when that recovery fails and the + Postgres "cached plan must not change result type" error escapes (raised by + prisma as a data-layer RawQueryError), it is a transient stale-DB-state + condition, not an invalid key, so it must classify as service-unavailable + -> 503 rather than fall through to 401.""" + cached_plan_error = RawQueryError( + data={ + "user_facing_error": { + "message": "cached plan must not change result type", + "meta": {"table": "t"}, + } + } + ) + assert ( + PrismaDBExceptionHandler.is_database_service_unavailable_error( + cached_plan_error + ) + is True + ) + + +def test_is_database_service_unavailable_error_prisma_engine_malformed_payload(): + """Real-world regression: at the instant the DB socket drops, the prisma + query engine returns a malformed error payload (``user_facing_error.meta`` + is ``null``). prisma-client-py's ``handle_response_errors`` then crashes + with ``AttributeError: 'NoneType' object has no attribute 'get'`` before it + can raise the proper P1001 error. That bare AttributeError has no + connection keyword, so without the prisma-engine-origin check it falls + through to 401 on the first request of an outage. Reproduce the exact + prisma crash and assert it classifies as service-unavailable -> 503.""" + from prisma.engine import utils as prisma_engine_utils + + malformed_payload = [ + { + "error": "Can't reach database server", + "user_facing_error": { + "error_code": "P1001", + "message": "Can't reach database server at `localhost`:`5503`", + "meta": None, + }, + } + ] + with pytest.raises(AttributeError) as exc_info: + prisma_engine_utils.handle_response_errors(None, malformed_payload) + + assert "no attribute 'get'" in str(exc_info.value) + assert ( + PrismaDBExceptionHandler.is_database_service_unavailable_error(exc_info.value) + is True + ) + + +def test_is_prisma_engine_internal_error_excludes_application_attributeerror(): + """The prisma-engine-origin check must stay narrow: a genuine AttributeError + raised by application code (a real bug) must NOT be classified as + service-unavailable, otherwise real bugs would silently become 503s.""" + + def application_bug(): + none_value = None + return none_value.get("oops") + + with pytest.raises(AttributeError) as exc_info: + application_bug() + + assert ( + PrismaDBExceptionHandler.is_prisma_engine_internal_error(exc_info.value) + is False + ) + assert ( + PrismaDBExceptionHandler.is_database_service_unavailable_error(exc_info.value) + is False + ) + + +def test_is_prisma_engine_internal_error_excludes_data_layer_prisma_error(): + """A data-layer ``PrismaError`` (the DB IS reachable and rejected the data) + must stay 401. These are always raised from prisma internals, so the check + excludes any ``PrismaError`` by type before inspecting the traceback.""" + data_layer_error = UniqueViolationError( + data={"user_facing_error": {"meta": {"table": "t"}}} + ) + try: + raise data_layer_error + except UniqueViolationError as e: + assert PrismaDBExceptionHandler.is_prisma_engine_internal_error(e) is False + + +@pytest.mark.parametrize( + "error", + [ + DataError(data={"user_facing_error": {"meta": {"table": "t"}}}), + UniqueViolationError(data={"user_facing_error": {"meta": {"table": "t"}}}), + RecordNotFoundError(data={"user_facing_error": {"meta": {"table": "t"}}}), + Exception("some unrelated error"), + ValueError("bad value"), + ], +) +def test_is_database_service_unavailable_error_excludes_non_infra(error): + """Data-layer errors (the DB IS reachable and answered) and generic + non-DB errors must NOT be classified as service-unavailable, otherwise a + genuine 401 would be masked as a transient 503.""" + assert ( + PrismaDBExceptionHandler.is_database_service_unavailable_error(error) is False + ) + + +def test_is_database_service_unavailable_error_asyncpg(monkeypatch): + """asyncpg connection/interface errors map to service-unavailable. asyncpg + is not a hard dependency, so inject a stand-in module to exercise the + branch deterministically regardless of the install environment.""" + import sys + import types + + fake_asyncpg = types.ModuleType("asyncpg") + fake_exceptions = types.ModuleType("asyncpg.exceptions") + + class PostgresConnectionError(Exception): + pass + + class InterfaceError(Exception): + pass + + class UniqueViolationError(Exception): # data-layer, must stay False + pass + + fake_exceptions.PostgresConnectionError = PostgresConnectionError + fake_exceptions.InterfaceError = InterfaceError + fake_exceptions.UniqueViolationError = UniqueViolationError + fake_asyncpg.exceptions = fake_exceptions + + monkeypatch.setitem(sys.modules, "asyncpg", fake_asyncpg) + monkeypatch.setitem(sys.modules, "asyncpg.exceptions", fake_exceptions) + + assert ( + PrismaDBExceptionHandler.is_database_service_unavailable_error( + PostgresConnectionError("connection reset") + ) + is True + ) + assert ( + PrismaDBExceptionHandler.is_database_service_unavailable_error( + InterfaceError("connection was closed") + ) + is True + ) + assert ( + PrismaDBExceptionHandler.is_database_service_unavailable_error( + UniqueViolationError("duplicate key") + ) + is False + ) + + # Test should_allow_request_on_db_unavailable method @patch( "litellm.proxy.proxy_server.general_settings", diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py index c58c94cbbc7d..014362554e4c 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py @@ -412,6 +412,171 @@ async def test_apply_guardrail_response_ok( assert result["texts"] == inputs["texts"] +@pytest.mark.asyncio +async def test_apply_guardrail_sends_user_id_model_and_extra_info( + crowdstrike_aidr_guardrail: CrowdStrikeAIDRHandler, +) -> None: + inputs: GenericGuardrailAPIInputs = { + "texts": ["Hello"], + "structured_messages": [{"role": "user", "content": "Hello"}], + "model": "gpt-4o", + } + request_data = { + "messages": inputs["structured_messages"], + "model": "gpt-4o", + "litellm_metadata": { + "user_api_key_user_id": "uid-abc", + "user_api_key_user_email": "alice@example.com", + }, + } + guardrail_endpoint = ( + f"{crowdstrike_aidr_guardrail.api_base}/v1/guard_chat_completions" + ) + + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + return_value=httpx.Response( + status_code=200, + json={"result": {"blocked": False, "transformed": False}}, + request=httpx.Request(method="POST", url=guardrail_endpoint), + ), + ) as mock_method: + await crowdstrike_aidr_guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + ) + + payload = mock_method.call_args.kwargs["json"] + assert payload["user_id"] == "uid-abc" + assert payload["model"] == "gpt-4o" + assert payload["extra_info"] == {"user_name": "alice@example.com"} + + +@pytest.mark.asyncio +async def test_apply_guardrail_empty_extra_info_when_no_email( + crowdstrike_aidr_guardrail: CrowdStrikeAIDRHandler, +) -> None: + inputs: GenericGuardrailAPIInputs = { + "texts": ["Hello"], + "structured_messages": [{"role": "user", "content": "Hello"}], + "model": "gemini-flash", + } + request_data = { + "messages": inputs["structured_messages"], + "model": "gemini-flash", + "litellm_metadata": { + "user_api_key_user_id": "uid-no-email", + "user_api_key_user_email": None, + }, + } + guardrail_endpoint = ( + f"{crowdstrike_aidr_guardrail.api_base}/v1/guard_chat_completions" + ) + + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + return_value=httpx.Response( + status_code=200, + json={"result": {"blocked": False, "transformed": False}}, + request=httpx.Request(method="POST", url=guardrail_endpoint), + ), + ) as mock_method: + await crowdstrike_aidr_guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + ) + + payload = mock_method.call_args.kwargs["json"] + assert payload["user_id"] == "uid-no-email" + assert payload["model"] == "gemini-flash" + assert payload["extra_info"] == {} + + +@pytest.mark.asyncio +async def test_apply_guardrail_no_metadata_skips_user_fields( + crowdstrike_aidr_guardrail: CrowdStrikeAIDRHandler, +) -> None: + inputs: GenericGuardrailAPIInputs = { + "texts": ["Hello"], + "structured_messages": [{"role": "user", "content": "Hello"}], + } + request_data = {"messages": inputs["structured_messages"]} + guardrail_endpoint = ( + f"{crowdstrike_aidr_guardrail.api_base}/v1/guard_chat_completions" + ) + + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + return_value=httpx.Response( + status_code=200, + json={"result": {"blocked": False, "transformed": False}}, + request=httpx.Request(method="POST", url=guardrail_endpoint), + ), + ) as mock_method: + await crowdstrike_aidr_guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + ) + + payload = mock_method.call_args.kwargs["json"] + assert "user_id" not in payload + assert "model" not in payload + assert "extra_info" not in payload + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "litellm_metadata, metadata", + [ + (None, {"user_api_key_user_id": "uid-abc", "user_api_key_user_email": "alice@example.com"}), + ({"trace_id": "t1"}, {"user_api_key_user_id": "uid-abc", "user_api_key_user_email": "alice@example.com"}), + (["unexpected"], {"user_api_key_user_id": "uid-abc", "user_api_key_user_email": "alice@example.com"}), + ({"user_api_key_user_id": "uid-abc", "user_api_key_user_email": "alice@example.com"}, {"trace_id": "t1"}), + ], + ids=["identity_in_metadata_llm_none", "identity_in_metadata_llm_user_dict", "identity_in_metadata_llm_non_mapping", "identity_in_litellm_metadata"], +) +async def test_apply_guardrail_reads_identity_from_either_metadata_bag( + crowdstrike_aidr_guardrail: CrowdStrikeAIDRHandler, + litellm_metadata, + metadata, +) -> None: + inputs: GenericGuardrailAPIInputs = { + "texts": ["Hello"], + "structured_messages": [{"role": "user", "content": "Hello"}], + "model": "gpt-4o", + } + request_data = { + "messages": inputs["structured_messages"], + "model": "gpt-4o", + "litellm_metadata": litellm_metadata, + "metadata": metadata, + } + guardrail_endpoint = ( + f"{crowdstrike_aidr_guardrail.api_base}/v1/guard_chat_completions" + ) + + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + return_value=httpx.Response( + status_code=200, + json={"result": {"blocked": False, "transformed": False}}, + request=httpx.Request(method="POST", url=guardrail_endpoint), + ), + ) as mock_method: + await crowdstrike_aidr_guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + ) + + payload = mock_method.call_args.kwargs["json"] + assert payload["user_id"] == "uid-abc" + assert payload["extra_info"] == {"user_name": "alice@example.com"} + + @pytest.mark.asyncio async def test_apply_guardrail_request_skipped_messages_stay_aligned( crowdstrike_aidr_guardrail: CrowdStrikeAIDRHandler, diff --git a/tests/test_litellm/proxy/hooks/test_batch_file_validation.py b/tests/test_litellm/proxy/hooks/test_batch_file_validation.py index ae71d10b3788..a6f6e651487a 100644 --- a/tests/test_litellm/proxy/hooks/test_batch_file_validation.py +++ b/tests/test_litellm/proxy/hooks/test_batch_file_validation.py @@ -713,7 +713,8 @@ async def test_count_input_file_usage_decodes_model_embedded_file_id(): @pytest.mark.asyncio async def test_pre_call_allows_stripped_provider_model_when_key_has_proxy_alias(): """After replace_model_in_jsonl, body.model is the provider id (e.g. gpt-5.5). - Auth must check the proxy model_name the key was granted, not the stripped id.""" + Auth must check target_model_names from the unified file id, not reverse-map + the stripped id.""" from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter rate_limiter = _PROXY_BatchRateLimiter( @@ -732,7 +733,6 @@ async def test_pre_call_allows_stripped_provider_model_when_key_has_proxy_alias( ) mock_router = MagicMock() mock_router.model_list = [] - mock_router.resolve_model_name_from_model_id.return_value = proxy_alias can_key_call_model = AsyncMock(return_value=True) with ( @@ -745,10 +745,105 @@ async def test_pre_call_allows_stripped_provider_model_when_key_has_proxy_alias( await rate_limiter._enforce_batch_file_model_access( user_api_key_dict=user, file_content_as_dict=file_dict, + target_model_names=[proxy_alias], ) can_key_call_model.assert_awaited_once() assert can_key_call_model.await_args.kwargs["model"] == proxy_alias + mock_router.resolve_model_name_from_model_id.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model_list_order", + [ + [ + "openai/openai/gpt-5.5", + "openai/openai/gpt-5.5-batch", + "us/azure/openai/gpt-5.5", + ], + [ + "us/azure/openai/gpt-5.5", + "openai/openai/gpt-5.5", + "openai/openai/gpt-5.5-batch", + ], + [ + "openai/openai/gpt-5.5-batch", + "us/azure/openai/gpt-5.5", + "openai/openai/gpt-5.5", + ], + ], +) +async def test_pre_call_uses_target_model_names_not_stripped_reverse_lookup( + model_list_order, +): + """LIT-3593: three deployments strip to gpt-5.5; auth must use the upload + target alias from target_model_names, not first-match reverse lookup.""" + from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter + + rate_limiter = _PROXY_BatchRateLimiter( + internal_usage_cache=MagicMock(), + parallel_request_limiter=MagicMock(), + ) + batch_alias = "openai/openai/gpt-5.5-batch" + deployment_templates = { + "openai/openai/gpt-5.5": { + "model_name": "openai/openai/gpt-5.5", + "litellm_params": {"model": "openai/gpt-5.5"}, + "model_info": {"id": "openai/openai/gpt-5.5", "mode": "chat"}, + }, + "openai/openai/gpt-5.5-batch": { + "model_name": "openai/openai/gpt-5.5-batch", + "litellm_params": {"model": "openai/gpt-5.5"}, + "model_info": {"id": "openai/openai/gpt-5.5-batch", "mode": "batch"}, + }, + "us/azure/openai/gpt-5.5": { + "model_name": "us/azure/openai/gpt-5.5", + "litellm_params": {"model": "azure/gpt-5.5"}, + "model_info": {"id": "openai/openai/gpt-5.5", "mode": "chat"}, + }, + } + mock_router = MagicMock() + mock_router.model_list = [deployment_templates[name] for name in model_list_order] + + def _resolve(model_id): + for deployment in mock_router.model_list: + actual_model = deployment.get("litellm_params", {}).get("model") + if actual_model == model_id or ( + actual_model and actual_model.endswith(f"/{model_id}") + ): + return deployment.get("model_name") + return None + + mock_router.resolve_model_name_from_model_id.side_effect = _resolve + + file_dict = [ + {"body": {"model": "gpt-5.5", "messages": [{"role": "user", "content": "x"}]}} + ] + user = UserAPIKeyAuth( + api_key="sk-ok", + user_id="alice", + models=[batch_alias], + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + can_key_call_model = AsyncMock(return_value=True) + + with ( + patch( + "litellm.proxy.auth.auth_checks.can_key_call_model", + new=can_key_call_model, + ), + patch("litellm.proxy.proxy_server.llm_router", mock_router), + ): + await rate_limiter._enforce_batch_file_model_access( + user_api_key_dict=user, + file_content_as_dict=file_dict, + target_model_names=[batch_alias], + ) + + can_key_call_model.assert_awaited_once() + assert can_key_call_model.await_args.kwargs["model"] == batch_alias + mock_router.resolve_model_name_from_model_id.assert_not_called() @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index b5eb091bb81f..06986472d565 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -1482,6 +1482,10 @@ async def test_get_cached_temporary_mcp_server_non_admin_denied(self): "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", mock_manager, ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts", + AsyncMock(return_value=[non_admin]), + ), ): with pytest.raises(HTTPException) as exc_info: await _get_cached_temporary_mcp_server_or_404("server-x", non_admin) @@ -1514,6 +1518,10 @@ async def test_get_cached_temporary_mcp_server_non_admin_allowed(self): "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", mock_manager, ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts", + AsyncMock(return_value=[non_admin]), + ), ): result = await _get_cached_temporary_mcp_server_or_404( "server-x", non_admin @@ -1521,6 +1529,58 @@ async def test_get_cached_temporary_mcp_server_non_admin_allowed(self): assert result is registry_server + @pytest.mark.asyncio + async def test_get_cached_temporary_mcp_server_non_admin_allowed_via_team_access_group( + self, + ): + """Internal user whose only grant to the server flows through a team + access-group must pass the authorize/token access check. The check has to + expand the UI session into per-team contexts (build_effective_auth_contexts), + the same way the server-list grid does; checking only the bare session + context leaves the team grant invisible and 403s the user.""" + from litellm.constants import UI_SESSION_TOKEN_TEAM_ID + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + _get_cached_temporary_mcp_server_or_404, + ) + + registry_server = generate_mock_mcp_server_config_record(server_id="server-x") + ui_session_auth = generate_mock_user_api_key_auth( + user_role=LitellmUserRoles.INTERNAL_USER, + team_id=UI_SESSION_TOKEN_TEAM_ID, + ) + team_context = ui_session_auth.model_copy() + team_context.team_id = "team-with-mcp-grant" + + mock_manager = MagicMock() + mock_manager.get_mcp_server_by_id.return_value = registry_server + mock_manager.get_mcp_server_by_name.return_value = None + + def allowed_for(auth): + return ["server-x"] if auth.team_id == "team-with-mcp-grant" else [] + + mock_manager.get_allowed_mcp_servers = AsyncMock(side_effect=allowed_for) + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_cached_temporary_mcp_server", + return_value=None, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts", + AsyncMock(return_value=[ui_session_auth, team_context]), + ), + ): + result = await _get_cached_temporary_mcp_server_or_404( + "server-x", ui_session_auth + ) + + assert result is registry_server + assert mock_manager.get_allowed_mcp_servers.await_count == 2 + @pytest.mark.asyncio async def test_get_cached_temporary_mcp_server_temp_cache_non_admin_denied(self): """Servers resolved from the admin-only temp cache reject non-admins.""" diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index d580f1f7703b..90940c42dfd2 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -1615,6 +1615,7 @@ async def test_team_model_add_delete_refresh_team_cache(endpoint_name): mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( return_value=existing_team ) + mock_prisma_client.db.execute_raw = AsyncMock(return_value=None) mock_prisma_client.db.litellm_teamtable.update = AsyncMock( return_value=updated_team ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_model_alias_merge.py b/tests/test_litellm/proxy/management_endpoints/test_team_model_alias_merge.py new file mode 100644 index 000000000000..73b6e98d94f6 --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/test_team_model_alias_merge.py @@ -0,0 +1,77 @@ +""" +Tests for atomic team model operations during BYOK model creation. + +Regression tests for https://github.com/BerriAI/litellm/issues/22594 +Concurrent BYOK model creates must not overwrite each other's entries +in team.models. +""" + +import os +import sys +from unittest.mock import AsyncMock, MagicMock + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.proxy._types import ( + LitellmUserRoles, + TeamModelAddRequest, + UserAPIKeyAuth, +) + + +class TestTeamModelAddAtomicAppend: + """Verify team_model_add uses atomic SQL for the models array append.""" + + @pytest.mark.asyncio + async def test_uses_atomic_array_append_with_dedup(self): + """team_model_add must call execute_raw with DISTINCT unnest SQL.""" + from unittest.mock import patch + + from litellm.proxy.management_endpoints.team_endpoints import team_model_add + + mock_request = MagicMock() + mock_user = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test_user" + ) + + existing_team = MagicMock() + existing_team.model_dump.return_value = { + "team_id": "team-1", + "models": ["existing-model"], + } + + updated_team = MagicMock() + updated_team.team_id = "team-1" + + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch( + "litellm.proxy.management_endpoints.team_endpoints._refresh_cached_team", + new=AsyncMock(return_value=None), + ), + ): + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock( + return_value=existing_team + ) + mock_prisma.db.execute_raw = AsyncMock(return_value=None) + mock_prisma.db.litellm_teamtable.update = AsyncMock( + return_value=updated_team + ) + + await team_model_add( + data=TeamModelAddRequest(team_id="team-1", models=["new-model"]), + http_request=mock_request, + user_api_key_dict=mock_user, + ) + + mock_prisma.db.execute_raw.assert_called_once() + sql = mock_prisma.db.execute_raw.call_args[0][0] + assert "DISTINCT unnest" in sql + assert "all-proxy-models" in sql + assert mock_prisma.db.execute_raw.call_args[0][1] == ["new-model"] + assert mock_prisma.db.execute_raw.call_args[0][2] == "team-1" + + # Should use write-routed update to re-fetch, not find_unique + mock_prisma.db.litellm_teamtable.update.assert_called_once() diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py index f8b6fbde3dc4..1eab8184fa86 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py @@ -321,6 +321,307 @@ def test_cost_calculation_does_not_duplicate_provider_prefix( assert call_kwargs["model"] == "azure_ai/claude-sonnet-4-5_gb_20250929" assert call_kwargs["custom_llm_provider"] == "azure_ai" + @patch("litellm.completion_cost") + def test_cost_calculation_resolves_unknown_model_from_litellm_params( + self, mock_completion_cost + ): + """When the body model is the "unknown" sentinel, the deployment model + from litellm_params must be used for costing, not "unknown" (which makes + completion_cost raise and the cost silently fall back to $0).""" + from datetime import datetime + + from litellm.types.utils import ModelResponse + + mock_completion_cost.return_value = 0.001 + + logging_obj = self._create_mock_logging_obj(model="unknown") + logging_obj.model_call_details["litellm_params"] = { + "model": "anthropic/claude-3-5-haiku-20241022", + "metadata": { + "model_group": "passthrough/anthropic/claude-3-5-haiku-20241022" + }, + } + logging_obj.litellm_params = logging_obj.model_call_details["litellm_params"] + + mock_response = MagicMock(spec=ModelResponse) + mock_response.id = "test-id" + mock_response.model = "unknown" + + kwargs = AnthropicPassthroughLoggingHandler._create_anthropic_response_logging_payload( + litellm_model_response=mock_response, + model="unknown", + kwargs={}, + start_time=datetime.now(), + end_time=datetime.now(), + logging_obj=logging_obj, + ) + + mock_completion_cost.assert_called_once() + assert ( + mock_completion_cost.call_args[1]["model"] + == "anthropic/claude-3-5-haiku-20241022" + ) + assert kwargs["response_cost"] == 0.001 + assert kwargs["model"] == "anthropic/claude-3-5-haiku-20241022" + + @patch("litellm.completion_cost") + def test_cost_calculation_resolves_unknown_model_from_model_group( + self, mock_completion_cost + ): + """With only model_group available (no deployment litellm_params.model), + the leading passthrough/ prefix must be stripped so the cost map can + resolve the model.""" + from datetime import datetime + + from litellm.types.utils import ModelResponse + + mock_completion_cost.return_value = 0.002 + + logging_obj = self._create_mock_logging_obj(model="unknown") + logging_obj.model_call_details["litellm_params"] = { + "metadata": { + "model_group": "passthrough/anthropic/claude-3-5-haiku-20241022" + } + } + logging_obj.litellm_params = logging_obj.model_call_details["litellm_params"] + + mock_response = MagicMock(spec=ModelResponse) + mock_response.id = "test-id" + mock_response.model = "unknown" + + kwargs = AnthropicPassthroughLoggingHandler._create_anthropic_response_logging_payload( + litellm_model_response=mock_response, + model="unknown", + kwargs={}, + start_time=datetime.now(), + end_time=datetime.now(), + logging_obj=logging_obj, + ) + + mock_completion_cost.assert_called_once() + assert ( + mock_completion_cost.call_args[1]["model"] + == "anthropic/claude-3-5-haiku-20241022" + ) + assert kwargs["response_cost"] == 0.002 + + @patch("litellm.completion_cost") + def test_cost_calculation_skips_unknown_litellm_params_model_for_model_group( + self, mock_completion_cost + ): + """When litellm_params.model is itself the "unknown" sentinel, the + deployment-model branch must not short-circuit; resolution falls through + to model_group so costing still prices the real model instead of "unknown".""" + from datetime import datetime + + from litellm.types.utils import ModelResponse + + mock_completion_cost.return_value = 0.003 + + logging_obj = self._create_mock_logging_obj(model="unknown") + logging_obj.model_call_details["litellm_params"] = { + "model": "unknown", + "metadata": { + "model_group": "passthrough/anthropic/claude-3-5-haiku-20241022" + }, + } + logging_obj.litellm_params = logging_obj.model_call_details["litellm_params"] + + mock_response = MagicMock(spec=ModelResponse) + mock_response.id = "test-id" + mock_response.model = "unknown" + + kwargs = AnthropicPassthroughLoggingHandler._create_anthropic_response_logging_payload( + litellm_model_response=mock_response, + model="unknown", + kwargs={}, + start_time=datetime.now(), + end_time=datetime.now(), + logging_obj=logging_obj, + ) + + mock_completion_cost.assert_called_once() + assert ( + mock_completion_cost.call_args[1]["model"] + == "anthropic/claude-3-5-haiku-20241022" + ) + assert kwargs["response_cost"] == 0.003 + assert kwargs["model"] == "anthropic/claude-3-5-haiku-20241022" + + @patch("litellm.completion_cost") + def test_streaming_cost_calculation_resolves_model_from_message_start_chunk( + self, mock_completion_cost + ): + """On the bare /anthropic passthrough path litellm_params carries no model + or model_group and the body model is the "unknown" sentinel; the model + must be recovered from the message_start SSE event so completion_cost + prices the real model instead of failing on "unknown" and logging $0.""" + from datetime import datetime + + from litellm.litellm_core_utils.litellm_logging import ( + Logging as RealLoggingObj, + ) + from litellm.proxy.pass_through_endpoints.streaming_handler import ( + PassThroughStreamingHandler, + ) + + mock_completion_cost.return_value = 0.001 + + def _sse(event, data): + return f"event: {event}\ndata: {json.dumps(data)}\n\n".encode() + + frames = [ + _sse( + "message_start", + { + "type": "message_start", + "message": { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-3-5-haiku-20241022", + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 10, "output_tokens": 0}, + }, + }, + ), + _sse( + "content_block_start", + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "text", "text": ""}, + }, + ), + _sse( + "content_block_delta", + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "hi"}, + }, + ), + _sse("content_block_stop", {"type": "content_block_stop", "index": 0}), + _sse( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 1}, + }, + ), + _sse("message_stop", {"type": "message_stop"}), + ] + all_chunks = list( + PassThroughStreamingHandler._convert_raw_bytes_to_str_lines(frames) + ) + + logging_obj = RealLoggingObj( + model="unknown", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="pass_through_endpoint", + start_time=datetime.now(), + litellm_call_id="test-call-id", + function_id="1", + ) + logging_obj.model_call_details["model"] = "unknown" + logging_obj.model_call_details["stream"] = True + logging_obj.model_call_details["litellm_params"] = {} + logging_obj.litellm_params = {} + + result = AnthropicPassthroughLoggingHandler._handle_logging_anthropic_collected_chunks( + litellm_logging_obj=logging_obj, + passthrough_success_handler_obj=MagicMock(), + url_route="/anthropic/v1/messages", + request_body={"stream": True}, + endpoint_type="messages", + start_time=datetime.now(), + all_chunks=all_chunks, + end_time=datetime.now(), + ) + + assert result["result"] is not None + mock_completion_cost.assert_called_once() + assert mock_completion_cost.call_args[1]["model"] == "claude-3-5-haiku-20241022" + assert result["kwargs"]["response_cost"] == 0.001 + assert result["kwargs"]["model"] == "claude-3-5-haiku-20241022" + + def test_extract_model_skips_non_dict_data_payload(self): + """A scalar data: payload (e.g. `data: null`) must be skipped, not crash + the streaming log handler with AttributeError, which would propagate out + and break spend logging for the whole request.""" + chunks = [ + "event: ping\ndata: null\n\n", + 'event: message_start\ndata: {"type": "message_start", "message": ' + '{"model": "claude-3-5-haiku-20241022"}}\n\n', + ] + + assert ( + AnthropicPassthroughLoggingHandler._extract_model_from_anthropic_chunks( + chunks + ) + == "claude-3-5-haiku-20241022" + ) + + def test_extract_model_parses_per_line_not_first_data_substring(self): + """A raw multi-line SSE event whose non-data line contains the substring + "data:" must not derail parsing: matching only lines that start with + "data:" recovers the message_start model, whereas a first-substring slice + would consume the wrong offset, fail to parse JSON, and return None.""" + raw_event = ( + "event: ping data: not-json\n" + 'data: {"type": "message_start", "message": ' + '{"model": "claude-3-5-haiku-20241022"}}\n\n' + ) + + assert ( + AnthropicPassthroughLoggingHandler._extract_model_from_anthropic_chunks( + [raw_event] + ) + == "claude-3-5-haiku-20241022" + ) + + def test_passthrough_logging_sets_response_cost_with_server_tool_use_dict(self): + from litellm.types.utils import Choices, Message, ModelResponse + + logging_obj = self._create_mock_logging_obj(model="claude-3-7-sonnet-20250219") + logging_obj.get_router_model_id.return_value = None + logging_obj.litellm_params = {} + + response = ModelResponse( + id="test-id", + choices=[ + Choices( + finish_reason="stop", + index=0, + message=Message(content="test", role="assistant"), + ) + ], + created=1234567890, + model="claude-3-7-sonnet-20250219", + usage={ + "prompt_tokens": 10, + "completion_tokens": 5, + "total_tokens": 15, + "server_tool_use": {"web_search_requests": 1}, + }, + ) + + kwargs = AnthropicPassthroughLoggingHandler._create_anthropic_response_logging_payload( + litellm_model_response=response, + model="claude-3-7-sonnet-20250219", + kwargs={}, + start_time=datetime.now(), + end_time=datetime.now(), + logging_obj=logging_obj, + ) + + assert "response_cost" in kwargs + assert kwargs["response_cost"] > 0 + class TestAnthropicBatchPassthroughCostTracking: """Test cases for Anthropic batch passthrough cost tracking functionality""" @@ -1045,6 +1346,74 @@ def test_collapse_returns_none_for_interleaved_block_indexes(self): ) +class TestBuildCompleteStreamingResponseRobustness: + """_build_complete_streaming_response must tolerate non-standard SSE frames.""" + + def _build(self, chunks: List[str]): + return AnthropicPassthroughLoggingHandler._build_complete_streaming_response( + all_chunks=chunks, + litellm_logging_obj=MagicMock(), + model="claude-3-sonnet-20240229", + ) + + def test_done_frame_is_skipped(self): + """A bare 'data: [DONE]' control frame must not break reconstruction.""" + chunks = [ + 'event: message_start\ndata: {"type":"message_start","message":{"id":"msg_1","type":"message","role":"assistant","content":[],"model":"claude-3-sonnet-20240229","stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":10,"output_tokens":1}}}', + 'event: content_block_start\ndata: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}', + 'event: content_block_delta\ndata: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hi"}}', + 'event: content_block_stop\ndata: {"type":"content_block_stop","index":0}', + 'event: message_delta\ndata: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":2}}', + 'event: message_stop\ndata: {"type":"message_stop"}', + "data: [DONE]", + ] + result = self._build(chunks) + assert result is not None + assert result.choices[0].message.content == "Hi" + + def test_non_json_sse_line_is_skipped(self): + """Non-JSON SSE lines (comments, keep-alive pings) must be skipped.""" + chunks = [ + ": ping", + 'event: message_start\ndata: {"type":"message_start","message":{"id":"msg_1","type":"message","role":"assistant","content":[],"model":"claude-3-sonnet-20240229","stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":10,"output_tokens":1}}}', + "this is not json at all", + ] + # Must not raise; a malformed stream simply yields no usable response. + result = self._build(chunks) + assert result is None or hasattr(result, "choices") + + def test_mixed_valid_and_invalid_frames(self): + """Valid events are still collected when interleaved with invalid ones.""" + chunks = [ + 'event: message_start\ndata: {"type":"message_start","message":{"id":"msg_1","type":"message","role":"assistant","content":[],"model":"claude-3-sonnet-20240229","stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":10,"output_tokens":1}}}', + "data: [DONE]", + ": keep-alive", + "not-json", + 'event: content_block_start\ndata: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}', + 'event: content_block_delta\ndata: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello"}}', + 'event: content_block_stop\ndata: {"type":"content_block_stop","index":0}', + 'event: message_delta\ndata: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":2}}', + 'event: message_stop\ndata: {"type":"message_stop"}', + ] + result = self._build(chunks) + assert result is not None + assert result.choices[0].message.content == "Hello" + + def test_done_in_text_payload_is_not_dropped(self): + """A valid event whose text content contains '[DONE]' must NOT be skipped.""" + chunks = [ + 'event: message_start\ndata: {"type":"message_start","message":{"id":"msg_1","type":"message","role":"assistant","content":[],"model":"claude-3-sonnet-20240229","stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":10,"output_tokens":1}}}', + 'event: content_block_start\ndata: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}', + 'event: content_block_delta\ndata: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"The stream ends with [DONE]"}}', + 'event: content_block_stop\ndata: {"type":"content_block_stop","index":0}', + 'event: message_delta\ndata: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":8}}', + 'event: message_stop\ndata: {"type":"message_stop"}', + ] + result = self._build(chunks) + assert result is not None + assert result.choices[0].message.content == "The stream ends with [DONE]" + + class TestStreamFalseDeduplication: """ Regression tests for the duplicate-callback bug where a streaming pass-through diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py b/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py index 982598243785..017f4bd4368e 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py @@ -60,6 +60,15 @@ def test_v2_model_info_invalid_page_returns_422(client, auth_as, empty_router): assert "detail" in response.json() +def test_v2_model_info_in_openapi_schema(): + """``GET /v2/model/info`` is published in the proxy OpenAPI/Swagger spec.""" + from litellm.proxy.proxy_server import get_openapi_schema + + schema = get_openapi_schema() + assert "/v2/model/info" in schema["paths"] + assert "get" in schema["paths"]["/v2/model/info"] + + # --------------------------------------------------------------------------- # GET /v1/model/info, GET /model/info # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py b/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py index 97e5c494916d..6a8e0d15d8bd 100644 --- a/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py +++ b/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py @@ -151,21 +151,24 @@ async def test_model_info_v2_translates_team_model_name(monkeypatch): @pytest.mark.asyncio async def test_model_info_v1_list_path_translates_team_model_name(monkeypatch): - """/v1/model/info list path (no litellm_model_id) must surface the public - name. Covers the list comprehension that assigns _get_proxy_model_info's - return back into all_models (#28382 review).""" + """/v1/model/info list path (no litellm_model_id) must include team-scoped + deployments from the router model list and surface the public name (#28382).""" + team_row = _team_row() + global_row = { + "model_name": "gpt-4o", + "litellm_params": {"model": "gpt-4o"}, + "model_info": {"id": "normal-id-1", "db_model": False}, + } router = MagicMock() - router.get_model_names.return_value = ["team-claude-sonnet"] + router.model_list = [team_row, global_row] + router.get_model_names.return_value = ["gpt-4o"] router.get_model_access_groups.return_value = {} - router.get_model_list.return_value = [_team_row()] monkeypatch.setattr(ps, "user_model", None) - monkeypatch.setattr(ps, "llm_model_list", [_team_row()]) + monkeypatch.setattr(ps, "llm_model_list", router.model_list) monkeypatch.setattr(ps, "llm_router", router) - monkeypatch.setattr(ps, "get_key_models", lambda **kw: []) - monkeypatch.setattr(ps, "get_team_models", lambda **kw: []) monkeypatch.setattr( - ps, "get_complete_model_list", lambda **kw: ["team-claude-sonnet"] + ps, "_enrich_model_info_with_litellm_data", lambda model, **kw: model ) admin = UserAPIKeyAuth( @@ -176,3 +179,417 @@ async def test_model_info_v1_list_path_translates_team_model_name(monkeypatch): names = [m["model_name"] for m in resp["data"]] assert "team-claude-sonnet" in names assert "model_name_team-abc-123_4a6b8" not in names + + +@pytest.mark.asyncio +async def test_model_info_v1_unrestricted_key_returns_all_deployments(monkeypatch): + """Unrestricted keys must see all router deployments (legacy v1 access logic).""" + deployment = { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "global-id-1", "db_model": False}, + } + router = MagicMock() + router.model_list = [deployment] + router.get_model_names.return_value = ["gpt-4"] + router.get_model_access_groups.return_value = {} + + monkeypatch.setattr(ps, "user_model", None) + monkeypatch.setattr(ps, "llm_model_list", router.model_list) + monkeypatch.setattr(ps, "llm_router", router) + monkeypatch.setattr( + ps, "_enrich_model_info_with_litellm_data", lambda model, **kw: model + ) + + caller = UserAPIKeyAuth( + user_id="user-1", + user_role=LitellmUserRoles.INTERNAL_USER, + models=[], + team_models=[], + ) + resp = await ps.model_info_v1(user_api_key_dict=caller, litellm_model_id=None) + + assert [m["model_name"] for m in resp["data"]] == ["gpt-4"] + + +@pytest.mark.asyncio +async def test_model_info_v1_restricted_key_filters_deployments(monkeypatch): + """Key-level model allowlists must filter router deployments.""" + team_row = _team_row() + global_row = { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "global-id-1", "db_model": False}, + } + router = MagicMock() + router.model_list = [team_row, global_row] + router.get_model_names.return_value = ["gpt-4", "team-claude-sonnet"] + router.get_model_access_groups.return_value = {} + + monkeypatch.setattr(ps, "user_model", None) + monkeypatch.setattr(ps, "llm_model_list", router.model_list) + monkeypatch.setattr(ps, "llm_router", router) + monkeypatch.setattr( + ps, "_enrich_model_info_with_litellm_data", lambda model, **kw: model + ) + + caller = UserAPIKeyAuth( + user_id="user-1", + user_role=LitellmUserRoles.INTERNAL_USER, + models=["gpt-4"], + team_models=[], + ) + resp = await ps.model_info_v1(user_api_key_dict=caller, litellm_model_id=None) + + assert [m["model_name"] for m in resp["data"]] == ["gpt-4"] + + +def _other_team_row() -> dict: + return { + "model_name": "model_name_team-other_9f2c1", + "litellm_params": { + "model": "azure/gpt-5.2-low-rpm-testing", + "api_base": "https://team-other-private.example.com", + }, + "model_info": { + "id": "byok-id-other", + "team_id": "team-other", + "team_public_model_name": "team-claude-sonnet", + "db_model": True, + }, + } + + +@pytest.mark.asyncio +async def test_model_info_v1_unrestricted_key_hides_other_team_byok(monkeypatch): + """Unrestricted non-admin keys must not enumerate other teams' BYOK + deployments, but must still see global models and their own team's.""" + team_row = _team_row() + other_team_row = _other_team_row() + global_row = { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "global-id-1", "db_model": False}, + } + router = MagicMock() + router.model_list = [team_row, other_team_row, global_row] + router.get_model_names.return_value = ["gpt-4"] + router.get_model_access_groups.return_value = {} + + prisma_client = MagicMock() + caller_user_row = MagicMock() + caller_user_row.teams = ["team-abc-123"] + caller_user_row.model_dump.return_value = { + "user_id": "user-1", + "teams": ["team-abc-123"], + "models": [], + } + prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=caller_user_row + ) + + monkeypatch.setattr(ps, "user_model", None) + monkeypatch.setattr(ps, "llm_model_list", router.model_list) + monkeypatch.setattr(ps, "llm_router", router) + monkeypatch.setattr(ps, "prisma_client", prisma_client) + monkeypatch.setattr(ps, "get_all_team_models", AsyncMock(return_value={})) + monkeypatch.setattr( + ps, "_enrich_model_info_with_litellm_data", lambda model, **kw: model + ) + + caller = UserAPIKeyAuth( + user_id="user-1", + user_role=LitellmUserRoles.INTERNAL_USER, + models=[], + team_models=[], + ) + resp = await ps.model_info_v1(user_api_key_dict=caller, litellm_model_id=None) + + returned_ids = {m["model_info"]["id"] for m in resp["data"]} + assert returned_ids == {"global-id-1", "byok-id-1"} + assert "byok-id-other" not in returned_ids + names = [m["model_name"] for m in resp["data"]] + assert "team-claude-sonnet" in names + assert "gpt-4" in names + + +@pytest.mark.asyncio +async def test_model_info_v1_service_key_hides_all_team_byok(monkeypatch): + """A key without a resolvable user (e.g. CI/service token) sees only + global deployments, never any team-scoped BYOK rows.""" + team_row = _team_row() + other_team_row = _other_team_row() + global_row = { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "global-id-1", "db_model": False}, + } + router = MagicMock() + router.model_list = [team_row, other_team_row, global_row] + router.get_model_names.return_value = ["gpt-4"] + router.get_model_access_groups.return_value = {} + + prisma_client = MagicMock() + + monkeypatch.setattr(ps, "user_model", None) + monkeypatch.setattr(ps, "llm_model_list", router.model_list) + monkeypatch.setattr(ps, "llm_router", router) + monkeypatch.setattr(ps, "prisma_client", prisma_client) + monkeypatch.setattr( + ps, "_enrich_model_info_with_litellm_data", lambda model, **kw: model + ) + + caller = UserAPIKeyAuth( + user_id=None, + user_role=LitellmUserRoles.INTERNAL_USER, + team_id="team-abc-123", + models=[], + team_models=[], + ) + resp = await ps.model_info_v1(user_api_key_dict=caller, litellm_model_id=None) + + assert [m["model_info"]["id"] for m in resp["data"]] == ["global-id-1"] + + +@pytest.mark.asyncio +async def test_model_info_v1_populates_access_via_team_ids(monkeypatch): + """`/v1/model/info` must populate access_via_team_ids when the DB is connected.""" + team_id = "team-abc-123" + team_row = _team_row() + global_row = { + "model_name": "gpt-4o", + "litellm_params": {"model": "gpt-4o"}, + "model_info": {"id": "global-id-1", "db_model": False}, + } + router = MagicMock() + router.model_list = [team_row, global_row] + router.get_model_names.return_value = ["gpt-4o", "team-claude-sonnet"] + router.get_model_access_groups.return_value = {} + router.get_model_ids.return_value = ["global-id-1"] + + prisma_client = MagicMock() + + async def _fake_populate(**kwargs): + for model in kwargs["all_models"]: + model_id = model["model_info"]["id"] + if model_id == "byok-id-1": + model["model_info"]["access_via_team_ids"] = [team_id] + model["model_info"]["direct_access"] = False + elif model_id == "global-id-1": + model["model_info"]["direct_access"] = True + return kwargs["all_models"] + + monkeypatch.setattr(ps, "user_model", None) + monkeypatch.setattr(ps, "llm_model_list", router.model_list) + monkeypatch.setattr(ps, "llm_router", router) + monkeypatch.setattr(ps, "prisma_client", prisma_client) + monkeypatch.setattr(ps, "_populate_team_access_on_models", _fake_populate) + monkeypatch.setattr( + ps, "_enrich_model_info_with_litellm_data", lambda model, **kw: model + ) + + admin = UserAPIKeyAuth( + user_id="u", user_role=LitellmUserRoles.PROXY_ADMIN, team_models=[] + ) + resp = await ps.model_info_v1(user_api_key_dict=admin, litellm_model_id=None) + + by_id = {m["model_info"]["id"]: m for m in resp["data"]} + assert by_id["byok-id-1"]["model_info"]["access_via_team_ids"] == [team_id] + assert by_id["byok-id-1"]["model_info"]["direct_access"] is False + assert by_id["global-id-1"]["model_info"]["direct_access"] is True + + +@pytest.mark.asyncio +async def test_populate_team_access_sets_direct_access_false_by_default(monkeypatch): + """Team-accessible models without direct access must return direct_access=false.""" + team_row = _team_row() + global_row = { + "model_name": "gpt-4o", + "litellm_params": {"model": "gpt-4o"}, + "model_info": {"id": "global-id-1", "db_model": False}, + } + router = MagicMock() + router.get_model_ids.return_value = ["global-id-1"] + monkeypatch.setattr( + ps, + "get_all_team_models", + AsyncMock(return_value={"byok-id-1": ["team-abc-123"]}), + ) + + admin = UserAPIKeyAuth( + user_id="u", user_role=LitellmUserRoles.PROXY_ADMIN, team_models=[] + ) + result = await ps._populate_team_access_on_models( + user_api_key_dict=admin, + prisma_client=MagicMock(), + llm_router=router, + all_models=[team_row, global_row], + ) + + by_id = {m["model_info"]["id"]: m for m in result} + assert by_id["byok-id-1"]["model_info"]["direct_access"] is False + assert by_id["global-id-1"]["model_info"]["direct_access"] is True + + +@pytest.mark.asyncio +async def test_model_info_v1_team_id_without_db_fails_fast(monkeypatch): + """`teamId` without a connected DB raises 500 before any enrichment work runs.""" + router = MagicMock() + router.model_list = [_team_row()] + + enrich_spy = MagicMock(side_effect=lambda model, **kw: model) + + monkeypatch.setattr(ps, "user_model", None) + monkeypatch.setattr(ps, "llm_model_list", router.model_list) + monkeypatch.setattr(ps, "llm_router", router) + monkeypatch.setattr(ps, "prisma_client", None) + monkeypatch.setattr(ps, "_enrich_model_info_with_litellm_data", enrich_spy) + + admin = UserAPIKeyAuth( + user_id="u", user_role=LitellmUserRoles.PROXY_ADMIN, team_models=[] + ) + + with pytest.raises(ps.HTTPException) as exc_info: + await ps.model_info_v1( + user_api_key_dict=admin, litellm_model_id=None, teamId="team-abc-123" + ) + + assert exc_info.value.status_code == 500 + assert "DB not connected" in exc_info.value.detail["error"] + enrich_spy.assert_not_called() + + +@pytest.mark.asyncio +async def test_model_info_v1_include_team_models_without_db_fails_fast(monkeypatch): + """`include_team_models` without a connected DB raises 500 instead of silently + returning an empty list (the access fields can only be populated from the DB).""" + router = MagicMock() + router.model_list = [_team_row()] + + enrich_spy = MagicMock(side_effect=lambda model, **kw: model) + + monkeypatch.setattr(ps, "user_model", None) + monkeypatch.setattr(ps, "llm_model_list", router.model_list) + monkeypatch.setattr(ps, "llm_router", router) + monkeypatch.setattr(ps, "prisma_client", None) + monkeypatch.setattr(ps, "_enrich_model_info_with_litellm_data", enrich_spy) + + admin = UserAPIKeyAuth( + user_id="u", user_role=LitellmUserRoles.PROXY_ADMIN, team_models=[] + ) + + with pytest.raises(ps.HTTPException) as exc_info: + await ps.model_info_v1( + user_api_key_dict=admin, litellm_model_id=None, include_team_models=True + ) + + assert exc_info.value.status_code == 500 + assert "DB not connected" in exc_info.value.detail["error"] + enrich_spy.assert_not_called() + + +@pytest.mark.asyncio +async def test_model_info_v1_litellm_model_id_team_id_without_db_fails_fast( + monkeypatch, +): + """`litellm_model_id` + `teamId` without a connected DB must raise 500 too, not + return 200 with a model dict missing direct_access/access_via_team_ids.""" + router = MagicMock() + router.model_list = [_team_row()] + + monkeypatch.setattr(ps, "user_model", None) + monkeypatch.setattr(ps, "llm_model_list", router.model_list) + monkeypatch.setattr(ps, "llm_router", router) + monkeypatch.setattr(ps, "prisma_client", None) + + admin = UserAPIKeyAuth( + user_id="u", user_role=LitellmUserRoles.PROXY_ADMIN, team_models=[] + ) + + with pytest.raises(ps.HTTPException) as exc_info: + await ps.model_info_v1( + user_api_key_dict=admin, + litellm_model_id="byok-id-1", + teamId="team-abc-123", + ) + + assert exc_info.value.status_code == 500 + assert "DB not connected" in exc_info.value.detail["error"] + router.get_deployment.assert_not_called() + + +@pytest.mark.asyncio +async def test_model_info_v1_litellm_model_id_include_team_models_filters_inaccessible( + monkeypatch, +): + """`litellm_model_id` + `include_team_models` must drop a model the caller cannot + use instead of returning it unconditionally from the single-model lookup.""" + team_row = _team_row() + + router = MagicMock() + deployment = MagicMock() + deployment.model_dump.return_value = team_row + router.get_deployment.return_value = deployment + + async def _fake_populate(**kwargs): + for model in kwargs["all_models"]: + model["model_info"]["direct_access"] = False + model["model_info"]["access_via_team_ids"] = [] + return kwargs["all_models"] + + monkeypatch.setattr(ps, "user_model", None) + monkeypatch.setattr(ps, "llm_model_list", [team_row]) + monkeypatch.setattr(ps, "llm_router", router) + monkeypatch.setattr(ps, "prisma_client", MagicMock()) + monkeypatch.setattr(ps, "_get_proxy_model_info", lambda model: team_row) + monkeypatch.setattr(ps, "_populate_team_access_on_models", _fake_populate) + + caller = UserAPIKeyAuth( + user_id="u", user_role=LitellmUserRoles.INTERNAL_USER, team_models=[] + ) + resp = await ps.model_info_v1( + user_api_key_dict=caller, + litellm_model_id="byok-id-1", + include_team_models=True, + ) + + assert resp["data"] == [] + + +@pytest.mark.asyncio +async def test_model_info_v1_litellm_model_id_team_id_applies_team_filter(monkeypatch): + """`litellm_model_id` + `teamId` must run the teamId filter on the single model + rather than returning it regardless of the team's access.""" + team_row = _team_row() + + router = MagicMock() + deployment = MagicMock() + deployment.model_dump.return_value = team_row + router.get_deployment.return_value = deployment + + async def _fake_populate(**kwargs): + return kwargs["all_models"] + + team_filter = AsyncMock(return_value=[]) + + monkeypatch.setattr(ps, "user_model", None) + monkeypatch.setattr(ps, "llm_model_list", [team_row]) + monkeypatch.setattr(ps, "llm_router", router) + monkeypatch.setattr(ps, "prisma_client", MagicMock()) + monkeypatch.setattr(ps, "_get_proxy_model_info", lambda model: team_row) + monkeypatch.setattr(ps, "_populate_team_access_on_models", _fake_populate) + monkeypatch.setattr(ps, "_filter_models_by_team_id", team_filter) + + admin = UserAPIKeyAuth( + user_id="u", user_role=LitellmUserRoles.PROXY_ADMIN, team_models=[] + ) + resp = await ps.model_info_v1( + user_api_key_dict=admin, + litellm_model_id="byok-id-1", + teamId="other-team", + ) + + assert resp["data"] == [] + team_filter.assert_awaited_once() + assert team_filter.await_args.kwargs["team_id"] == "other-team" + assert team_filter.await_args.kwargs["all_models"] == [team_row] diff --git a/tests/test_litellm/proxy/test_model_info_default_limits.py b/tests/test_litellm/proxy/test_model_info_default_limits.py index 641199c96f05..8111a7af0065 100644 --- a/tests/test_litellm/proxy/test_model_info_default_limits.py +++ b/tests/test_litellm/proxy/test_model_info_default_limits.py @@ -146,9 +146,9 @@ async def test_model_info_endpoint_returns_defaults_in_full_model_list(self): deployment_dict = deployment.model_dump(exclude_none=True) mock_router = MagicMock() + mock_router.model_list = [deployment_dict] mock_router.get_model_names.return_value = ["model1"] mock_router.get_model_access_groups.return_value = {} - mock_router.get_model_list.return_value = [deployment_dict] user_api_key_dict = UserAPIKeyAuth(api_key="sk-test") @@ -156,6 +156,7 @@ async def test_model_info_endpoint_returns_defaults_in_full_model_list(self): patch("litellm.proxy.proxy_server.llm_router", mock_router), patch("litellm.proxy.proxy_server.llm_model_list", [deployment_dict]), patch("litellm.proxy.proxy_server.user_model", None), + patch("litellm.proxy.proxy_server.prisma_client", None), patch("litellm.proxy.proxy_server.get_key_models", return_value=["model1"]), patch( "litellm.proxy.proxy_server.get_team_models", return_value=["model1"] diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index 4fb725b7ef3c..34c88e2fd33a 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -795,6 +795,127 @@ def test_db_connection_extra_params_forwarded_to_url( assert appended_params["pgbouncer"] == "true" assert appended_params["statement_cache_size"] == 0 + def test_build_db_connection_url_params_disable_prepared_statements(self): + from litellm.proxy.proxy_cli import _build_db_connection_url_params + + params = _build_db_connection_url_params( + connection_limit=10, + pool_timeout=60, + disable_prepared_statements=True, + ) + assert params["pgbouncer"] == "true" + + def test_build_db_connection_url_params_no_pgbouncer_by_default(self): + from litellm.proxy.proxy_cli import _build_db_connection_url_params + + params = _build_db_connection_url_params( + connection_limit=10, + pool_timeout=60, + ) + assert "pgbouncer" not in params + + def test_build_db_connection_url_params_extra_pgbouncer_overrides_flag(self): + from litellm.proxy.proxy_cli import _build_db_connection_url_params + + params = _build_db_connection_url_params( + connection_limit=10, + pool_timeout=60, + disable_prepared_statements=True, + extra_params={"pgbouncer": "false"}, + ) + assert params["pgbouncer"] == "false" + + @pytest.mark.parametrize( + "config_value, expect_pgbouncer", + [ + (True, True), + (False, False), + ("true", True), + ("false", False), + ("not-a-bool", False), + ], + ) + @patch("subprocess.run") + @patch("atexit.register") + @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") + @patch( + "litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False + ) + def test_disable_prepared_statements_forwarded_to_url( + self, + mock_should_update, + mock_setup_db, + mock_atexit_register, + mock_subprocess_run, + config_value, + expect_pgbouncer, + ): + from click.testing import CliRunner + + from litellm.proxy.proxy_cli import run_server + + runner = CliRunner() + mock_subprocess_run.return_value = MagicMock(returncode=0) + + mock_proxy_module = MagicMock( + app=MagicMock(), + ProxyConfig=MagicMock(), + KeyManagementSettings=MagicMock(), + save_worker_config=MagicMock(), + ) + mock_proxy_module.ProxyConfig.return_value.get_config = AsyncMock( + return_value={ + "general_settings": { + "database_url": "postgresql://test:test@localhost:5432/test", + "database_disable_prepared_statements": config_value, + } + } + ) + + clean_env = { + k: v + for k, v in os.environ.items() + if k not in ("DATABASE_URL", "DIRECT_URL") + } + + with ( + patch.dict(os.environ, clean_env, clear=True), + patch.dict( + "sys.modules", + { + "proxy_server": mock_proxy_module, + "litellm.proxy.proxy_server": mock_proxy_module, + }, + ), + patch( + "litellm.proxy.proxy_cli.ProxyInitializationHelpers._get_default_unvicorn_init_args" + ) as mock_get_args, + patch( + "litellm.proxy.proxy_cli.append_query_params", + side_effect=lambda url, params: str(url), + ) as mock_append_query_params, + ): + mock_get_args.return_value = { + "app": "litellm.proxy.proxy_server:app", + "host": "localhost", + "port": 8000, + } + + result = runner.invoke( + run_server, + ["--local", "--config", "test-config.yaml", "--skip_server_startup"], + ) + + assert ( + result.exit_code == 0 + ), f"exit_code={result.exit_code}, output={result.output}" + mock_append_query_params.assert_called() + appended_params = mock_append_query_params.call_args.args[1] + if expect_pgbouncer: + assert appended_params["pgbouncer"] == "true" + else: + assert "pgbouncer" not in appended_params + @patch("uvicorn.run") @patch("atexit.register") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 8aa839cdfcb6..13fb125fc27b 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -3810,14 +3810,15 @@ async def test_model_info_v1_oci_secrets_not_leaked(): # Mock the llm_router to return our test data mock_router = MagicMock() + mock_router.model_list = [mock_model_data] mock_router.get_model_names.return_value = ["oci-grok-test"] mock_router.get_model_access_groups.return_value = {} - mock_router.get_model_list.return_value = [mock_model_data] # Mock global variables with ( patch("litellm.proxy.proxy_server.llm_router", mock_router), patch("litellm.proxy.proxy_server.llm_model_list", [mock_model_data]), + patch("litellm.proxy.proxy_server.prisma_client", None), patch( "litellm.proxy.proxy_server.general_settings", {"infer_model_from_keys": False}, diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py index 7e7e98d13607..08d1ef619a79 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py @@ -15,6 +15,7 @@ import hashlib import json +from datetime import datetime, timedelta, timezone from types import SimpleNamespace from typing import Any from unittest.mock import AsyncMock, MagicMock @@ -22,6 +23,7 @@ import pytest from fastapi import HTTPException +from litellm.proxy._types import LiteLLM_VerificationTokenView from litellm.proxy.utils import PrismaClient @@ -193,6 +195,7 @@ async def test_query_first_with_cached_plan_fallback_happy_returns_row( ) -> None: expected = {"token": "abc", "team_spend": 1.0, "team_max_budget": 5.0} prisma_client.db.query_first = AsyncMock(return_value=expected) + prisma_client.attempt_db_reconnect = AsyncMock(return_value=True) result = await prisma_client._query_first_with_cached_plan_fallback( "SELECT * FROM x WHERE token = $1", "abc" ) @@ -208,35 +211,110 @@ async def test_query_first_with_cached_plan_fallback_happy_returns_row( "args": ("SELECT * FROM x WHERE token = $1", "abc"), "matches": True, } + prisma_client.attempt_db_reconnect.assert_not_awaited() @pytest.mark.asyncio -async def test_query_first_with_cached_plan_fallback_retries_on_cached_plan_error( +async def test_query_first_with_cached_plan_fallback_reconnects_then_retries_identical_query( prisma_client: PrismaClient, ) -> None: + original_query = 'SELECT * FROM "LiteLLM_VerificationToken" WHERE v.token = $1' expected = {"token": "abc", "team_spend": 1.0, "team_max_budget": 5.0} - prisma_client.db.query_first = AsyncMock( + manager = MagicMock() + query_first = AsyncMock( side_effect=[ RuntimeError("cached plan must not change result type"), expected, ] ) + reconnect = AsyncMock(return_value=True) + manager.attach_mock(query_first, "query_first") + manager.attach_mock(reconnect, "attempt_db_reconnect") + prisma_client.db.query_first = query_first + prisma_client.attempt_db_reconnect = reconnect + result = await prisma_client._query_first_with_cached_plan_fallback( - "SELECT * FROM x WHERE token = $1", "abc" + original_query, "abc" ) + + assert result == expected + assert query_first.await_count == 2 + first_call, retry_call = query_first.await_args_list + assert retry_call.args == first_call.args == (original_query, "abc") + reconnect.assert_awaited_once() + assert reconnect.await_args.kwargs.get("force", False) is False + assert [name for name, *_ in manager.mock_calls] == [ + "query_first", + "attempt_db_reconnect", + "query_first", + ] + + +@pytest.mark.asyncio +async def test_query_first_with_cached_plan_fallback_never_deallocates( + prisma_client: PrismaClient, +) -> None: + expected = {"token": "abc"} + prisma_client.db.query_first = AsyncMock( + side_effect=[ + RuntimeError("cached plan must not change result type"), + expected, + ] + ) + prisma_client.db.execute_raw = AsyncMock(return_value=0) + prisma_client.attempt_db_reconnect = AsyncMock(return_value=True) + + await prisma_client._query_first_with_cached_plan_fallback("SELECT 1") + + prisma_client.db.execute_raw.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_query_first_with_cached_plan_fallback_propagates_when_retry_also_fails( + prisma_client: PrismaClient, +) -> None: + plan_error = RuntimeError("cached plan must not change result type") + prisma_client.db.query_first = AsyncMock(side_effect=[plan_error, plan_error]) + prisma_client.attempt_db_reconnect = AsyncMock(return_value=True) + + with pytest.raises(RuntimeError, match="cached plan must not change result type"): + await prisma_client._query_first_with_cached_plan_fallback("SELECT 1") + + assert prisma_client.db.query_first.await_count == 2 + prisma_client.attempt_db_reconnect.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_query_first_with_cached_plan_fallback_retries_when_reconnect_returns_false( + prisma_client: PrismaClient, +) -> None: + expected = {"token": "abc"} + prisma_client.db.query_first = AsyncMock( + side_effect=[ + RuntimeError("cached plan must not change result type"), + expected, + ] + ) + prisma_client.attempt_db_reconnect = AsyncMock(return_value=False) + + result = await prisma_client._query_first_with_cached_plan_fallback("SELECT 1") + assert result == expected assert prisma_client.db.query_first.await_count == 2 - second_call_sql = prisma_client.db.query_first.await_args_list[1].args[0] - assert "cache_invalidated_" in second_call_sql @pytest.mark.asyncio async def test_query_first_with_cached_plan_fallback_reraises_non_plan_errors( prisma_client: PrismaClient, ) -> None: - prisma_client.db.query_first = AsyncMock(side_effect=RuntimeError("totally unrelated")) + prisma_client.db.query_first = AsyncMock( + side_effect=RuntimeError("totally unrelated") + ) + prisma_client.attempt_db_reconnect = AsyncMock(return_value=True) with pytest.raises(RuntimeError, match="totally unrelated"): await prisma_client._query_first_with_cached_plan_fallback("SELECT 1") + assert prisma_client.db.query_first.await_count == 1 + prisma_client.attempt_db_reconnect.assert_not_awaited() @pytest.mark.asyncio @@ -351,7 +429,9 @@ async def test_get_data_token_find_unique_returns_record( async def test_get_data_token_find_unique_missing_token_raises_401( prisma_client: PrismaClient, ) -> None: - prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=None) + prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + return_value=None + ) with pytest.raises(HTTPException) as excinfo: await prisma_client.get_data(token="sk-missing", table_name="key") err = excinfo.value @@ -398,3 +478,39 @@ async def test_get_data_logs_and_raises_on_db_error( ) with pytest.raises(RuntimeError, match="network split"): await prisma_client.get_data(token="sk-broken", table_name="key") + + +@pytest.mark.asyncio +async def test_get_data_combined_view_returns_view_for_deprecated_key( + prisma_client: PrismaClient, +) -> None: + """Grace-period rotation, full get_data flow: the old hash misses the + combined view, the deprecated-key table resolves it to the active token, + and get_data must return the recursive lookup's finished view instead of + re-running dict normalization on it (which raised TypeError and turned + every grace-period request into a 401).""" + old_hash = "hashed-old-token-grace-e2e" + active_hash = "hashed-active-token-grace-e2e" + active_row = { + "token": active_hash, + "team_models": None, + "team_blocked": None, + "team_members_with_roles": None, + "user_id": None, + "expires": None, + } + prisma_client.db.query_first = AsyncMock(side_effect=[None, active_row]) + prisma_client.db.litellm_deprecatedverificationtoken = MagicMock() + prisma_client.db.litellm_deprecatedverificationtoken.find_first = AsyncMock( + return_value=SimpleNamespace( + active_token_id=active_hash, + revoke_at=datetime.now(timezone.utc) + timedelta(hours=1), + ) + ) + + response = await prisma_client.get_data( + token=old_hash, table_name="combined_view", query_type="find_unique" + ) + + assert isinstance(response, LiteLLM_VerificationTokenView) + assert response.token == active_hash diff --git a/tests/test_litellm/test_claude_fable_5_config.py b/tests/test_litellm/test_claude_fable_5_config.py new file mode 100644 index 000000000000..d8d95fba0da2 --- /dev/null +++ b/tests/test_litellm/test_claude_fable_5_config.py @@ -0,0 +1,230 @@ +""" +Validate Claude Fable 5 model configuration entries. + +Fable 5 is a new tier above Opus ($10/$50 per MTok) with the same adaptive-only +API surface as Opus 4.7/4.8. The cost-map entries below are what make the model +resolvable across Anthropic, Bedrock, Vertex AI, and Azure AI (Microsoft +Foundry), and the ``supports_adaptive_thinking`` flag is what makes LiteLLM send +``thinking.type='adaptive'`` instead of the legacy ``enabled``/``budget_tokens`` +shape, which Fable 5 rejects with a 400. +""" + +import json +import os + +import pytest + +import litellm +from litellm.constants import BEDROCK_CONVERSE_MODELS +from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap + +REPO_ROOT = os.path.join(os.path.dirname(__file__), "../..") + + +def _load_root_cost_map() -> dict: + json_path = os.path.join(REPO_ROOT, "model_prices_and_context_window.json") + with open(json_path) as f: + return json.load(f) + + +@pytest.fixture +def local_model_cost_map(monkeypatch): + """Force the bundled backup cost map so assertions don't depend on the + network-fetched ``main`` copy (which lags this branch until merge).""" + original_model_cost = litellm.model_cost + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + litellm.model_cost = litellm.get_model_cost_map(url="") + litellm.get_model_info.cache_clear() + try: + yield + finally: + litellm.model_cost = original_model_cost + litellm.get_model_info.cache_clear() + + +def test_fable_5_model_pricing_and_capabilities(): + model_data = _load_root_cost_map() + + expected_models = [ + ("claude-fable-5", "anthropic"), + ("anthropic.claude-fable-5", "bedrock_converse"), + ("vertex_ai/claude-fable-5", "vertex_ai-anthropic_models"), + # Unlike Opus 4.8 (200k on Foundry), Fable 5 has the full 1M context + # window on Microsoft Foundry. + ("azure_ai/claude-fable-5", "azure_ai"), + ] + + for model_name, provider in expected_models: + assert model_name in model_data, f"Missing model entry: {model_name}" + info = model_data[model_name] + + assert info["litellm_provider"] == provider + assert info["mode"] == "chat" + assert info["max_input_tokens"] == 1000000 + assert info["max_output_tokens"] == 128000 + assert info["max_tokens"] == 128000 + + # $10 / $50 per MTok (2x Opus 4.8), with the standard 1.25x 5m + # cache-write, 2x 1h cache-write, and 0.1x cache-read multipliers. + assert info["input_cost_per_token"] == 1e-05 + assert info["output_cost_per_token"] == 5e-05 + assert info["cache_creation_input_token_cost"] == 1.25e-05 + assert info["cache_creation_input_token_cost_above_1hr"] == 2e-05 + assert info["cache_read_input_token_cost"] == 1e-06 + + # Flat-rate across the full 1M context window. + assert "input_cost_per_token_above_200k_tokens" not in info + assert "output_cost_per_token_above_200k_tokens" not in info + + assert info["supports_assistant_prefill"] is False + assert info["supports_function_calling"] is True + assert info["supports_prompt_caching"] is True + assert info["supports_reasoning"] is True + assert info["supports_tool_choice"] is True + assert info["supports_vision"] is True + assert info["supports_xhigh_reasoning_effort"] is True + assert info["supports_max_reasoning_effort"] is True + + +def test_fable_5_bedrock_regional_model_pricing(): + model_data = _load_root_cost_map() + + # Fable 5 launched with us/eu geo inference profiles plus a global profile + # (no au/apac/jp). Global uses base pricing; geo profiles carry the + # standard 10% regional premium. + expected_models = { + "global.anthropic.claude-fable-5": { + "input_cost_per_token": 1e-05, + "output_cost_per_token": 5e-05, + "cache_creation_input_token_cost": 1.25e-05, + "cache_read_input_token_cost": 1e-06, + }, + "us.anthropic.claude-fable-5": { + "input_cost_per_token": 1.1e-05, + "output_cost_per_token": 5.5e-05, + "cache_creation_input_token_cost": 1.375e-05, + "cache_read_input_token_cost": 1.1e-06, + }, + "eu.anthropic.claude-fable-5": { + "input_cost_per_token": 1.1e-05, + "output_cost_per_token": 5.5e-05, + "cache_creation_input_token_cost": 1.375e-05, + "cache_read_input_token_cost": 1.1e-06, + }, + } + + for model_name, expected in expected_models.items(): + assert model_name in model_data, f"Missing model entry: {model_name}" + info = model_data[model_name] + assert info["litellm_provider"] == "bedrock_converse" + assert info["max_input_tokens"] == 1000000 + assert info["max_output_tokens"] == 128000 + assert info["bedrock_output_config_effort_ceiling"] == "xhigh" + for key, value in expected.items(): + assert info[key] == value + + +def test_fable_5_geo_multiplier_without_fast_mode(): + """First-party ``inference_geo='us'`` carries the 1.1x premium, but unlike + the Opus line there is no fast-mode variant for Fable 5; a ``fast`` key + here would silently misprice ``speed='fast'`` requests.""" + model_data = _load_root_cost_map() + entry = model_data["claude-fable-5"]["provider_specific_entry"] + assert entry == {"us": 1.1} + + +def test_fable_5_present_in_bundled_backup(): + """The bundled backup is the runtime fallback (and what tests load with + ``LITELLM_LOCAL_MODEL_COST_MAP=True``) — it must carry the same entries as + the root cost map, otherwise the model resolves on one path but not the + other.""" + backup = GetModelCostMap.load_local_model_cost_map() + root = _load_root_cost_map() + for model_name in ( + "claude-fable-5", + "anthropic.claude-fable-5", + "global.anthropic.claude-fable-5", + "us.anthropic.claude-fable-5", + "eu.anthropic.claude-fable-5", + "vertex_ai/claude-fable-5", + "vertex_ai/claude-fable-5@default", + "azure_ai/claude-fable-5", + ): + assert model_name in backup, f"Missing from backup cost map: {model_name}" + assert backup[model_name] == root[model_name], model_name + + +def test_fable_5_registered_for_bedrock_converse(): + assert "anthropic.claude-fable-5" in BEDROCK_CONVERSE_MODELS + + +def test_fable_5_provider_resolves_via_model_info(local_model_cost_map): + info = litellm.get_model_info(model="claude-fable-5") + assert info["litellm_provider"] == "anthropic" + assert info["max_input_tokens"] == 1000000 + assert info["max_output_tokens"] == 128000 + + +@pytest.mark.parametrize( + "cost_map", + [_load_root_cost_map(), GetModelCostMap.load_local_model_cost_map()], + ids=["root", "bundled_backup"], +) +def test_fable_5_all_variants_carry_adaptive_thinking_flag(cost_map): + """Every Fable 5 entry must advertise ``supports_adaptive_thinking``. + + Adaptive-thinking detection is cost-map driven, so a single variant missing + the flag silently sends the legacy ``thinking.type='enabled'`` shape and the + provider 400s (issue #29188 for the Opus 4.8 equivalent). Fable 5 is even + stricter than Opus 4.8: an explicit ``thinking.type='disabled'`` also 400s, + so adaptive is the only valid thinking shape LiteLLM can emit for it.""" + variants = [k for k in cost_map if "claude-fable-5" in k] + assert variants, "no claude-fable-5 entries found in cost map" + missing = [ + k for k in variants if cost_map[k].get("supports_adaptive_thinking") is not True + ] + assert not missing, f"missing supports_adaptive_thinking: {missing}" + + +@pytest.mark.parametrize( + "model", + [ + "claude-fable-5", + "anthropic/claude-fable-5", + "anthropic.claude-fable-5", + "bedrock/us.anthropic.claude-fable-5", + "bedrock/invoke/eu.anthropic.claude-fable-5", + "bedrock/global.anthropic.claude-fable-5", + "vertex_ai/claude-fable-5", + "azure_ai/claude-fable-5", + ], +) +def test_adaptive_thinking_detected_for_fable_5(local_model_cost_map, model): + """Provider-routed ids must resolve to a flagged entry so ``reasoning_effort`` + maps to ``thinking.type='adaptive'`` + ``output_config.effort``.""" + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + assert AnthropicModelInfo._is_adaptive_thinking_model(model) is True + + +@pytest.mark.parametrize( + "cost_map", + [_load_root_cost_map(), GetModelCostMap.load_local_model_cost_map()], + ids=["root", "bundled_backup"], +) +def test_sampling_params_flag_on_all_models_that_removed_them(cost_map): + """Fable 5 and Opus 4.7/4.8 reject ``top_p``/``top_k``/``temperature != 1``; + the drop/raise gating is cost-map driven, so every variant must carry an + explicit ``supports_sampling_params: false``. The perplexity route is + exempt: it is OpenAI-compatible and maps sampling params upstream.""" + variants = [ + k + for k in cost_map + if any(v in k for v in ("claude-fable-5", "claude-opus-4-7", "claude-opus-4-8")) + and not k.startswith("perplexity/") + ] + assert variants, "no matching entries found in cost map" + missing = [ + k for k in variants if cost_map[k].get("supports_sampling_params") is not False + ] + assert not missing, f"missing supports_sampling_params=false: {missing}" diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index f179e9c8f93b..4c4d9e1133bc 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -858,6 +858,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "supports_xhigh_reasoning_effort": {"type": "boolean"}, "supports_max_reasoning_effort": {"type": "boolean"}, "supports_adaptive_thinking": {"type": "boolean"}, + "supports_sampling_params": {"type": "boolean"}, "supports_service_tier": {"type": "boolean"}, "supports_preset": {"type": "boolean"}, "supports_output_config": {"type": "boolean"}, diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index bbab73b07f1d..233741652a95 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -1351,11 +1351,6 @@ "count": 1 } }, - "src/components/mcp_tools/mcp_server_edit.test.tsx": { - "unused-imports/no-unused-imports": { - "count": 1 - } - }, "src/components/mcp_tools/mcp_server_edit.tsx": { "no-restricted-imports": { "count": 1 @@ -1517,11 +1512,6 @@ "count": 1 } }, - "src/components/organisms/RegenerateKeyModal.tsx": { - "react-hooks/set-state-in-effect": { - "count": 1 - } - }, "src/components/organisms/create_key_button.test.tsx": { "@typescript-eslint/no-require-imports": { "count": 2 diff --git a/ui/litellm-dashboard/package-lock.json b/ui/litellm-dashboard/package-lock.json index 7000efd63b49..70f8151cdf8d 100644 --- a/ui/litellm-dashboard/package-lock.json +++ b/ui/litellm-dashboard/package-lock.json @@ -52,8 +52,8 @@ "@types/react-dom": "18.3.7", "@types/react-syntax-highlighter": "15.5.13", "@types/uuid": "10.0.0", - "@vitest/coverage-v8": "3.2.4", - "@vitest/ui": "3.2.4", + "@vitest/coverage-v8": "3.2.6", + "@vitest/ui": "3.2.6", "autoprefixer": "10.4.24", "dotenv": "17.2.3", "eslint": "9.39.2", @@ -69,7 +69,7 @@ "typescript": "5.9.3", "typescript-eslint": "8.60.1", "vite": "7.3.2", - "vitest": "3.2.4" + "vitest": "3.2.6" }, "engines": { "node": ">=20.9.0", @@ -4276,9 +4276,9 @@ ] }, "node_modules/@vitest/coverage-v8": { - "version": "3.2.4", - "resolved": "https://registry.npmjs.org/@vitest/coverage-v8/-/coverage-v8-3.2.4.tgz", - "integrity": "sha512-EyF9SXU6kS5Ku/U82E259WSnvg6c8KTjppUncuNdm5QHpe17mwREHnjDzozC8x9MZ0xfBUFSaLkRv4TMA75ALQ==", + "version": "3.2.6", + "resolved": "https://registry.npmjs.org/@vitest/coverage-v8/-/coverage-v8-3.2.6.tgz", + "integrity": "sha512-LsAdmUapA0qSN306d8+zOyawM0hFm2m2Hg9IwVNIKBm+qJV8cijiq2c+gxKZcB1HCfIWAy+0qEZDCUQA58A1cw==", "dev": true, "license": "MIT", "dependencies": { @@ -4300,8 +4300,8 @@ "url": "https://opencollective.com/vitest" }, "peerDependencies": { - "@vitest/browser": "3.2.4", - "vitest": "3.2.4" + "@vitest/browser": "3.2.6", + "vitest": "3.2.6" }, "peerDependenciesMeta": { "@vitest/browser": { @@ -4310,15 +4310,15 @@ } }, "node_modules/@vitest/expect": { - "version": "3.2.4", - "resolved": "https://registry.npmjs.org/@vitest/expect/-/expect-3.2.4.tgz", - "integrity": "sha512-Io0yyORnB6sikFlt8QW5K7slY4OjqNX9jmJQ02QDda8lyM6B5oNgVWoSoKPac8/kgnCUzuHQKrSLtu/uOqqrig==", + "version": "3.2.6", + "resolved": "https://registry.npmjs.org/@vitest/expect/-/expect-3.2.6.tgz", + "integrity": "sha512-1+7q9BtaKzEmO+fmNT3kYvoNn5Y71XWAx2Q5HRim4tTVRQVRv4uJFAQ5FbK0OPUeNP/WmVCpxYxoJdvuHVjzBQ==", "dev": true, "license": "MIT", "dependencies": { "@types/chai": "^5.2.2", - "@vitest/spy": "3.2.4", - "@vitest/utils": "3.2.4", + "@vitest/spy": "3.2.6", + "@vitest/utils": "3.2.6", "chai": "^5.2.0", "tinyrainbow": "^2.0.0" }, @@ -4327,13 +4327,13 @@ } }, "node_modules/@vitest/mocker": { - "version": "3.2.4", - "resolved": "https://registry.npmjs.org/@vitest/mocker/-/mocker-3.2.4.tgz", - "integrity": "sha512-46ryTE9RZO/rfDd7pEqFl7etuyzekzEhUbTW3BvmeO/BcCMEgq59BKhek3dXDWgAj4oMK6OZi+vRr1wPW6qjEQ==", + "version": "3.2.6", + "resolved": "https://registry.npmjs.org/@vitest/mocker/-/mocker-3.2.6.tgz", + "integrity": "sha512-EZOrpDbkKotFAP7wPAQV1UIyoGOk4oX7ynWhBhLB7v+meMHbQhU16oPpIYGTTe4oFlhpryGpgpcZP/sin3hYuw==", "dev": true, "license": "MIT", "dependencies": { - "@vitest/spy": "3.2.4", + "@vitest/spy": "3.2.6", "estree-walker": "^3.0.3", "magic-string": "^0.30.17" }, @@ -4354,9 +4354,9 @@ } }, "node_modules/@vitest/pretty-format": { - "version": "3.2.4", - "resolved": "https://registry.npmjs.org/@vitest/pretty-format/-/pretty-format-3.2.4.tgz", - "integrity": "sha512-IVNZik8IVRJRTr9fxlitMKeJeXFFFN0JaB9PHPGQ8NKQbGpfjlTx9zO4RefN8gp7eqjNy8nyK3NZmBzOPeIxtA==", + "version": "3.2.6", + "resolved": "https://registry.npmjs.org/@vitest/pretty-format/-/pretty-format-3.2.6.tgz", + "integrity": "sha512-lb7XXXzmm2h2ASzFnRvQpDo6onT1NmMJA3tkGTWiBFtRJ9lxGY3d3mm/Apt36gej2bkkOVLL/yTOtufDaFa/jA==", "dev": true, "license": "MIT", "dependencies": { @@ -4367,13 +4367,13 @@ } }, "node_modules/@vitest/runner": { - "version": "3.2.4", - "resolved": "https://registry.npmjs.org/@vitest/runner/-/runner-3.2.4.tgz", - "integrity": "sha512-oukfKT9Mk41LreEW09vt45f8wx7DordoWUZMYdY/cyAk7w5TWkTRCNZYF7sX7n2wB7jyGAl74OxgwhPgKaqDMQ==", + "version": "3.2.6", + "resolved": "https://registry.npmjs.org/@vitest/runner/-/runner-3.2.6.tgz", + "integrity": "sha512-HYcoSj1w5tcgUnzoF0HcyaAQjpA1gj9ftUJ7iSJSuipc02jW9gKkigwZbjFldAfYHA1fa8UZVRftdMY5msWM9Q==", "dev": true, "license": "MIT", "dependencies": { - "@vitest/utils": "3.2.4", + "@vitest/utils": "3.2.6", "pathe": "^2.0.3", "strip-literal": "^3.0.0" }, @@ -4382,13 +4382,13 @@ } }, "node_modules/@vitest/snapshot": { - "version": "3.2.4", - "resolved": "https://registry.npmjs.org/@vitest/snapshot/-/snapshot-3.2.4.tgz", - "integrity": "sha512-dEYtS7qQP2CjU27QBC5oUOxLE/v5eLkGqPE0ZKEIDGMs4vKWe7IjgLOeauHsR0D5YuuycGRO5oSRXnwnmA78fQ==", + "version": "3.2.6", + "resolved": "https://registry.npmjs.org/@vitest/snapshot/-/snapshot-3.2.6.tgz", + "integrity": "sha512-H+ZjNTWGpObenh0YnlBctAPnJSI20P81PL8BPzWpx54YXLLTm8hEsWawtcYLMrwvpK48hGxLLbCS+1KRXhsKhw==", "dev": true, "license": "MIT", "dependencies": { - "@vitest/pretty-format": "3.2.4", + "@vitest/pretty-format": "3.2.6", "magic-string": "^0.30.17", "pathe": "^2.0.3" }, @@ -4397,9 +4397,9 @@ } }, "node_modules/@vitest/spy": { - "version": "3.2.4", - "resolved": "https://registry.npmjs.org/@vitest/spy/-/spy-3.2.4.tgz", - "integrity": "sha512-vAfasCOe6AIK70iP5UD11Ac4siNUNJ9i/9PZ3NKx07sG6sUxeag1LWdNrMWeKKYBLlzuK+Gn65Yd5nyL6ds+nw==", + "version": "3.2.6", + "resolved": "https://registry.npmjs.org/@vitest/spy/-/spy-3.2.6.tgz", + "integrity": "sha512-oq6BbH68WzcWmwtBrU9nqLeaXTR4XwJF7FSLkKEZo4i6eoXcrxjcwSuTvWBIRUTC6VC72nXYunzqgZA+IKdtxg==", "dev": true, "license": "MIT", "dependencies": { @@ -4410,13 +4410,13 @@ } }, "node_modules/@vitest/ui": { - "version": "3.2.4", - "resolved": "https://registry.npmjs.org/@vitest/ui/-/ui-3.2.4.tgz", - "integrity": "sha512-hGISOaP18plkzbWEcP/QvtRW1xDXF2+96HbEX6byqQhAUbiS5oH6/9JwW+QsQCIYON2bI6QZBF+2PvOmrRZ9wA==", + "version": "3.2.6", + "resolved": "https://registry.npmjs.org/@vitest/ui/-/ui-3.2.6.tgz", + "integrity": "sha512-mATfG3zVdhobE9U1rIpvtYD3DGuSSxqZ3Aj/8ityGqKXy8YDJ9BoAjZmAz6dZ1IZ1xI5V+MerkCczvVa+3QK9Q==", "dev": true, "license": "MIT", "dependencies": { - "@vitest/utils": "3.2.4", + "@vitest/utils": "3.2.6", "fflate": "^0.8.2", "flatted": "^3.3.3", "pathe": "^2.0.3", @@ -4428,17 +4428,17 @@ "url": "https://opencollective.com/vitest" }, "peerDependencies": { - "vitest": "3.2.4" + "vitest": "3.2.6" } }, "node_modules/@vitest/utils": { - "version": "3.2.4", - "resolved": "https://registry.npmjs.org/@vitest/utils/-/utils-3.2.4.tgz", - "integrity": "sha512-fB2V0JFrQSMsCo9HiSq3Ezpdv4iYaXRG1Sx8edX3MwxfyNn83mKiGzOcH+Fkxt4MHxr3y42fQi1oeAInqgX2QA==", + "version": "3.2.6", + "resolved": "https://registry.npmjs.org/@vitest/utils/-/utils-3.2.6.tgz", + "integrity": "sha512-lI23nIs4bnT3T8NIoh+vFaz5s2/DdP0Jgt2jxwgWljvwn82cLJtyi/If+fjFyoLMGIOz0U/fKvWE0d4jsNQEfg==", "dev": true, "license": "MIT", "dependencies": { - "@vitest/pretty-format": "3.2.4", + "@vitest/pretty-format": "3.2.6", "loupe": "^3.1.4", "tinyrainbow": "^2.0.0" }, @@ -5027,9 +5027,9 @@ } }, "node_modules/brace-expansion": { - "version": "5.0.5", - "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-5.0.5.tgz", - "integrity": "sha512-VZznLgtwhn+Mact9tfiwx64fA9erHH/MCXEUfB/0bX/6Fz6ny5EGTXYltMocqg4xFAQZtnO3DHWWXi8RiuN7cQ==", + "version": "5.0.6", + "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-5.0.6.tgz", + "integrity": "sha512-kLpxurY4Z4r9sgMsyG0Z9uzsBlgiU/EFKhj/h91/8yHu0edo7XuixOIH3VcJ8kkxs6/jPzoI6U9Vj3WqbMQ94g==", "dev": true, "license": "MIT", "dependencies": { @@ -6840,9 +6840,9 @@ } }, "node_modules/fflate": { - "version": "0.8.2", - "resolved": "https://registry.npmjs.org/fflate/-/fflate-0.8.2.tgz", - "integrity": "sha512-cPJU47OaAoCbg0pBvzsgpTPhmhqI5eJjh/JIu8tPj5q+T7iLvW/JAYUqmE7KOB4R1ZyEhzBaIQpQpardBF5z8A==", + "version": "0.8.3", + "resolved": "https://registry.npmjs.org/fflate/-/fflate-0.8.3.tgz", + "integrity": "sha512-tbZNuJrLwGUp3zshBtdy4W+ORxZuIh8a5ilyIEQDC5rY1f3U20JMry0Ll3WBzU58EZKsEuJFXhb5gwv8CsPvgA==", "dev": true, "license": "MIT" }, @@ -13499,20 +13499,20 @@ } }, "node_modules/vitest": { - "version": "3.2.4", - "resolved": "https://registry.npmjs.org/vitest/-/vitest-3.2.4.tgz", - "integrity": "sha512-LUCP5ev3GURDysTWiP47wRRUpLKMOfPh+yKTx3kVIEiu5KOMeqzpnYNsKyOoVrULivR8tLcks4+lga33Whn90A==", + "version": "3.2.6", + "resolved": "https://registry.npmjs.org/vitest/-/vitest-3.2.6.tgz", + "integrity": "sha512-xejya+bT/j/+R/AGa1XOfRxLmNUlLtlwjRsFUILF+xHfzElmGcmFydy2gqqIrd62ptIEfwVMofd19uNWD9L7Nw==", "dev": true, "license": "MIT", "dependencies": { "@types/chai": "^5.2.2", - "@vitest/expect": "3.2.4", - "@vitest/mocker": "3.2.4", - "@vitest/pretty-format": "^3.2.4", - "@vitest/runner": "3.2.4", - "@vitest/snapshot": "3.2.4", - "@vitest/spy": "3.2.4", - "@vitest/utils": "3.2.4", + "@vitest/expect": "3.2.6", + "@vitest/mocker": "3.2.6", + "@vitest/pretty-format": "^3.2.6", + "@vitest/runner": "3.2.6", + "@vitest/snapshot": "3.2.6", + "@vitest/spy": "3.2.6", + "@vitest/utils": "3.2.6", "chai": "^5.2.0", "debug": "^4.4.1", "expect-type": "^1.2.1", @@ -13542,8 +13542,8 @@ "@edge-runtime/vm": "*", "@types/debug": "^4.1.12", "@types/node": "^18.0.0 || ^20.0.0 || >=22.0.0", - "@vitest/browser": "3.2.4", - "@vitest/ui": "3.2.4", + "@vitest/browser": "3.2.6", + "@vitest/ui": "3.2.6", "happy-dom": "*", "jsdom": "*" }, diff --git a/ui/litellm-dashboard/package.json b/ui/litellm-dashboard/package.json index ca753e59dc5c..da6869c538d7 100644 --- a/ui/litellm-dashboard/package.json +++ b/ui/litellm-dashboard/package.json @@ -65,8 +65,8 @@ "@types/react-dom": "18.3.7", "@types/react-syntax-highlighter": "15.5.13", "@types/uuid": "10.0.0", - "@vitest/coverage-v8": "3.2.4", - "@vitest/ui": "3.2.4", + "@vitest/coverage-v8": "3.2.6", + "@vitest/ui": "3.2.6", "autoprefixer": "10.4.24", "dotenv": "17.2.3", "eslint": "9.39.2", @@ -82,7 +82,7 @@ "typescript": "5.9.3", "typescript-eslint": "8.60.1", "vite": "7.3.2", - "vitest": "3.2.4" + "vitest": "3.2.6" }, "overrides": { "prismjs": "1.30.0", @@ -92,6 +92,7 @@ "lodash": "4.18.1", "ws": "8.19.0", "braces": "3.0.3", + "brace-expansion": "5.0.6", "axios": "1.13.6", "postcss": "8.5.13" }, diff --git a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.test.tsx b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.test.tsx index 042f02251b4c..e70548d6a96f 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.test.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.test.tsx @@ -17,15 +17,23 @@ vi.mock("@/utils/mcpTokenStore", () => ({ })); // Mutable holder so individual tests can simulate "Authorize & Fetch" having -// produced a token before submit. -const oauthHook = vi.hoisted(() => ({ tokenResponse: null as Record | null })); +// produced a token before submit, and inspect the reset wiring. +const oauthHook = vi.hoisted(() => ({ + tokenResponse: null as Record | null, + reset: vi.fn(), + onTokenReceived: null as ((token: Record | null) => void) | null, +})); vi.mock("@/hooks/useMcpOAuthFlow", () => ({ - useMcpOAuthFlow: () => ({ - startOAuthFlow: vi.fn(), - status: "idle", - error: null, - tokenResponse: oauthHook.tokenResponse, - }), + useMcpOAuthFlow: (opts: { onTokenReceived: (token: Record | null) => void }) => { + oauthHook.onTokenReceived = opts.onTokenReceived; + return { + startOAuthFlow: vi.fn(), + status: "idle", + error: null, + tokenResponse: oauthHook.tokenResponse, + reset: oauthHook.reset, + }; + }, })); vi.mock("./mcp_server_cost_config", () => ({ @@ -59,7 +67,9 @@ vi.mock("./mcp_tool_configuration", () => ({ })); vi.mock("./mcp_connection_status", () => ({ - default: () =>
, + default: ({ tools }: { tools?: any[] }) => ( +
+ ), })); vi.mock("./StdioConfiguration", () => ({ @@ -121,6 +131,7 @@ describe("CreateMCPServer", () => { beforeEach(() => { vi.clearAllMocks(); oauthHook.tokenResponse = null; + oauthHook.onTokenReceived = null; }); it("should render the modal with title when visible", () => { @@ -614,6 +625,100 @@ describe("CreateMCPServer", () => { expect(defaultProps.setModalVisible).toHaveBeenCalledWith(false); }); + + it("does not leak a previous server's OAuth token into the next add-server session", async () => { + const usedToken = (token: string) => + vi.mocked(networking.testMCPToolsListRequest).mock.calls.some((call) => call[2] === token); + + const { rerender } = render(); + + await selectAntOption("Transport Type", "Streamable HTTP"); + await waitFor(() => { + expect(screen.getByPlaceholderText("https://your-mcp-server.com")).toBeInTheDocument(); + }); + await selectAntOption("Authentication", "OAuth"); + await waitFor(() => { + expect(screen.getByText("OAuth Flow Type")).toBeInTheDocument(); + }); + + const urlInput = screen.getByPlaceholderText("https://your-mcp-server.com"); + await act(async () => { + fireEvent.change(urlInput, { target: { value: "https://server-a.example.com/mcp" } }); + }); + + // Simulate "Authorize & Fetch Token" completing for server A. + await act(async () => { + oauthHook.onTokenReceived?.({ access_token: "stale-token-A", expires_in: 3600 }); + }); + + // Precondition: the freshly fetched token drives the tool preview for server A. + await waitFor(() => { + expect(usedToken("stale-token-A")).toBe(true); + }); + + // Parent hides the modal (Cancel / successful create both flip this prop). + rerender(); + + // The OAuth flow state (source of the "Token fetched" badge) is reset on close. + expect(oauthHook.reset).toHaveBeenCalled(); + + vi.mocked(networking.testMCPToolsListRequest).mockClear(); + + // Reopen for a brand-new server and enter a different URL without re-authorizing. + rerender(); + const reopenedUrlInput = screen.getByPlaceholderText("https://your-mcp-server.com"); + await act(async () => { + fireEvent.change(reopenedUrlInput, { target: { value: "https://server-b.example.com/mcp" } }); + }); + + // The previous server's token must never be replayed for the new session. + expect(usedToken("stale-token-A")).toBe(false); + }); + + it("clears the tool list and form fields when a parent dismisses the modal", async () => { + vi.mocked(networking.testMCPToolsListRequest).mockResolvedValue({ + tools: [{ name: "tool_a" }], + error: null, + }); + const toolCount = () => screen.getByTestId("mcp-connection-status").getAttribute("data-tool-count"); + + const { rerender } = render(); + + await selectAntOption("Transport Type", "Streamable HTTP"); + await waitFor(() => { + expect(screen.getByPlaceholderText("https://your-mcp-server.com")).toBeInTheDocument(); + }); + await selectAntOption("Authentication", "OAuth"); + await waitFor(() => { + expect(screen.getByText("OAuth Flow Type")).toBeInTheDocument(); + }); + + const urlInput = screen.getByPlaceholderText("https://your-mcp-server.com"); + await act(async () => { + fireEvent.change(urlInput, { target: { value: "https://server-a.example.com/mcp" } }); + }); + await act(async () => { + oauthHook.onTokenReceived?.({ access_token: "stale-token-A", expires_in: 3600 }); + }); + + // Precondition: a tool list is shown for server A. + await waitFor(() => { + expect(toolCount()).toBe("1"); + }); + + // Parent dismisses the modal without routing through Cancel or create. + rerender(); + + // Stale tools are cleared even though neither handler ran. + await waitFor(() => { + expect(toolCount()).toBe("0"); + }); + + // Reopening starts clean: the URL the prior server left in the Ant form store is gone. + rerender(); + const reopenedUrlInput = screen.getByPlaceholderText("https://your-mcp-server.com") as HTMLInputElement; + expect(reopenedUrlInput.value).toBe(""); + }); }); describe("when stdio transport is selected", () => { diff --git a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx index 6a8dc353f212..ddcc9f65d38f 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx @@ -134,6 +134,7 @@ const CreateMCPServer: React.FC = ({ status: oauthStatus, error: oauthError, tokenResponse: oauthTokenResponse, + reset: resetOAuthFlow, } = useMcpOAuthFlow({ accessToken, getCredentials: () => form.getFieldValue("credentials"), @@ -188,6 +189,7 @@ const CreateMCPServer: React.FC = ({ } }, onBeforeRedirect: persistCreateUiState, + flowSource: "create", }); React.useEffect(() => { @@ -553,12 +555,19 @@ const CreateMCPServer: React.FC = ({ } }, [formValues.server_name]); - // Clear formValues when modal closes to reset child components + // Clear form, tools, and OAuth state when the modal closes so a previous server's + // authorization, credentials, or tool list never bleed into the next "Add New MCP + // Server" session, including when a parent dismisses the modal without routing + // through handleCancel or handleCreate. React.useEffect(() => { if (!isModalVisible) { + form.resetFields(); setFormValues({}); + setOauthAccessToken(null); + clearTools(); + resetOAuthFlow(); } - }, [isModalVisible]); + }, [isModalVisible, form, clearTools, resetOAuthFlow]); const isAdmin = isAdminRole(userRole); @@ -1088,7 +1097,6 @@ const CreateMCPServer: React.FC = ({
({ updateMCPServer: vi.fn(), listMCPTools: vi.fn().mockResolvedValue({ tools: [], error: null }), + storeMCPOAuthUserCredential: vi.fn().mockResolvedValue({}), })); vi.mock("../molecules/notifications_manager", () => ({ @@ -17,12 +18,13 @@ vi.mock("../molecules/notifications_manager", () => ({ }, })); +const mockOauth: { tokenResponse: any } = { tokenResponse: null }; vi.mock("@/hooks/useMcpOAuthFlow", () => ({ useMcpOAuthFlow: () => ({ startOAuthFlow: vi.fn(), status: "idle", error: null, - tokenResponse: null, + tokenResponse: mockOauth.tokenResponse, }), })); @@ -37,12 +39,19 @@ vi.mock("./MCPPermissionManagement", () => ({ vi.mock("./mcp_tool_configuration", () => ({ default: ({ existingAllowedTools, + externalTools, + externalError, onAllowedToolsChange, onToolAllowlistInteraction, onToolNameToDisplayNameChange, onToolNameToDescriptionChange, }: any) => ( -
+