diff --git a/docs/my-website/docs/providers/perplexity.md b/docs/my-website/docs/providers/perplexity.md index 68adf9939c..e3991c63bf 100644 --- a/docs/my-website/docs/providers/perplexity.md +++ b/docs/my-website/docs/providers/perplexity.md @@ -120,7 +120,7 @@ All models listed here https://docs.perplexity.ai/docs/model-cards are supported -## Agentic Research API (Responses API) +## Agent API (Responses API) Requires v1.72.6+ @@ -196,7 +196,7 @@ import os os.environ['PERPLEXITY_API_KEY'] = "" response = responses( - model="perplexity/openai/gpt-4o", + model="perplexity/openai/gpt-5.2", input="Explain quantum computing in simple terms", custom_llm_provider="perplexity", max_output_tokens=500, @@ -215,7 +215,7 @@ import os os.environ['PERPLEXITY_API_KEY'] = "" response = responses( - model="perplexity/anthropic/claude-3-5-sonnet-20241022", + model="perplexity/anthropic/claude-sonnet-4-5", input="Write a short story about a robot learning to paint", custom_llm_provider="perplexity", max_output_tokens=500, @@ -234,7 +234,7 @@ import os os.environ['PERPLEXITY_API_KEY'] = "" response = responses( - model="perplexity/google/gemini-2.0-flash-exp", + model="perplexity/google/gemini-2.5-flash", input="Explain the concept of neural networks", custom_llm_provider="perplexity", max_output_tokens=500, @@ -253,7 +253,7 @@ import os os.environ['PERPLEXITY_API_KEY'] = "" response = responses( - model="perplexity/xai/grok-2-1212", + model="perplexity/xai/grok-4-1-fast-non-reasoning", input="What makes a good AI assistant?", custom_llm_provider="perplexity", max_output_tokens=500, @@ -276,7 +276,7 @@ import os os.environ['PERPLEXITY_API_KEY'] = "" response = responses( - model="perplexity/openai/gpt-4o", + model="perplexity/openai/gpt-5.2", input="What's the weather in San Francisco today?", custom_llm_provider="perplexity", tools=[{"type": "web_search"}], @@ -286,6 +286,78 @@ response = responses( print(response.output) ``` +### Function Calling + +The Agent API supports custom function tools. Pass function tools through unchanged: + +```python +from litellm import responses +import os + +os.environ['PERPLEXITY_API_KEY'] = "" + +response = responses( + model="perplexity/openai/gpt-5.2", + input="What's the weather in San Francisco?", + custom_llm_provider="perplexity", + tools=[ + {"type": "web_search"}, + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather for a location", + "parameters": { + "type": "object", + "properties": { + "location": {"type": "string"}, + "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]}, + }, + }, + }, + }, + ], + instructions="Use tools when appropriate.", +) + +print(response.output) +``` + +### Structured Outputs + +Request JSON schema structured outputs via the `text` parameter: + +```python +from litellm import responses +import os + +os.environ['PERPLEXITY_API_KEY'] = "" + +response = responses( + model="perplexity/preset/pro-search", + input="Extract key facts about the Eiffel Tower", + custom_llm_provider="perplexity", + text={ + "format": { + "type": "json_schema", + "name": "facts", + "schema": { + "type": "object", + "properties": { + "name": {"type": "string"}, + "height_meters": {"type": "number"}, + "year_built": {"type": "integer"}, + }, + "required": ["name", "height_meters", "year_built"], + }, + "strict": True, + } + }, +) + +print(response.output) +``` + ### Reasoning Effort (Responses API) @@ -319,7 +391,7 @@ import os os.environ['PERPLEXITY_API_KEY'] = "" response = responses( - model="perplexity/anthropic/claude-3-5-sonnet-20241022", + model="perplexity/anthropic/claude-sonnet-4-5", input=[ {"type": "message", "role": "system", "content": "You are a helpful assistant."}, {"type": "message", "role": "user", "content": "What are the latest AI developments?"}, @@ -343,7 +415,7 @@ import os os.environ['PERPLEXITY_API_KEY'] = "" response = responses( - model="perplexity/openai/gpt-4o", + model="perplexity/openai/gpt-5.2", input="Tell me a story about space exploration", custom_llm_provider="perplexity", stream=True, @@ -360,23 +432,28 @@ for chunk in response: | Provider | Model Name | Function Call | |----------|------------|---------------| -| OpenAI | gpt-4o | `responses(model="perplexity/openai/gpt-4o", ...)` | -| OpenAI | gpt-4o-mini | `responses(model="perplexity/openai/gpt-4o-mini", ...)` | | OpenAI | gpt-5.2 | `responses(model="perplexity/openai/gpt-5.2", ...)` | -| Anthropic | claude-3-5-sonnet-20241022 | `responses(model="perplexity/anthropic/claude-3-5-sonnet-20241022", ...)` | -| Anthropic | claude-3-5-haiku-20241022 | `responses(model="perplexity/anthropic/claude-3-5-haiku-20241022", ...)` | -| Google | gemini-2.0-flash-exp | `responses(model="perplexity/google/gemini-2.0-flash-exp", ...)` | -| Google | gemini-2.0-flash-thinking-exp | `responses(model="perplexity/google/gemini-2.0-flash-thinking-exp", ...)` | -| xAI | grok-2-1212 | `responses(model="perplexity/xai/grok-2-1212", ...)` | -| xAI | grok-2-vision-1212 | `responses(model="perplexity/xai/grok-2-vision-1212", ...)` | +| OpenAI | gpt-5.1 | `responses(model="perplexity/openai/gpt-5.1", ...)` | +| OpenAI | gpt-5-mini | `responses(model="perplexity/openai/gpt-5-mini", ...)` | +| Anthropic | claude-opus-4-6 | `responses(model="perplexity/anthropic/claude-opus-4-6", ...)` | +| Anthropic | claude-opus-4-5 | `responses(model="perplexity/anthropic/claude-opus-4-5", ...)` | +| Anthropic | claude-sonnet-4-5 | `responses(model="perplexity/anthropic/claude-sonnet-4-5", ...)` | +| Anthropic | claude-haiku-4-5 | `responses(model="perplexity/anthropic/claude-haiku-4-5", ...)` | +| Google | gemini-3-pro-preview | `responses(model="perplexity/google/gemini-3-pro-preview", ...)` | +| Google | gemini-3-flash-preview | `responses(model="perplexity/google/gemini-3-flash-preview", ...)` | +| Google | gemini-2.5-pro | `responses(model="perplexity/google/gemini-2.5-pro", ...)` | +| Google | gemini-2.5-flash | `responses(model="perplexity/google/gemini-2.5-flash", ...)` | +| xAI | grok-4-1-fast-non-reasoning | `responses(model="perplexity/xai/grok-4-1-fast-non-reasoning", ...)` | +| Perplexity | sonar | `responses(model="perplexity/perplexity/sonar", ...)` | ### Available Presets -| Preset Name | Function Call | -|----------------|--------------------------------------------------------| -| fast-search | `responses(model="perplexity/preset/fast-search", ...)`| -| pro-search | `responses(model="perplexity/preset/pro-search", ...)` | -| deep-research | `responses(model="perplexity/preset/deep-research", ...)`| +| Preset Name | Function Call | +|-------------|---------------| +| fast-search | `responses(model="perplexity/preset/fast-search", ...)` | +| pro-search | `responses(model="perplexity/preset/pro-search", ...)` | +| deep-research | `responses(model="perplexity/preset/deep-research", ...)` | +| advanced-deep-research | `responses(model="perplexity/preset/advanced-deep-research", ...)` | ### Complete Example @@ -388,7 +465,7 @@ os.environ['PERPLEXITY_API_KEY'] = "" # Comprehensive example with multiple features response = responses( - model="perplexity/openai/gpt-4o", + model="perplexity/openai/gpt-5.2", input="Research the latest developments in quantum computing and provide sources", custom_llm_provider="perplexity", tools=[ diff --git a/docs/my-website/docs/realtime.md b/docs/my-website/docs/realtime.md index b191c82c67..c853a860de 100644 --- a/docs/my-website/docs/realtime.md +++ b/docs/my-website/docs/realtime.md @@ -85,7 +85,7 @@ const url = "ws://0.0.0.0:4000/v1/realtime?model=openai-gpt-4o-realtime-audio"; // const url = "wss://my-endpoint-sweden-berri992.openai.azure.com/openai/realtime?api-version=2024-10-01-preview&deployment=gpt-4o-realtime-preview"; const ws = new WebSocket(url, { headers: { - "api-key": `f28ab7b695af4154bc53498e5bdccb07`, + "api-key": `sk-1234`, "OpenAI-Beta": "realtime=v1", }, }); diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index ad6eb6b4f3..cc0f818b0a 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -1242,6 +1242,16 @@ def completion_cost( # noqa: PLR0915 ) elif call_type in _VIDEO_CALL_TYPES: ### VIDEO GENERATION COST CALCULATION ### + # Extract custom model_info for deployment-specific pricing + _video_model_info: Optional[ModelInfo] = None + if custom_pricing and litellm_logging_obj is not None: + _litellm_params = getattr( + litellm_logging_obj, "litellm_params", None + ) + if _litellm_params is not None: + _metadata = _litellm_params.get("metadata", {}) or {} + _video_model_info = _metadata.get("model_info", None) + usage_obj = getattr(completion_response, "usage", None) if completion_response is not None and usage_obj: # Handle both dict and Pydantic Usage object @@ -1262,12 +1272,14 @@ def completion_cost( # noqa: PLR0915 model=model, duration_seconds=duration_seconds, custom_llm_provider=custom_llm_provider, + model_info=_video_model_info, ) # Fallback to default video cost calculation if no duration available return default_video_cost_calculator( model=model, duration_seconds=0.0, # Default to 0 if no duration available custom_llm_provider=custom_llm_provider, + model_info=_video_model_info, ) elif call_type in _SPEECH_CALL_TYPES: prompt_characters = litellm.utils._count_characters(text=prompt) @@ -1892,6 +1904,7 @@ def default_video_cost_calculator( model: str, duration_seconds: float, custom_llm_provider: Optional[str] = None, + model_info: Optional[ModelInfo] = None, ) -> float: """ Default video cost calculator for video generation @@ -1900,6 +1913,9 @@ def default_video_cost_calculator( model (str): Model name duration_seconds (float): Duration of the generated video in seconds custom_llm_provider (Optional[str]): Custom LLM provider + model_info (Optional[ModelInfo]): Deployment-level model info containing + custom video pricing. When provided, used before falling back to + the global litellm.model_cost lookup. Returns: float: Cost in USD for the video generation @@ -1907,42 +1923,47 @@ def default_video_cost_calculator( Raises: Exception: If model pricing not found in cost map """ - # Build model names for cost lookup - base_model_name = model - model_name_without_custom_llm_provider: Optional[str] = None - if custom_llm_provider and model.startswith(f"{custom_llm_provider}/"): - model_name_without_custom_llm_provider = model.replace( - f"{custom_llm_provider}/", "" - ) - base_model_name = ( - f"{custom_llm_provider}/{model_name_without_custom_llm_provider}" - ) - - verbose_logger.debug(f"Looking up cost for video model: {base_model_name}") - - model_without_provider = model.split("/")[-1] - - # Try model with provider first, fall back to base model name + # Use custom model_info pricing if provided (deployment-specific pricing) cost_info: Optional[dict] = None - models_to_check: List[Optional[str]] = [ - base_model_name, - model, - model_without_provider, - model_name_without_custom_llm_provider, - ] - for _model in models_to_check: - if _model is not None and _model in litellm.model_cost: - cost_info = litellm.model_cost[_model] - break + if model_info is not None: + cost_info = dict(model_info) + else: + # Build model names for cost lookup + base_model_name = model + model_name_without_custom_llm_provider: Optional[str] = None + if custom_llm_provider and model.startswith(f"{custom_llm_provider}/"): + model_name_without_custom_llm_provider = model.replace( + f"{custom_llm_provider}/", "" + ) + base_model_name = ( + f"{custom_llm_provider}/{model_name_without_custom_llm_provider}" + ) + + verbose_logger.debug(f"Looking up cost for video model: {base_model_name}") + + model_without_provider = model.split("/")[-1] + + # Try model with provider first, fall back to base model name + models_to_check: List[Optional[str]] = [ + base_model_name, + model, + model_without_provider, + model_name_without_custom_llm_provider, + ] + for _model in models_to_check: + if _model is not None and _model in litellm.model_cost: + cost_info = litellm.model_cost[_model] + break + + # If still not found, try with custom_llm_provider prefix + if cost_info is None and custom_llm_provider: + prefixed_model = f"{custom_llm_provider}/{model}" + if prefixed_model in litellm.model_cost: + cost_info = litellm.model_cost[prefixed_model] - # If still not found, try with custom_llm_provider prefix - if cost_info is None and custom_llm_provider: - prefixed_model = f"{custom_llm_provider}/{model}" - if prefixed_model in litellm.model_cost: - cost_info = litellm.model_cost[prefixed_model] if cost_info is None: raise Exception( - f"Model not found in cost map. Tried checking {models_to_check}" + f"Model not found in cost map for model={model}" ) # Check for video-specific cost per second first diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 3648dc691e..0601e7e845 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -4709,6 +4709,7 @@ class StandardLoggingPayloadSetup: custom_pricing: Optional[bool], custom_llm_provider: Optional[str], init_response_obj: Union[Any, BaseModel, dict], + api_base: Optional[str] = None, ) -> StandardLoggingModelInformation: model_cost_name = _select_model_name_for_cost_calc( model=None, @@ -4723,7 +4724,9 @@ class StandardLoggingPayloadSetup: else: try: _model_cost_information = litellm.get_model_info( - model=model_cost_name, custom_llm_provider=custom_llm_provider + model=model_cost_name, + custom_llm_provider=custom_llm_provider, + api_base=api_base, ) model_cost_information = StandardLoggingModelInformation( model_map_key=model_cost_name, @@ -5236,6 +5239,7 @@ def get_standard_logging_object_payload( custom_pricing=custom_pricing, custom_llm_provider=kwargs.get("custom_llm_provider"), init_response_obj=init_response_obj, + api_base=litellm_params.get("api_base"), ) response_cost: float = kwargs.get("response_cost", 0) or 0.0 diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 7b485501f6..ba415af9a5 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -1848,9 +1848,10 @@ def convert_to_anthropic_tool_invoke( break else: # Regular tool_use + sanitized_tool_id = _sanitize_anthropic_tool_use_id(tool_id) _anthropic_tool_use_param = AnthropicMessagesToolUseParam( type="tool_use", - id=tool_id, + id=sanitized_tool_id, name=tool_name, input=tool_input, ) diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 3789f546d7..3dfef07d42 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -28,6 +28,7 @@ from litellm.constants import ( AIOHTTP_CONNECTOR_LIMIT, AIOHTTP_CONNECTOR_LIMIT_PER_HOST, AIOHTTP_KEEPALIVE_TIMEOUT, + AIOHTTP_NEEDS_CLEANUP_CLOSED, AIOHTTP_TTL_DNS_CACHE, DEFAULT_SSL_CIPHERS, ) @@ -876,9 +877,10 @@ class AsyncHTTPHandler: transport_connector_kwargs = { "keepalive_timeout": AIOHTTP_KEEPALIVE_TIMEOUT, "ttl_dns_cache": AIOHTTP_TTL_DNS_CACHE, - "enable_cleanup_closed": True, **connector_kwargs, } + if AIOHTTP_NEEDS_CLEANUP_CLOSED: + transport_connector_kwargs["enable_cleanup_closed"] = True if AIOHTTP_CONNECTOR_LIMIT > 0: transport_connector_kwargs["limit"] = AIOHTTP_CONNECTOR_LIMIT if AIOHTTP_CONNECTOR_LIMIT_PER_HOST > 0: diff --git a/litellm/llms/ollama/completion/transformation.py b/litellm/llms/ollama/completion/transformation.py index c4d08c83a2..ed14b6a331 100644 --- a/litellm/llms/ollama/completion/transformation.py +++ b/litellm/llms/ollama/completion/transformation.py @@ -6,7 +6,7 @@ from typing import TYPE_CHECKING, Any, AsyncIterator, Iterator, List, Optional, from httpx._models import Headers, Response import litellm -from litellm._logging import verbose_proxy_logger +from litellm._logging import verbose_logger, verbose_proxy_logger from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_str_from_messages, ) @@ -223,7 +223,9 @@ class OllamaConfig(BaseConfig): or get_secret_str("OLLAMA_API_KEY") ) - def get_model_info(self, model: str) -> ModelInfoBase: + def get_model_info( + self, model: str, api_base: Optional[str] = None + ) -> ModelInfoBase: """ curl http://localhost:11434/api/show -d '{ "name": "mistral" @@ -231,7 +233,11 @@ class OllamaConfig(BaseConfig): """ if model.startswith("ollama/") or model.startswith("ollama_chat/"): model = model.split("/", 1)[1] - api_base = get_secret_str("OLLAMA_API_BASE") or "http://localhost:11434" + api_base = ( + api_base + or get_secret_str("OLLAMA_API_BASE") + or "http://localhost:11434" + ) api_key = self.get_api_key() headers = {"Authorization": f"Bearer {api_key}"} if api_key else {} @@ -242,8 +248,21 @@ class OllamaConfig(BaseConfig): headers=headers, ) except Exception as e: - raise Exception( - f"OllamaError: Error getting model info for {model}. Set Ollama API Base via `OLLAMA_API_BASE` environment variable. Error: {e}" + verbose_logger.debug( + "OllamaError: Could not get model info for %s from %s. Error: %s", + model, + api_base, + e, + ) + return ModelInfoBase( + key=model, + litellm_provider="ollama", + mode="chat", + input_cost_per_token=0.0, + output_cost_per_token=0.0, + max_tokens=None, + max_input_tokens=None, + max_output_tokens=None, ) model_info = response.json() diff --git a/litellm/llms/openai/cost_calculation.py b/litellm/llms/openai/cost_calculation.py index e5349db3af..ac1e4a6b08 100644 --- a/litellm/llms/openai/cost_calculation.py +++ b/litellm/llms/openai/cost_calculation.py @@ -7,7 +7,7 @@ from typing import Literal, Optional, Tuple from litellm._logging import verbose_logger from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token -from litellm.types.utils import CallTypes, Usage +from litellm.types.utils import CallTypes, ModelInfo, Usage from litellm.utils import get_model_info @@ -129,7 +129,10 @@ def cost_per_second( def video_generation_cost( - model: str, duration_seconds: float, custom_llm_provider: Optional[str] = None + model: str, + duration_seconds: float, + custom_llm_provider: Optional[str] = None, + model_info: Optional[ModelInfo] = None, ) -> float: """ Calculates the cost for video generation based on duration in seconds. @@ -138,14 +141,18 @@ def video_generation_cost( - model: str, the model name without provider prefix - duration_seconds: float, the duration of the generated video in seconds - custom_llm_provider: str, the custom llm provider + - model_info: Optional[dict], deployment-level model info containing + custom video pricing. When provided, skips the global + get_model_info() lookup so that deployment-specific pricing is used. Returns: float - total_cost_in_usd """ ## GET MODEL INFO - model_info = get_model_info( - model=model, custom_llm_provider=custom_llm_provider or "openai" - ) + if model_info is None: + model_info = get_model_info( + model=model, custom_llm_provider=custom_llm_provider or "openai" + ) # Check for video-specific cost per second video_cost_per_second = model_info.get("output_cost_per_video_per_second") diff --git a/litellm/llms/perplexity/responses/__init__.py b/litellm/llms/perplexity/responses/__init__.py index 9bdf810e83..3285a47211 100644 --- a/litellm/llms/perplexity/responses/__init__.py +++ b/litellm/llms/perplexity/responses/__init__.py @@ -1,5 +1,5 @@ """ -Perplexity Agentic Research API (Responses API) module +Perplexity Agent API (Responses API) module """ from .transformation import PerplexityResponsesConfig diff --git a/litellm/llms/perplexity/responses/transformation.py b/litellm/llms/perplexity/responses/transformation.py index 178e76ea97..6d2ed51600 100644 --- a/litellm/llms/perplexity/responses/transformation.py +++ b/litellm/llms/perplexity/responses/transformation.py @@ -1,5 +1,5 @@ """ -Transformation logic for Perplexity Agentic Research API (Responses API) +Transformation logic for Perplexity Agent API (Responses API) This module handles the translation between OpenAI's Responses API format and Perplexity's Responses API format, which supports: @@ -32,10 +32,10 @@ from litellm.types.utils import LlmProviders class PerplexityResponsesConfig(OpenAIResponsesAPIConfig): """ - Configuration for Perplexity Agentic Research API (Responses API) + Configuration for Perplexity Agent API (Responses API) - - Reference: https://docs.perplexity.ai/agentic-research/quickstart + + Reference: https://docs.perplexity.ai/docs/agent-api/overview """ @property @@ -45,8 +45,9 @@ class PerplexityResponsesConfig(OpenAIResponsesAPIConfig): def get_supported_openai_params(self, model: str) -> list: """ Perplexity Responses API supports a different set of parameters - + Ref: https://docs.perplexity.ai/api-reference/responses-post + Params aligned with response-echo fields and Open Responses spec. """ return [ "max_output_tokens", @@ -58,6 +59,23 @@ class PerplexityResponsesConfig(OpenAIResponsesAPIConfig): "preset", "instructions", "models", # Model fallback support + "tool_choice", + "parallel_tool_calls", + "max_tool_calls", + "text", + "previous_response_id", + "store", + "background", + "truncation", + "metadata", + "safety_identifier", + "user", + "stream_options", + "top_logprobs", + "prompt_cache_key", + "frequency_penalty", + "presence_penalty", + "service_tier", ] def validate_environment( @@ -65,16 +83,15 @@ class PerplexityResponsesConfig(OpenAIResponsesAPIConfig): ) -> dict: """Validate environment and set up headers""" # Get API key from environment - api_key = ( - get_secret_str("PERPLEXITYAI_API_KEY") - or get_secret_str("PERPLEXITY_API_KEY") + api_key = get_secret_str("PERPLEXITYAI_API_KEY") or get_secret_str( + "PERPLEXITY_API_KEY" ) - + if api_key: headers["Authorization"] = f"Bearer {api_key}" - + headers["Content-Type"] = "application/json" - + return headers def get_complete_url( @@ -84,15 +101,17 @@ class PerplexityResponsesConfig(OpenAIResponsesAPIConfig): ) -> str: """Get the complete URL for the Perplexity Responses API""" if api_base is None: - api_base = get_secret_str("PERPLEXITY_API_BASE") or "https://api.perplexity.ai" - + api_base = ( + get_secret_str("PERPLEXITY_API_BASE") or "https://api.perplexity.ai" + ) + # Ensure api_base doesn't end with a slash api_base = api_base.rstrip("/") - + # Add the responses endpoint return f"{api_base}/v1/responses" - def map_openai_params( + def map_openai_params( # noqa: PLR0915 self, response_api_optional_params: ResponsesAPIOptionalRequestParams, model: str, @@ -100,78 +119,136 @@ class PerplexityResponsesConfig(OpenAIResponsesAPIConfig): ) -> Dict: """ Map OpenAI Responses API parameters to Perplexity format - + Key differences: - Supports 'preset' parameter for predefined configurations - Supports 'instructions' parameter for system-level guidance - Tools are specified differently (web_search, fetch_url) """ mapped_params: Dict[str, Any] = {} - + # Map standard parameters if response_api_optional_params.get("max_output_tokens"): - mapped_params["max_output_tokens"] = response_api_optional_params["max_output_tokens"] - + mapped_params["max_output_tokens"] = response_api_optional_params[ + "max_output_tokens" + ] + if response_api_optional_params.get("temperature"): mapped_params["temperature"] = response_api_optional_params["temperature"] - + if response_api_optional_params.get("top_p"): mapped_params["top_p"] = response_api_optional_params["top_p"] - + if response_api_optional_params.get("stream"): mapped_params["stream"] = response_api_optional_params["stream"] - + if response_api_optional_params.get("stream_options"): - mapped_params["stream_options"] = response_api_optional_params["stream_options"] - + mapped_params["stream_options"] = response_api_optional_params[ + "stream_options" + ] + # Map Perplexity-specific parameters (using .get() with Any dict access) preset = response_api_optional_params.get("preset") # type: ignore if preset: mapped_params["preset"] = preset - + instructions = response_api_optional_params.get("instructions") # type: ignore if instructions: mapped_params["instructions"] = instructions - + if response_api_optional_params.get("reasoning"): mapped_params["reasoning"] = response_api_optional_params["reasoning"] - + tools = response_api_optional_params.get("tools") if tools: # Convert tools to list of dicts for transformation - tools_list = [dict(tool) if hasattr(tool, '__dict__') else tool for tool in tools] # type: ignore + tools_list = [dict(tool) if hasattr(tool, "__dict__") else tool for tool in tools] # type: ignore mapped_params["tools"] = self._transform_tools(tools_list) # type: ignore - + + # Tool control + if response_api_optional_params.get("tool_choice"): + mapped_params["tool_choice"] = response_api_optional_params["tool_choice"] + if response_api_optional_params.get("parallel_tool_calls") is not None: + mapped_params["parallel_tool_calls"] = response_api_optional_params[ + "parallel_tool_calls" + ] + if response_api_optional_params.get("max_tool_calls"): + mapped_params["max_tool_calls"] = response_api_optional_params[ + "max_tool_calls" + ] + + # Structured outputs + text_param = response_api_optional_params.get("text") + if text_param: + mapped_params["text"] = text_param + + # Conversation continuity + if response_api_optional_params.get("previous_response_id"): + mapped_params["previous_response_id"] = response_api_optional_params[ + "previous_response_id" + ] + + # Storage and lifecycle + if response_api_optional_params.get("store") is not None: + mapped_params["store"] = response_api_optional_params["store"] + if response_api_optional_params.get("background") is not None: + mapped_params["background"] = response_api_optional_params["background"] + if response_api_optional_params.get("truncation"): + mapped_params["truncation"] = response_api_optional_params["truncation"] + + # Metadata + if response_api_optional_params.get("metadata"): + mapped_params["metadata"] = response_api_optional_params["metadata"] + if response_api_optional_params.get("safety_identifier"): + mapped_params["safety_identifier"] = response_api_optional_params[ + "safety_identifier" + ] + if response_api_optional_params.get("user"): + mapped_params["user"] = response_api_optional_params["user"] + + # Additional + if response_api_optional_params.get("top_logprobs") is not None: + mapped_params["top_logprobs"] = response_api_optional_params["top_logprobs"] + if response_api_optional_params.get("prompt_cache_key"): + mapped_params["prompt_cache_key"] = response_api_optional_params[ + "prompt_cache_key" + ] + if response_api_optional_params.get("frequency_penalty") is not None: + mapped_params["frequency_penalty"] = response_api_optional_params[ + "frequency_penalty" # type: ignore[typeddict-item] + ] + if response_api_optional_params.get("presence_penalty") is not None: + mapped_params["presence_penalty"] = response_api_optional_params[ + "presence_penalty" # type: ignore[typeddict-item] + ] + if response_api_optional_params.get("service_tier"): + mapped_params["service_tier"] = response_api_optional_params["service_tier"] + return mapped_params def _transform_tools(self, tools: List[Dict[str, Any]]) -> List[Dict[str, Any]]: """ - Transform tools to Perplexity format - - Perplexity supports: + Transform tools to Perplexity format. + + Perplexity supports (per public OpenAPI spec): - web_search: Performs web searches - fetch_url: Fetches content from URLs + - function: Function Calling """ perplexity_tools = [] - + for tool in tools: if isinstance(tool, dict): - tool_type = tool.get("type") - + tool_type = tool.get("type", "") + # Direct Perplexity tool format if tool_type in ["web_search", "fetch_url"]: perplexity_tools.append(tool) - - # OpenAI function format - try to map to Perplexity tools + + # Function tools: Perplexity supports them natively elif tool_type == "function": - function = tool.get("function", {}) - function_name = function.get("name", "") - - if function_name == "web_search" or "search" in function_name.lower(): - perplexity_tools.append({"type": "web_search"}) - elif function_name == "fetch_url" or "fetch" in function_name.lower(): - perplexity_tools.append({"type": "fetch_url"}) - + perplexity_tools.append(tool) + return perplexity_tools def transform_responses_api_request( @@ -204,24 +281,26 @@ class PerplexityResponsesConfig(OpenAIResponsesAPIConfig): "model": model, "input": self._format_input(input), } - + # Add all optional parameters for key, value in response_api_optional_request_params.items(): data[key] = value - + return data - def _format_input(self, input: Union[str, ResponseInputParam]) -> Union[str, List[Dict[str, Any]]]: + def _format_input( + self, input: Union[str, ResponseInputParam] + ) -> Union[str, List[Dict[str, Any]]]: """ Format input for Perplexity Responses API - + The API accepts either: - A simple string for single-turn queries - An array of message objects for multi-turn conversations """ if isinstance(input, str): return input - + # Handle ResponseInputParam format if isinstance(input, list): formatted_messages = [] @@ -234,7 +313,7 @@ class PerplexityResponsesConfig(OpenAIResponsesAPIConfig): } formatted_messages.append(formatted_message) return formatted_messages - + return str(input) def transform_response_api_response( @@ -267,10 +346,14 @@ class PerplexityResponsesConfig(OpenAIResponsesAPIConfig): # Transform usage to handle Perplexity's cost structure usage_data = raw_response_json.get("usage", {}) transformed_usage_dict = self._transform_usage(usage_data) - + # Convert usage dict to ResponseAPIUsage object - usage_obj = ResponseAPIUsage(**transformed_usage_dict) if transformed_usage_dict else None - + usage_obj = ( + ResponseAPIUsage(**transformed_usage_dict) + if transformed_usage_dict + else None + ) + # Map Perplexity response to OpenAI Responses API format response = ResponsesAPIResponse( id=raw_response_json.get("id", ""), @@ -283,11 +366,11 @@ class PerplexityResponsesConfig(OpenAIResponsesAPIConfig): ) return response - + def _transform_usage(self, usage_data: Dict[str, Any]) -> Dict[str, Any]: """ Transform Perplexity usage data to OpenAI format - + Perplexity returns: { "input_tokens": 100, @@ -300,7 +383,7 @@ class PerplexityResponsesConfig(OpenAIResponsesAPIConfig): "total_cost": 0.0003 } } - + OpenAI expects: { "input_tokens": 100, @@ -314,7 +397,7 @@ class PerplexityResponsesConfig(OpenAIResponsesAPIConfig): "output_tokens": usage_data.get("output_tokens", 0), "total_tokens": usage_data.get("total_tokens", 0), } - + # Transform cost from Perplexity format (dict) to OpenAI format (float) cost_obj = usage_data.get("cost") if isinstance(cost_obj, dict) and "total_cost" in cost_obj: @@ -322,20 +405,20 @@ class PerplexityResponsesConfig(OpenAIResponsesAPIConfig): verbose_logger.debug( "Transformed Perplexity cost object to float: %s -> %s", cost_obj, - cost_obj["total_cost"] + cost_obj["total_cost"], ) elif cost_obj is not None: # If cost is already a float/number, use it as-is transformed["cost"] = cost_obj - + # Add input_tokens_details if present if "input_tokens_details" in usage_data: transformed["input_tokens_details"] = usage_data["input_tokens_details"] - + # Add output_tokens_details if present if "output_tokens_details" in usage_data: transformed["output_tokens_details"] = usage_data["output_tokens_details"] - + return transformed def transform_streaming_response( @@ -353,10 +436,10 @@ class PerplexityResponsesConfig(OpenAIResponsesAPIConfig): event_pydantic_model = PerplexityResponsesConfig.get_event_model_class( event_type=event_type ) - + # Transform Perplexity-specific fields to OpenAI format parsed_chunk = self._transform_perplexity_chunk(parsed_chunk) - + # Defensive: Handle error.code being null (similar to OpenAI implementation) try: error_obj = parsed_chunk.get("error") @@ -375,13 +458,13 @@ class PerplexityResponsesConfig(OpenAIResponsesAPIConfig): def _transform_perplexity_chunk(self, chunk: dict) -> dict: """ Transform Perplexity-specific fields in a streaming chunk to OpenAI format. - + This handles: - Converting Perplexity's cost object to a simple float """ # Make a copy to avoid modifying the original chunk = dict(chunk) - + # Transform usage.cost from Perplexity format to OpenAI format # Perplexity: {"currency": "USD", "input_cost": 0.0001, "output_cost": 0.0002, "total_cost": 0.0003} # OpenAI: 0.0003 (just the total_cost as a float) @@ -400,10 +483,10 @@ class PerplexityResponsesConfig(OpenAIResponsesAPIConfig): verbose_logger.debug( "Transformed Perplexity cost object to float: %s -> %s", cost_obj, - cost_obj["total_cost"] + cost_obj["total_cost"], ) except Exception as e: # If transformation fails, log and continue with original chunk verbose_logger.debug("Failed to transform Perplexity cost object: %s", e) - + return chunk diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 5c72bfbc13..7e90c64efd 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -2463,8 +2463,8 @@ class MCPServerManager: # Check if we should skip health check based on auth configuration should_skip_health_check = False - # Skip if auth_type is oauth2 - if server.needs_user_oauth_token: + # Skip if server requires per-user authentication (OAuth2 or passthrough auth) + if server.requires_per_user_auth: should_skip_health_check = True # Skip if auth_type is not none and authentication_token is missing elif ( @@ -2604,6 +2604,7 @@ class MCPServerManager: server.mcp_info.get("description") if server.mcp_info else None ), url=server.url, + spec_path=server.spec_path, transport=server.transport, auth_type=server.auth_type, created_at=datetime.now(), diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 9f2765652f..35a75d4caf 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -714,8 +714,9 @@ async def _initialize_shared_aiohttp_session(): connector_kwargs = { "keepalive_timeout": AIOHTTP_KEEPALIVE_TIMEOUT, "ttl_dns_cache": AIOHTTP_TTL_DNS_CACHE, - "enable_cleanup_closed": True, } + if AIOHTTP_NEEDS_CLEANUP_CLOSED: + connector_kwargs["enable_cleanup_closed"] = True if AIOHTTP_CONNECTOR_LIMIT > 0: connector_kwargs["limit"] = AIOHTTP_CONNECTOR_LIMIT if AIOHTTP_CONNECTOR_LIMIT_PER_HOST > 0: diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 536cb73de9..c4ff325db1 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -13,8 +13,6 @@ from email.mime.text import MIMEText from typing import ( TYPE_CHECKING, Any, - Callable, - Coroutine, Dict, List, Literal, @@ -2277,6 +2275,11 @@ class PrismaClient: 0.0, float(os.getenv("PRISMA_AUTH_RECONNECT_LOCK_TIMEOUT_SECONDS", "0.1")), ) + self._engine_pidfd: int = -1 + self._engine_pid: int = 0 + self._watching_engine: bool = False + self._engine_confirmed_dead: bool = False + self._engine_wait_thread: Optional[threading.Thread] = None verbose_proxy_logger.debug("Success - Created Prisma Client") def get_request_status( @@ -3544,31 +3547,368 @@ class PrismaClient: ) raise e + def _get_engine_pid(self) -> int: + try: + engine = self.db._original_prisma._engine # type: ignore[attr-defined] + if engine is not None and engine.process is not None: + return engine.process.pid + except (AttributeError, TypeError): + pass + return 0 + + def _is_engine_alive(self) -> bool: + if self._engine_pid <= 0: + return True + try: + os.kill(self._engine_pid, 0) + return True + except ProcessLookupError: + return False + except (PermissionError, OSError): + return True + + @staticmethod + def _reap_all_zombies() -> set: + """Reap ALL zombie child processes via waitpid(-1, WNOHANG). + + Returns a set of reaped PIDs. As PID 1 in Docker (or any + process that spawns children), we must reap ALL terminated + children to prevent zombie accumulation. + """ + reaped: set = set() + while True: + try: + pid, _ = os.waitpid(-1, os.WNOHANG) + if pid == 0: + break + reaped.add(pid) + except ChildProcessError: + break + return reaped + + def _try_waitpid_watch(self, pid: int) -> bool: + """Watch engine PID via os.waitpid() in a dedicated thread. + + The thread blocks on os.waitpid(pid, 0) which is a kernel-level + wait and with zero CPU overhead, instant detection when the process exits. + When the process dies, the thread notifies the asyncio event loop + via call_soon_threadsafe. + + Returns True if the thread was started, False on failure. + """ + try: + probe_pid, _ = os.waitpid(pid, os.WNOHANG) + except ChildProcessError: + verbose_proxy_logger.debug( + "PID %s is not a child process; skipping waitpid watch.", pid, + ) + return False + + if probe_pid == pid: + verbose_proxy_logger.warning( + "prisma-query-engine PID %s already dead at watch start.", pid, + ) + self._engine_confirmed_dead = True + self._reap_all_zombies() + self._cleanup_engine_watcher() + asyncio.create_task( + self.attempt_db_reconnect( + reason="engine_process_death", + force=True, + ) + ) + return True + + try: + loop = asyncio.get_running_loop() + except RuntimeError: + return False + + thread = threading.Thread( + target=self._waitpid_thread_func, + args=(pid, loop), + daemon=True, + name=f"prisma-engine-waitpid-{pid}", + ) + thread.start() + self._engine_wait_thread = thread + return True + + def _waitpid_thread_func(self, pid: int, loop: asyncio.AbstractEventLoop) -> None: + """Thread function: block until engine PID exits, then notify event loop. + + Note: uvloop/libuv may reap the child first via waitpid(-1, WNOHANG) + in its SIGCHLD handler. In that case our waitpid raises ChildProcessError. + we still notify the event loop because the engine is dead either way. + """ + try: + os.waitpid(pid, 0) + except ChildProcessError: + pass + except OSError: + pass + try: + loop.call_soon_threadsafe(self._on_engine_death_from_thread, pid) + except RuntimeError: + pass + + def _on_engine_death_from_thread(self, dead_pid: int) -> None: + """Called on the event loop thread when the waitpid thread detects engine death.""" + if self._engine_confirmed_dead: + return + if dead_pid != self._engine_pid: + return + verbose_proxy_logger.error( + "prisma-query-engine PID %s exited (waitpid thread); triggering reconnect.", + dead_pid, + ) + self._engine_confirmed_dead = True + self._reap_all_zombies() + self._cleanup_engine_watcher() + asyncio.create_task( + self.attempt_db_reconnect( + reason="engine_process_death", + force=True, + ) + ) + + def _try_pidfd_watch(self, pid: int) -> bool: + """ + Watch engine PID via pidfd_open + asyncio event loop reader. + + Returns True if pidfd watch was set up, False if unavailable or failed. + Broad OSError catch handles both ENOSYS and SECCOMP-blocked syscalls. + """ + if not hasattr(os, "pidfd_open"): + return False + fd = -1 + try: + fd = os.pidfd_open(pid, 0) # type: ignore[attr-defined] + asyncio.get_running_loop().add_reader(fd, self._on_pidfd_readable) + self._engine_pidfd = fd + return True + except OSError: + if fd >= 0: + os.close(fd) + return False + + def _on_pidfd_readable(self) -> None: + """pidfd became readable: engine process exited or became zombie. + + Sets _engine_confirmed_dead BEFORE cleanup so _run_reconnect_cycle + takes the heavy path (recreate Prisma client + re-arm watcher). + """ + if self._engine_confirmed_dead: + # Already handled -- just clean up pidfd resources. + if self._engine_pidfd >= 0: + try: + asyncio.get_running_loop().remove_reader(self._engine_pidfd) + except Exception: + pass + try: + os.close(self._engine_pidfd) + except OSError: + pass + self._engine_pidfd = -1 + return + dead_pid = self._engine_pid + verbose_proxy_logger.error( + "prisma-query-engine PID %s exited (pidfd event); triggering reconnect.", + dead_pid, + ) + self._engine_confirmed_dead = True + self._reap_all_zombies() + self._cleanup_engine_watcher() + asyncio.create_task( + self.attempt_db_reconnect( + reason="engine_process_death", + force=True, + ) + ) + + async def _poll_engine_proc(self) -> None: + """poll via os.kill(pid, 0) every 1s. + Only used when BOTH waitpid thread and pidfd are unavailable + (e.g., PID is not our child process and pidfd_open fails) + """ + while self._watching_engine and self._engine_pid > 0: + try: + os.kill(self._engine_pid, 0) + except ProcessLookupError: + verbose_proxy_logger.error( + "prisma-query-engine PID %s gone; triggering reconnect.", + self._engine_pid, + ) + self._engine_confirmed_dead = True + self._reap_all_zombies() + self._cleanup_engine_watcher() + await self.attempt_db_reconnect( + reason="engine_process_death", + force=True, + ) + return + except (PermissionError, OSError): + verbose_proxy_logger.debug( + "Cannot signal PID %s; stopping engine poll.", + self._engine_pid, + ) + self._cleanup_engine_watcher() + return + await asyncio.sleep(1) + + def _cleanup_engine_watcher(self) -> None: + """Clean up pidfd reader, waitpid thread ref, or stop polling and reset state.""" + self._watching_engine = False + if self._engine_pidfd >= 0: + try: + asyncio.get_running_loop().remove_reader(self._engine_pidfd) + except Exception: + pass + try: + os.close(self._engine_pidfd) + except OSError: + pass + self._engine_pidfd = -1 + self._engine_wait_thread = None + self._engine_pid = 0 + + async def _start_engine_watcher(self) -> None: + """ + Start watching the Prisma query engine process for death. + + Detection priority: + 1. os.waitpid() in a dedicated thread, works with all event loops. + 2. pidfd_open kernel fd registered with asyncio. + 3. os.kill(pid, 0) polling (1s), last-resort fallback when neither + waitpid thread nor pidfd are available. + + """ + if self._watching_engine or self._engine_pidfd >= 0 or self._engine_wait_thread is not None: + return + pid = self._get_engine_pid() + if pid == 0: + verbose_proxy_logger.debug("Could not find prisma-query-engine PID; engine death detection unavailable.") + return + self._engine_pid = pid + self._engine_confirmed_dead = False + verbose_proxy_logger.info("Found prisma-query-engine at PID %s.", pid) + waitpid_ok = self._try_waitpid_watch(pid) + pidfd_ok = False if waitpid_ok else self._try_pidfd_watch(pid) + if waitpid_ok: + verbose_proxy_logger.info( + "Watching engine PID %s via waitpid thread.", pid, + ) + elif pidfd_ok: + verbose_proxy_logger.info( + "Watching engine PID %s via pidfd.", pid, + ) + else: + verbose_proxy_logger.info( + "Watching engine PID %s via os.kill polling.", pid, + ) + self._watching_engine = True + asyncio.create_task(self._poll_engine_proc()) + + def _stop_engine_watcher(self) -> None: + """Stop watching the engine process and clean up all resources.""" + self._cleanup_engine_watcher() + self._engine_confirmed_dead = False + verbose_proxy_logger.debug("Stopped engine process watcher.") + async def _run_reconnect_cycle( self, timeout_seconds: Optional[float] = None ) -> None: """ - Run a reconnect cycle with direct db operations and a single overall timeout - budget to avoid long retries on hot paths (e.g. auth). + Run a reconnect cycle with a single overall timeout budget. + + Uses the _engine_confirmed_dead flag (set by waitpid thread / pidfd / poll + handlers) to choose between heavy reconnect (engine dead -- recreate + Prisma client, re-arm watcher) and lightweight reconnect (network + blip -- disconnect, connect, SELECT 1). """ - async def _do_direct_reconnect() -> None: - try: - await self.db.disconnect() - except Exception as disconnect_err: - verbose_proxy_logger.debug( - "Prisma DB disconnect before reconnect failed (ignored): %s", - disconnect_err, - ) - - await self.db.connect() - await self.db.query_raw("SELECT 1") - effective_timeout = ( - timeout_seconds - if timeout_seconds is not None - else self._db_watchdog_reconnect_timeout_seconds + timeout_seconds if timeout_seconds is not None else self._db_watchdog_reconnect_timeout_seconds ) - await asyncio.wait_for(_do_direct_reconnect(), timeout=effective_timeout) + + engine_is_dead = self._engine_confirmed_dead or ( + self._engine_pid > 0 and not self._is_engine_alive() + ) + + if engine_is_dead: + dead_pid = self._engine_pid + verbose_proxy_logger.warning( + "prisma-query-engine PID %s is dead; reconnecting.", + dead_pid, + ) + self._reap_all_zombies() + self._cleanup_engine_watcher() + self._engine_confirmed_dead = False + + async def _do_heavy_reconnect() -> None: + db_url = os.getenv("DATABASE_URL", "") + if not db_url: + verbose_proxy_logger.error("DATABASE_URL not set; cannot recreate Prisma client.") + raise RuntimeError("DATABASE_URL not set") + await self.db.recreate_prisma_client(db_url) + await self._start_engine_watcher() + + await asyncio.wait_for(_do_heavy_reconnect(), timeout=effective_timeout) + else: + verbose_proxy_logger.debug("Performing Prisma DB reconnect (engine alive or unknown).") + + async def _do_direct_reconnect() -> None: + try: + await self.db.disconnect() + except Exception as disconnect_err: + verbose_proxy_logger.debug( + "Prisma DB disconnect before reconnect failed (ignored): %s", + disconnect_err, + ) + + await self.db.connect() + await self.db.query_raw("SELECT 1") + + await asyncio.wait_for(_do_direct_reconnect(), timeout=effective_timeout) + + async def _attempt_reconnect_inside_lock( + self, + force: bool, + reason: str, + timeout_seconds: Optional[float], + ) -> bool: + now = time.time() + if ( + force is False + and now - self._db_last_reconnect_attempt_ts + < self._db_reconnect_cooldown_seconds + ): + verbose_proxy_logger.debug( + "Skipping DB reconnect attempt inside lock due to cooldown. reason=%s", + reason, + ) + return False + + verbose_proxy_logger.warning( + "Attempting Prisma DB reconnect. reason=%s", reason + ) + + reconnect_succeeded = False + try: + await self._run_reconnect_cycle(timeout_seconds=timeout_seconds) + reconnect_succeeded = True + verbose_proxy_logger.info( + "Prisma DB reconnect succeeded. reason=%s", reason + ) + except Exception as reconnect_err: + verbose_proxy_logger.error( + "Prisma DB reconnect failed. reason=%s error=%s", + reason, + reconnect_err, + ) + finally: + self._db_last_reconnect_attempt_ts = time.time() + + return reconnect_succeeded async def attempt_db_reconnect( self, @@ -3595,59 +3935,10 @@ class PrismaClient: ) return False - async def _attempt_reconnect_inside_lock() -> bool: - now = time.time() - if ( - force is False - and now - self._db_last_reconnect_attempt_ts - < self._db_reconnect_cooldown_seconds - ): - verbose_proxy_logger.debug( - "Skipping DB reconnect attempt inside lock due to cooldown. reason=%s", - reason, - ) - return False - - verbose_proxy_logger.warning( - "Attempting Prisma DB reconnect. reason=%s", reason - ) - - reconnect_succeeded = False - try: - await self._run_reconnect_cycle(timeout_seconds=timeout_seconds) - reconnect_succeeded = True - verbose_proxy_logger.info( - "Prisma DB reconnect succeeded. reason=%s", reason - ) - except Exception as reconnect_err: - verbose_proxy_logger.error( - "Prisma DB reconnect failed. reason=%s error=%s", - reason, - reconnect_err, - ) - finally: - # Start cooldown after reconnect attempt has completed. - self._db_last_reconnect_attempt_ts = time.time() - - return reconnect_succeeded - if lock_timeout_seconds is None: async with self._db_reconnect_lock: - return await _attempt_reconnect_inside_lock() + return await self._attempt_reconnect_inside_lock(force, reason, timeout_seconds) - return await self._attempt_reconnect_with_lock_timeout( - _attempt_reconnect_inside_lock, - reason=reason, - lock_timeout_seconds=lock_timeout_seconds, - ) - - async def _attempt_reconnect_with_lock_timeout( - self, - reconnect_fn: Callable[[], Coroutine[Any, Any, bool]], - reason: str, - lock_timeout_seconds: float, - ) -> bool: - """Acquire the reconnect lock with a timeout, then run reconnect_fn.""" lock_acquired_by_timeout_task = False async def _acquire_reconnect_lock() -> bool: @@ -3695,14 +3986,14 @@ class PrismaClient: return False try: - return await reconnect_fn() + return await self._attempt_reconnect_inside_lock(force, reason, timeout_seconds) finally: self._db_reconnect_lock.release() async def start_db_health_watchdog_task(self) -> None: - """ - Start a background task that probes DB health and attempts reconnect on failure. - """ + """Start background tasks that monitor DB health: + - A periodic SELECT 1 probe that triggers reconnect on network/connection failure. + - A process-level watcher that detects engine death via waitpid thread, pidfd, or os.kill polling.""" if self._db_health_watchdog_enabled is not True: verbose_proxy_logger.debug( "Prisma DB health watchdog disabled via PRISMA_HEALTH_WATCHDOG_ENABLED" @@ -3720,11 +4011,11 @@ class PrismaClient: self._db_health_watchdog_probe_timeout_seconds, self._db_watchdog_reconnect_timeout_seconds, ) + await self._start_engine_watcher() async def stop_db_health_watchdog_task(self) -> None: - """ - Stop DB health watchdog task gracefully. - """ + """Stop DB health watchdog task and engine watcher gracefully.""" + self._stop_engine_watcher() if self._db_health_watchdog_task is None: return self._db_health_watchdog_task.cancel() diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index 2cd385c5bf..7f99fd526c 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -65,3 +65,25 @@ class MCPServer(BaseModel): def needs_user_oauth_token(self) -> bool: """True if this is an OAuth2 server that relies on per-user tokens (no client_credentials).""" return self.auth_type == MCPAuth.oauth2 and not self.has_client_credentials + + @property + def requires_per_user_auth(self) -> bool: + """ + True if this server requires per-user/per-request authentication. + This includes: + - OAuth2 servers without client credentials + - Servers with auth_type=none but extra_headers configured for auth passthrough + + Health checks should be skipped for these servers since they cannot + authenticate without user-provided credentials. + """ + # OAuth2 without client credentials + if self.needs_user_oauth_token: + return True + + # PAT passthrough: auth_type is none but extra_headers includes auth headers + if self.auth_type == MCPAuth.none and self.extra_headers: + auth_header_names = {"authorization", "x-api-key", "api-key", "apikey"} + return any(h.lower() in auth_header_names for h in self.extra_headers) + + return False diff --git a/litellm/utils.py b/litellm/utils.py index fef99c8b20..5046929257 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5385,14 +5385,18 @@ def _get_max_position_embeddings(model_name: str) -> Optional[int]: @lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE) def _cached_get_model_info_helper( - model: str, custom_llm_provider: Optional[str] + model: str, + custom_llm_provider: Optional[str], + api_base: Optional[str] = None, ) -> ModelInfoBase: """ _get_model_info_helper wrapped with lru_cache Speed Optimization to hit high RPS """ - return _get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider) + return _get_model_info_helper( + model=model, custom_llm_provider=custom_llm_provider, api_base=api_base + ) def get_provider_info( @@ -5428,7 +5432,9 @@ def _is_potential_model_name_in_model_cost( def _get_model_info_helper( # noqa: PLR0915 - model: str, custom_llm_provider: Optional[str] = None + model: str, + custom_llm_provider: Optional[str] = None, + api_base: Optional[str] = None, ) -> ModelInfoBase: """ Helper for 'get_model_info'. Separated out to avoid infinite loop caused by returning 'supported_openai_param's @@ -5486,7 +5492,7 @@ def _get_model_info_helper( # noqa: PLR0915 elif ( custom_llm_provider == "ollama" or custom_llm_provider == "ollama_chat" ) and not _is_potential_model_name_in_model_cost(potential_model_names): - return litellm.OllamaConfig().get_model_info(model) + return litellm.OllamaConfig().get_model_info(model, api_base=api_base) else: """ Check if: (in order of specificity) @@ -5725,8 +5731,6 @@ def _get_model_info_helper( # noqa: PLR0915 ) except Exception as e: verbose_logger.debug(f"Error getting model info: {e}") - if "OllamaError" in str(e): - raise e raise Exception( "This model isn't mapped yet. model={}, custom_llm_provider={}. Add it here - https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_window.json.".format( model, custom_llm_provider @@ -5735,7 +5739,11 @@ def _get_model_info_helper( # noqa: PLR0915 @lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE) -def get_model_info(model: str, custom_llm_provider: Optional[str] = None) -> ModelInfo: +def get_model_info( + model: str, + custom_llm_provider: Optional[str] = None, + api_base: Optional[str] = None, +) -> ModelInfo: """ Get a dict for the maximum tokens (context window), input_cost_per_token, output_cost_per_token for a given model. @@ -5813,6 +5821,7 @@ def get_model_info(model: str, custom_llm_provider: Optional[str] = None) -> Mod _model_info = _get_model_info_helper( model=model, custom_llm_provider=custom_llm_provider, + api_base=api_base, ) provider_info = get_provider_info( @@ -5823,8 +5832,8 @@ def get_model_info(model: str, custom_llm_provider: Optional[str] = None) -> Mod if value is not None: _model_info[key] = value # type: ignore - if verbose_logger.isEnabledFor(logging.DEBUG): - verbose_logger.debug(f"model_info: {_model_info}") + # if verbose_logger.isEnabledFor(logging.DEBUG): + # verbose_logger.debug(f"model_info: {_model_info}") returned_model_info = ModelInfo( **_model_info, supported_openai_params=supported_openai_params diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 3909ce4c8b..624714feb8 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -26555,65 +26555,124 @@ "supports_function_calling": true, "supports_tool_choice": true }, + "perplexity/preset/fast-search": { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_web_search": true, + "supports_preset": true, + "supports_function_calling": true + }, "perplexity/preset/pro-search": { "litellm_provider": "perplexity", "mode": "responses", "supports_web_search": true, - "supports_preset": true + "supports_preset": true, + "supports_function_calling": true }, - "perplexity/openai/gpt-4o": { + "perplexity/preset/deep-research": { "litellm_provider": "perplexity", "mode": "responses", "supports_web_search": true, - "supports_reasoning": false + "supports_preset": true, + "supports_function_calling": true }, - "perplexity/openai/gpt-4o-mini": { + "perplexity/preset/advanced-deep-research": { "litellm_provider": "perplexity", "mode": "responses", "supports_web_search": true, - "supports_reasoning": false + "supports_preset": true, + "supports_function_calling": true }, "perplexity/openai/gpt-5.2": { "litellm_provider": "perplexity", "mode": "responses", "supports_web_search": true, - "supports_reasoning": true + "supports_reasoning": true, + "supports_function_calling": true }, - "perplexity/anthropic/claude-3-5-sonnet-20241022": { + "perplexity/openai/gpt-5.1": { "litellm_provider": "perplexity", "mode": "responses", "supports_web_search": true, - "supports_reasoning": false + "supports_reasoning": false, + "supports_function_calling": true }, - "perplexity/anthropic/claude-3-5-haiku-20241022": { + "perplexity/openai/gpt-5-mini": { "litellm_provider": "perplexity", "mode": "responses", "supports_web_search": true, - "supports_reasoning": false + "supports_reasoning": false, + "supports_function_calling": true }, - "perplexity/google/gemini-2.0-flash-exp": { + "perplexity/anthropic/claude-opus-4-6": { "litellm_provider": "perplexity", "mode": "responses", "supports_web_search": true, - "supports_reasoning": false + "supports_reasoning": false, + "supports_function_calling": true }, - "perplexity/google/gemini-2.0-flash-thinking-exp": { + "perplexity/anthropic/claude-opus-4-5": { "litellm_provider": "perplexity", "mode": "responses", "supports_web_search": true, - "supports_reasoning": true + "supports_reasoning": false, + "supports_function_calling": true }, - "perplexity/xai/grok-2-1212": { + "perplexity/anthropic/claude-sonnet-4-5": { "litellm_provider": "perplexity", "mode": "responses", "supports_web_search": true, - "supports_reasoning": false + "supports_reasoning": false, + "supports_function_calling": true }, - "perplexity/xai/grok-2-vision-1212": { + "perplexity/anthropic/claude-haiku-4-5": { "litellm_provider": "perplexity", "mode": "responses", "supports_web_search": true, - "supports_reasoning": false + "supports_reasoning": false, + "supports_function_calling": true + }, + "perplexity/google/gemini-3-pro-preview": { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_web_search": true, + "supports_reasoning": false, + "supports_function_calling": true + }, + "perplexity/google/gemini-3-flash-preview": { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_web_search": true, + "supports_reasoning": false, + "supports_function_calling": true + }, + "perplexity/google/gemini-2.5-pro": { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_web_search": true, + "supports_reasoning": false, + "supports_function_calling": true + }, + "perplexity/google/gemini-2.5-flash": { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_web_search": true, + "supports_reasoning": false, + "supports_function_calling": true + }, + "perplexity/xai/grok-4-1-fast-non-reasoning": { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_web_search": true, + "supports_reasoning": false, + "supports_function_calling": true + }, + "perplexity/perplexity/sonar": { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_web_search": true, + "supports_reasoning": false, + "supports_function_calling": true }, "publicai/aisingapore/Qwen-SEA-LION-v4-32B-IT": { "input_cost_per_token": 0.0, diff --git a/tests/litellm/proxy/test_prisma_engine_watchdog.py b/tests/litellm/proxy/test_prisma_engine_watchdog.py new file mode 100644 index 0000000000..011b8002db --- /dev/null +++ b/tests/litellm/proxy/test_prisma_engine_watchdog.py @@ -0,0 +1,446 @@ +""" +Tests for PrismaClient engine watchdog: death detection and automatic reconnect. + +Covers: +- Engine PID discovery and liveness check +- Engine process gone (os.kill raises ProcessLookupError) → reconnect triggered +- PermissionError from os.kill → treated as alive (process exists but not ours) +- pidfd handler → schedules attempt_db_reconnect even when lock is held +- waitpid thread → instant cross-platform detection, triggers reconnect +- _run_reconnect_cycle branches: heavy path (engine dead) vs lightweight path (engine alive) +- _engine_confirmed_dead flag ensures heavy reconnect even after _engine_pid reset +- Successful heavy reconnect → watcher re-armed for new process +- Missing DATABASE_URL → graceful RuntimeError in reconnect cycle +- Shutdown → polling loop exits cleanly +""" + +import asyncio +import os +import threading +import time +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from litellm.proxy.utils import PrismaClient, ProxyLogging + + +@pytest.fixture(autouse=True) +def mock_prisma_binary(): + """Mock prisma.Prisma to avoid requiring generated Prisma binaries for unit tests.""" + import sys + + mock_module = MagicMock() + with patch.dict(sys.modules, {"prisma": mock_module}): + yield + + +@pytest.fixture +def mock_proxy_logging(): + proxy_logging = AsyncMock(spec=ProxyLogging) + proxy_logging.failure_handler = AsyncMock() + return proxy_logging + + +@pytest.fixture +def engine_client(mock_proxy_logging) -> PrismaClient: + """ + Minimal PrismaClient fixture for engine watchdog tests. + Uses the real constructor pattern from PR #21706 (database_url). + """ + client = PrismaClient(database_url="mock://test", proxy_logging_obj=mock_proxy_logging) + client.db = MagicMock() + client.db.recreate_prisma_client = AsyncMock() + client.db.disconnect = AsyncMock(return_value=None) + client.db.connect = AsyncMock(return_value=None) + client.db.query_raw = AsyncMock(return_value=[{"result": 1}]) + return client + + +# --------------------------------------------------------------------------- +# _is_engine_alive +# --------------------------------------------------------------------------- + + +def test_is_engine_alive_returns_true_when_pid_unknown(engine_client): + """_is_engine_alive returns True when no engine PID is tracked.""" + engine_client._engine_pid = 0 + assert engine_client._is_engine_alive() is True + + +def test_is_engine_alive_returns_false_when_process_gone(engine_client): + """_is_engine_alive returns False when os.kill raises ProcessLookupError.""" + engine_client._engine_pid = 9999 + with patch("os.kill", side_effect=ProcessLookupError): + assert engine_client._is_engine_alive() is False + + +def test_is_engine_alive_returns_true_on_permission_error(engine_client): + """_is_engine_alive returns True when os.kill raises PermissionError (process exists but not ours).""" + engine_client._engine_pid = 1234 + with patch("os.kill", side_effect=PermissionError): + assert engine_client._is_engine_alive() is True + + +def test_is_engine_alive_returns_true_for_running_process(engine_client): + """_is_engine_alive returns True when os.kill succeeds (process running).""" + engine_client._engine_pid = 1234 + with patch("os.kill"): + assert engine_client._is_engine_alive() is True + + +# --------------------------------------------------------------------------- +# _poll_engine_proc — calls attempt_db_reconnect on death +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_poll_missing_process_triggers_reconnect(engine_client) -> None: + """Polling loop triggers attempt_db_reconnect when os.kill raises ProcessLookupError.""" + engine_client._engine_pid = 1234 + engine_client._watching_engine = True + engine_client.attempt_db_reconnect = AsyncMock(return_value=True) + + with patch("os.kill", side_effect=ProcessLookupError): + await engine_client._poll_engine_proc() + + engine_client.attempt_db_reconnect.assert_awaited_once_with( + reason="engine_process_death", + force=True, + ) + + +@pytest.mark.asyncio +async def test_poll_permission_error_stops_polling(engine_client) -> None: + """Polling loop stops cleanly when os.kill raises PermissionError (process not ours).""" + engine_client._engine_pid = 1234 + engine_client._watching_engine = True + engine_client.attempt_db_reconnect = AsyncMock(return_value=True) + + with patch("os.kill", side_effect=PermissionError): + await engine_client._poll_engine_proc() + + # PermissionError means process exists but isn't ours — no reconnect, just stop polling + engine_client.attempt_db_reconnect.assert_not_awaited() + assert engine_client._watching_engine is False + assert engine_client._engine_pid == 0 + + +@pytest.mark.asyncio +async def test_stop_loop_halts_polling(engine_client) -> None: + """Polling loop exits cleanly when _stop_engine_watcher is called.""" + engine_client._engine_pid = 1234 + engine_client._watching_engine = True + + async def stop_during_sleep(_duration: float) -> None: + engine_client._stop_engine_watcher() + + with ( + patch("os.kill"), + patch("asyncio.sleep", side_effect=stop_during_sleep), + ): + await engine_client._poll_engine_proc() + + assert engine_client._watching_engine is False + assert engine_client._engine_pid == 0 + + +# --------------------------------------------------------------------------- +# _on_pidfd_readable — calls attempt_db_reconnect +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_pidfd_readable_schedules_reconnect(engine_client) -> None: + """pidfd handler schedules attempt_db_reconnect via asyncio.create_task.""" + engine_client._engine_pid = 1234 + engine_client.attempt_db_reconnect = AsyncMock(return_value=True) + + created_coros = [] + + def capture_task(coro): + created_coros.append(coro) + return MagicMock() + + with patch("asyncio.create_task", side_effect=capture_task): + engine_client._on_pidfd_readable() + + # Run the captured coroutine to completion + assert len(created_coros) == 1 + await created_coros[0] + + engine_client.attempt_db_reconnect.assert_awaited_once_with( + reason="engine_process_death", + force=True, + ) + + +@pytest.mark.asyncio +async def test_pidfd_schedules_reconnect_task_when_lock_held(engine_client) -> None: + """pidfd handler schedules reconnect task even when _db_reconnect_lock is held.""" + engine_client._engine_pid = 1234 + + created_coros = [] + + def capture_task(coro): + created_coros.append(coro) + return MagicMock() + + async with engine_client._db_reconnect_lock: + with patch("asyncio.create_task", side_effect=capture_task): + engine_client._on_pidfd_readable() + + for coro in created_coros: + coro.close() + + assert len(created_coros) == 1 + + +# --------------------------------------------------------------------------- +# _run_reconnect_cycle — engine liveness branching +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_run_reconnect_cycle_uses_heavy_path_when_engine_dead( + engine_client, +) -> None: + """_run_reconnect_cycle calls recreate_prisma_client when engine is dead.""" + engine_client._engine_pid = 1234 + engine_client._start_engine_watcher = AsyncMock() + + with ( + patch.object(engine_client, "_is_engine_alive", return_value=False), + patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}), + patch("os.waitpid", side_effect=ChildProcessError), + ): + await engine_client._run_reconnect_cycle(timeout_seconds=5.0) + + engine_client.db.recreate_prisma_client.assert_awaited_once_with("postgresql://test") + engine_client._start_engine_watcher.assert_awaited_once() + engine_client.db.connect.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_run_reconnect_cycle_uses_heavy_path_when_confirmed_dead( + engine_client, +) -> None: + """_run_reconnect_cycle takes heavy path when _engine_confirmed_dead is set. + + This is the critical race-condition fix: SIGCHLD/pidfd handlers set + _engine_confirmed_dead BEFORE _cleanup_engine_watcher resets _engine_pid + to 0, so the heavy path executes even after cleanup. + """ + engine_client._engine_pid = 0 # Already reset by cleanup! + engine_client._engine_confirmed_dead = True # But flag survives cleanup + engine_client._start_engine_watcher = AsyncMock() + + with ( + patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}), + patch("os.waitpid", side_effect=ChildProcessError), + ): + await engine_client._run_reconnect_cycle(timeout_seconds=5.0) + + engine_client.db.recreate_prisma_client.assert_awaited_once_with("postgresql://test") + engine_client._start_engine_watcher.assert_awaited_once() + engine_client.db.connect.assert_not_awaited() + assert engine_client._engine_confirmed_dead is False # Reset after use + + +@pytest.mark.asyncio +async def test_run_reconnect_cycle_uses_lightweight_path_when_engine_alive( + engine_client, +) -> None: + """_run_reconnect_cycle uses disconnect/connect when engine is alive.""" + engine_client._engine_pid = 1234 + + with patch.object(engine_client, "_is_engine_alive", return_value=True): + await engine_client._run_reconnect_cycle(timeout_seconds=5.0) + + engine_client.db.connect.assert_awaited_once() + engine_client.db.query_raw.assert_awaited_once_with("SELECT 1") + engine_client.db.recreate_prisma_client.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_run_reconnect_cycle_uses_lightweight_path_when_pid_unknown( + engine_client, +) -> None: + """_run_reconnect_cycle uses lightweight path when engine PID is not tracked.""" + engine_client._engine_pid = 0 + + await engine_client._run_reconnect_cycle(timeout_seconds=5.0) + + engine_client.db.connect.assert_awaited_once() + engine_client.db.query_raw.assert_awaited_once_with("SELECT 1") + engine_client.db.recreate_prisma_client.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_run_reconnect_cycle_heavy_path_raises_without_database_url( + engine_client, +) -> None: + """Heavy reconnect raises RuntimeError when DATABASE_URL is not set.""" + engine_client._engine_pid = 1234 + + with ( + patch.object(engine_client, "_is_engine_alive", return_value=False), + patch.dict(os.environ, {}, clear=True), + patch("os.waitpid", side_effect=ChildProcessError), + ): + with pytest.raises(RuntimeError, match="DATABASE_URL not set"): + await engine_client._run_reconnect_cycle(timeout_seconds=5.0) + + engine_client.db.recreate_prisma_client.assert_not_awaited() + + +# --------------------------------------------------------------------------- +# start/stop lifecycle integration +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_start_watchdog_task_also_starts_engine_watcher( + engine_client, +) -> None: + """start_db_health_watchdog_task() also starts engine watcher.""" + engine_client._start_engine_watcher = AsyncMock() + + loop = asyncio.get_running_loop() + dummy_task = loop.create_task(asyncio.sleep(3600)) + + def fake_create_task(coro): + coro.close() + return dummy_task + + with patch("asyncio.create_task", side_effect=fake_create_task): + await engine_client.start_db_health_watchdog_task() + + engine_client._start_engine_watcher.assert_awaited_once() + dummy_task.cancel() + try: + await dummy_task + except asyncio.CancelledError: + pass + + +@pytest.mark.asyncio +async def test_stop_watchdog_task_also_stops_engine_watcher( + engine_client, +) -> None: + """stop_db_health_watchdog_task() also stops engine watcher.""" + engine_client._stop_engine_watcher = MagicMock() + + loop = asyncio.get_running_loop() + dummy_task = loop.create_task(asyncio.sleep(3600)) + engine_client._db_health_watchdog_task = dummy_task + + await engine_client.stop_db_health_watchdog_task() + + engine_client._stop_engine_watcher.assert_called_once() + assert engine_client._db_health_watchdog_task is None + + +# --------------------------------------------------------------------------- +# waitpid thread (cross-platform) +# --------------------------------------------------------------------------- + + +def test_try_waitpid_watch_returns_false_when_not_child(engine_client): + """_try_waitpid_watch returns False when PID is not our child process.""" + engine_client._engine_pid = 9999 + with patch("os.waitpid", side_effect=ChildProcessError): + assert engine_client._try_waitpid_watch(9999) is False + assert engine_client._engine_wait_thread is None + + +def test_try_waitpid_watch_starts_thread_for_child(engine_client): + """_try_waitpid_watch starts a daemon thread when PID is our child.""" + engine_client._engine_pid = 1234 + mock_thread = MagicMock() + mock_loop = MagicMock() + with ( + patch("os.waitpid", return_value=(0, 0)), + patch("asyncio.get_running_loop", return_value=mock_loop), + patch("threading.Thread", return_value=mock_thread) as mock_thread_cls, + ): + result = engine_client._try_waitpid_watch(1234) + assert result is True + mock_thread.start.assert_called_once() + assert engine_client._engine_wait_thread is mock_thread + + +@pytest.mark.asyncio +async def test_try_waitpid_watch_handles_already_dead_engine(engine_client) -> None: + """_try_waitpid_watch detects engine already dead at watch start.""" + engine_client._engine_pid = 1234 + engine_client.attempt_db_reconnect = AsyncMock(return_value=True) + + created_coros = [] + + def capture_task(coro): + created_coros.append(coro) + return MagicMock() + + waitpid_calls = iter([(1234, 0)]) + + def mock_waitpid(pid, flags): + if pid == -1: + raise ChildProcessError + return next(waitpid_calls) + + with ( + patch("os.waitpid", side_effect=mock_waitpid), + patch("asyncio.create_task", side_effect=capture_task), + ): + result = engine_client._try_waitpid_watch(1234) + + assert result is True + assert engine_client._engine_confirmed_dead is True + assert len(created_coros) == 1 + created_coros[0].close() + + +@pytest.mark.asyncio +async def test_on_engine_death_from_thread_triggers_reconnect(engine_client) -> None: + """waitpid thread callback schedules attempt_db_reconnect.""" + engine_client._engine_pid = 1234 + engine_client.attempt_db_reconnect = AsyncMock(return_value=True) + + created_coros = [] + + def capture_task(coro): + created_coros.append(coro) + return MagicMock() + + with patch("asyncio.create_task", side_effect=capture_task): + engine_client._on_engine_death_from_thread(1234) + + assert len(created_coros) == 1 + await created_coros[0] + + engine_client.attempt_db_reconnect.assert_awaited_once_with( + reason="engine_process_death", + force=True, + ) + + +def test_on_engine_death_from_thread_no_double_trigger(engine_client): + """waitpid thread callback does not trigger reconnect if already confirmed dead.""" + engine_client._engine_pid = 1234 + engine_client._engine_confirmed_dead = True + + with patch("asyncio.create_task") as mock_create_task: + engine_client._on_engine_death_from_thread(1234) + + mock_create_task.assert_not_called() + + +def test_on_engine_death_from_thread_ignores_stale_pid(engine_client): + """waitpid thread callback ignores death notification for a stale PID.""" + engine_client._engine_pid = 5678 + + with patch("asyncio.create_task") as mock_create_task: + engine_client._on_engine_death_from_thread(1234) + + mock_create_task.assert_not_called() diff --git a/tests/llm_translation/test_prompt_factory.py b/tests/llm_translation/test_prompt_factory.py index 8974632631..88bca00774 100644 --- a/tests/llm_translation/test_prompt_factory.py +++ b/tests/llm_translation/test_prompt_factory.py @@ -973,6 +973,54 @@ def test_convert_to_anthropic_tool_invoke_regular_tool(): assert result[0]["input"] == {"location": "San Francisco"} +def test_convert_to_anthropic_tool_invoke_sanitizes_invalid_ids(): + """Test that tool_use IDs with invalid characters are sanitized. + + Anthropic requires tool_use_id to match ^[a-zA-Z0-9_-]+$. + IDs from external frameworks (e.g. MiniMax) may contain characters + like colons that violate this pattern. + """ + tool_calls = [ + { + "id": "sessions_history:183", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"location": "Boston"}', + }, + }, + { + "id": "composio.NOTION_SEARCH", + "type": "function", + "function": { + "name": "search_notes", + "arguments": '{"query": "test"}', + }, + }, + ] + + result = convert_to_anthropic_tool_invoke(tool_calls) + + assert len(result) == 2 + # Colons replaced with underscores + assert result[0]["id"] == "sessions_history_183" + # Dots replaced with underscores + assert result[1]["id"] == "composio_NOTION_SEARCH" + # Valid IDs should pass through unchanged + valid_tool_calls = [ + { + "id": "toolu_01ABC-xyz_123", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"location": "NYC"}', + }, + } + ] + valid_result = convert_to_anthropic_tool_invoke(valid_tool_calls) + assert valid_result[0]["id"] == "toolu_01ABC-xyz_123" + + def test_convert_to_anthropic_tool_invoke_server_tool(): """ Test that server_tool_use (srvtoolu_) is reconstructed as server_tool_use. diff --git a/tests/test_litellm/llms/custom_httpx/test_aiohttp_cleanup_closed.py b/tests/test_litellm/llms/custom_httpx/test_aiohttp_cleanup_closed.py new file mode 100644 index 0000000000..c8c0e09c08 --- /dev/null +++ b/tests/test_litellm/llms/custom_httpx/test_aiohttp_cleanup_closed.py @@ -0,0 +1,31 @@ +from unittest.mock import MagicMock, patch + + +def test_create_aiohttp_transport_sets_enable_cleanup_closed_when_needed(monkeypatch): + from litellm.llms.custom_httpx import http_handler as http_handler_module + + connector_mock = MagicMock(name="connector") + session_mock = MagicMock(name="session") + monkeypatch.setattr(http_handler_module, "AIOHTTP_NEEDS_CLEANUP_CLOSED", True) + + with patch.object(http_handler_module, "TCPConnector", return_value=connector_mock) as mock_tcp_connector: + with patch.object(http_handler_module, "ClientSession", return_value=session_mock): + transport = http_handler_module.AsyncHTTPHandler._create_aiohttp_transport(shared_session=None) + transport._get_valid_client_session() + + assert mock_tcp_connector.call_args.kwargs["enable_cleanup_closed"] is True + + +def test_create_aiohttp_transport_omits_enable_cleanup_closed_when_not_needed(monkeypatch): + from litellm.llms.custom_httpx import http_handler as http_handler_module + + connector_mock = MagicMock(name="connector") + session_mock = MagicMock(name="session") + monkeypatch.setattr(http_handler_module, "AIOHTTP_NEEDS_CLEANUP_CLOSED", False) + + with patch.object(http_handler_module, "TCPConnector", return_value=connector_mock) as mock_tcp_connector: + with patch.object(http_handler_module, "ClientSession", return_value=session_mock): + transport = http_handler_module.AsyncHTTPHandler._create_aiohttp_transport(shared_session=None) + transport._get_valid_client_session() + + assert "enable_cleanup_closed" not in mock_tcp_connector.call_args.kwargs diff --git a/tests/test_litellm/llms/ollama/test_ollama_model_info.py b/tests/test_litellm/llms/ollama/test_ollama_model_info.py index 7eef15cd4d..5585e9d1e0 100644 --- a/tests/test_litellm/llms/ollama/test_ollama_model_info.py +++ b/tests/test_litellm/llms/ollama/test_ollama_model_info.py @@ -138,6 +138,88 @@ class TestOllamaModelInfo: assert models == ["ollama/llama2"] +class TestOllamaGetModelInfo: + """Tests for OllamaConfig.get_model_info() api_base threading and graceful fallback.""" + + def test_get_model_info_uses_provided_api_base(self, monkeypatch): + """When api_base is passed, get_model_info should use it instead of env var or default.""" + from litellm.llms.ollama.completion.transformation import OllamaConfig + + captured_urls = [] + + def mock_post(url, json, headers=None): + captured_urls.append(url) + resp = DummyResponse( + {"template": "{{ .System }} tools {{ .Prompt }}", "model_info": {"context_length": 4096}}, + status_code=200, + ) + return resp + + monkeypatch.setattr("litellm.module_level_client.post", mock_post) + + config = OllamaConfig() + result = config.get_model_info("llama3", api_base="http://my-remote-server:11434") + + assert captured_urls[0] == "http://my-remote-server:11434/api/show" + assert result["max_tokens"] == 4096 + + def test_get_model_info_falls_back_to_env_var(self, monkeypatch): + """When no api_base is passed, should fall back to OLLAMA_API_BASE env var.""" + from litellm.llms.ollama.completion.transformation import OllamaConfig + + captured_urls = [] + + def mock_post(url, json, headers=None): + captured_urls.append(url) + return DummyResponse({"template": "", "model_info": {}}, status_code=200) + + monkeypatch.setattr("litellm.module_level_client.post", mock_post) + monkeypatch.setenv("OLLAMA_API_BASE", "http://env-server:11434") + + config = OllamaConfig() + config.get_model_info("llama3") + + assert captured_urls[0] == "http://env-server:11434/api/show" + + def test_get_model_info_graceful_fallback_on_connection_error(self, monkeypatch): + """When the Ollama server is unreachable, should return defaults instead of raising.""" + from litellm.llms.ollama.completion.transformation import OllamaConfig + + def mock_post(url, json, headers=None): + raise ConnectionError("Connection refused") + + monkeypatch.setattr("litellm.module_level_client.post", mock_post) + monkeypatch.delenv("OLLAMA_API_BASE", raising=False) + + config = OllamaConfig() + result = config.get_model_info("llama3", api_base="http://unreachable:11434") + + assert result["key"] == "llama3" + assert result["litellm_provider"] == "ollama" + assert result["input_cost_per_token"] == 0.0 + assert result["output_cost_per_token"] == 0.0 + assert result["max_tokens"] is None + + def test_get_model_info_strips_ollama_prefix(self, monkeypatch): + """Should strip 'ollama/' or 'ollama_chat/' prefix from model name.""" + from litellm.llms.ollama.completion.transformation import OllamaConfig + + captured_json = [] + + def mock_post(url, json, headers=None): + captured_json.append(json) + return DummyResponse({"template": "", "model_info": {}}, status_code=200) + + monkeypatch.setattr("litellm.module_level_client.post", mock_post) + + config = OllamaConfig() + config.get_model_info("ollama/llama3", api_base="http://localhost:11434") + assert captured_json[0]["name"] == "llama3" + + config.get_model_info("ollama_chat/llama3", api_base="http://localhost:11434") + assert captured_json[1]["name"] == "llama3" + + class TestOllamaAuthHeaders: """Tests for Ollama authentication header handling in completion calls.""" diff --git a/tests/test_litellm/llms/perplexity/responses/test_perplexity_responses_transformation.py b/tests/test_litellm/llms/perplexity/responses/test_perplexity_responses_transformation.py new file mode 100644 index 0000000000..cdd4ef913f --- /dev/null +++ b/tests/test_litellm/llms/perplexity/responses/test_perplexity_responses_transformation.py @@ -0,0 +1,381 @@ +""" +Tests for Perplexity Responses API transformation + +Tests the PerplexityResponsesConfig class that handles Perplexity-specific +transformations for the Agent API (Responses API). + +Source: litellm/llms/perplexity/responses/transformation.py +""" + +import os +import sys + +sys.path.insert(0, os.path.abspath("../../../../..")) + +from litellm.llms.perplexity.responses.transformation import PerplexityResponsesConfig +from litellm.types.llms.openai import ResponsesAPIOptionalRequestParams +from litellm.types.utils import LlmProviders +from litellm.utils import ProviderConfigManager + + +class TestPerplexityResponsesTransformation: + """Test Perplexity Responses API configuration and transformations""" + + def test_function_tool_passthrough(self): + """Function tools with name/description/parameters are preserved""" + config = PerplexityResponsesConfig() + + params = ResponsesAPIOptionalRequestParams( + tools=[ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather", + "parameters": { + "type": "object", + "properties": { + "location": {"type": "string"}, + "unit": { + "type": "string", + "enum": ["celsius", "fahrenheit"], + }, + }, + }, + }, + } + ] + ) + + result = config.map_openai_params( + response_api_optional_params=params, + model="perplexity/openai/gpt-5.2", + drop_params=False, + ) + + assert "tools" in result + assert len(result["tools"]) == 1 + assert result["tools"][0]["type"] == "function" + assert result["tools"][0]["function"]["name"] == "get_weather" + assert ( + result["tools"][0]["function"]["description"] == "Get the current weather" + ) + assert "parameters" in result["tools"][0]["function"] + + def test_web_search_tool_passthrough(self): + """web_search tools are passed through unchanged""" + config = PerplexityResponsesConfig() + + params = ResponsesAPIOptionalRequestParams(tools=[{"type": "web_search"}]) + + result = config.map_openai_params( + response_api_optional_params=params, + model="perplexity/openai/gpt-5.2", + drop_params=False, + ) + + assert "tools" in result + assert len(result["tools"]) == 1 + assert result["tools"][0]["type"] == "web_search" + + def test_fetch_url_tool_passthrough(self): + """fetch_url tools are passed through""" + config = PerplexityResponsesConfig() + + params = ResponsesAPIOptionalRequestParams(tools=[{"type": "fetch_url"}]) + + result = config.map_openai_params( + response_api_optional_params=params, + model="perplexity/openai/gpt-5.2", + drop_params=False, + ) + + assert "tools" in result + assert len(result["tools"]) == 1 + assert result["tools"][0]["type"] == "fetch_url" + + def test_mixed_tools_function_and_web_search(self): + """Mixed function and web_search tools are transformed correctly""" + config = PerplexityResponsesConfig() + + params = ResponsesAPIOptionalRequestParams( + tools=[ + {"type": "web_search"}, + { + "type": "function", + "function": { + "name": "custom_tool", + "description": "A custom tool", + "parameters": {"type": "object"}, + }, + }, + ] + ) + + result = config.map_openai_params( + response_api_optional_params=params, + model="perplexity/openai/gpt-5.2", + drop_params=False, + ) + + assert len(result["tools"]) == 2 + assert result["tools"][0]["type"] == "web_search" + assert result["tools"][1]["type"] == "function" + assert result["tools"][1]["function"]["name"] == "custom_tool" + + def test_tool_choice_mapping(self): + """tool_choice passes through""" + config = PerplexityResponsesConfig() + + params = ResponsesAPIOptionalRequestParams( + tool_choice="required", temperature=0.7 + ) + + result = config.map_openai_params( + response_api_optional_params=params, + model="perplexity/openai/gpt-5.2", + drop_params=False, + ) + + assert result.get("tool_choice") == "required" + + def test_parallel_tool_calls(self): + """parallel_tool_calls passes through""" + config = PerplexityResponsesConfig() + + params = ResponsesAPIOptionalRequestParams( + parallel_tool_calls=True, temperature=0.7 + ) + + result = config.map_openai_params( + response_api_optional_params=params, + model="perplexity/openai/gpt-5.2", + drop_params=False, + ) + + assert result.get("parallel_tool_calls") is True + + def test_max_tool_calls_mapping(self): + """max_tool_calls passes through""" + config = PerplexityResponsesConfig() + + params = ResponsesAPIOptionalRequestParams(max_tool_calls=5, temperature=0.7) + + result = config.map_openai_params( + response_api_optional_params=params, + model="perplexity/openai/gpt-5.2", + drop_params=False, + ) + + assert result.get("max_tool_calls") == 5 + + def test_text_passthrough(self): + """text param passes through as-is (Perplexity accepts Open Responses format directly)""" + config = PerplexityResponsesConfig() + + text_value = { + "format": { + "type": "json_schema", + "name": "weather_response", + "schema": { + "type": "object", + "properties": {"temp": {"type": "number"}}, + }, + "strict": True, + } + } + + params = ResponsesAPIOptionalRequestParams( + text=text_value, + temperature=0.7, + ) + + result = config.map_openai_params( + response_api_optional_params=params, + model="perplexity/openai/gpt-5.2", + drop_params=False, + ) + + assert "text" in result + assert result["text"] == text_value + assert "response_format" not in result + + def test_previous_response_id(self): + """previous_response_id passes through""" + config = PerplexityResponsesConfig() + + params = ResponsesAPIOptionalRequestParams( + previous_response_id="resp_abc123", + temperature=0.7, + ) + + result = config.map_openai_params( + response_api_optional_params=params, + model="perplexity/openai/gpt-5.2", + drop_params=False, + ) + + assert result.get("previous_response_id") == "resp_abc123" + + def test_store_background_truncation(self): + """Lifecycle params pass through""" + config = PerplexityResponsesConfig() + + params = ResponsesAPIOptionalRequestParams( + store=True, + background=False, + truncation="auto", + temperature=0.7, + ) + + result = config.map_openai_params( + response_api_optional_params=params, + model="perplexity/openai/gpt-5.2", + drop_params=False, + ) + + assert result.get("store") is True + assert result.get("background") is False + assert result.get("truncation") == "auto" + + def test_metadata_safety_identifier_user(self): + """Metadata params pass through""" + config = PerplexityResponsesConfig() + + params = ResponsesAPIOptionalRequestParams( + metadata={"request_id": "req_123"}, + safety_identifier="safety_123", + user="user_456", + temperature=0.7, + ) + + result = config.map_openai_params( + response_api_optional_params=params, + model="perplexity/openai/gpt-5.2", + drop_params=False, + ) + + assert result.get("metadata") == {"request_id": "req_123"} + assert result.get("safety_identifier") == "safety_123" + assert result.get("user") == "user_456" + + def test_all_supported_params_declared(self): + """get_supported_openai_params returns complete list""" + config = PerplexityResponsesConfig() + supported = config.get_supported_openai_params("perplexity/openai/gpt-5.2") + + expected = [ + "max_output_tokens", + "stream", + "temperature", + "top_p", + "tools", + "reasoning", + "preset", + "instructions", + "models", + "tool_choice", + "parallel_tool_calls", + "max_tool_calls", + "text", + "previous_response_id", + "store", + "background", + "truncation", + "metadata", + "safety_identifier", + "user", + "stream_options", + "top_logprobs", + "prompt_cache_key", + "frequency_penalty", + "presence_penalty", + "service_tier", + ] + + for param in expected: + assert param in supported, f"Missing supported param: {param}" + + def test_cost_transformation(self): + """Perplexity cost dict to OpenAI float""" + config = PerplexityResponsesConfig() + + usage_data = { + "input_tokens": 100, + "output_tokens": 200, + "total_tokens": 300, + "cost": { + "currency": "USD", + "input_cost": 0.0001, + "output_cost": 0.0002, + "total_cost": 0.0003, + }, + } + + result = config._transform_usage(usage_data) + + assert result["input_tokens"] == 100 + assert result["output_tokens"] == 200 + assert result["total_tokens"] == 300 + assert result["cost"] == 0.0003 + + def test_cost_transformation_float_passthrough(self): + """Cost already float passes through""" + config = PerplexityResponsesConfig() + + usage_data = { + "input_tokens": 100, + "output_tokens": 200, + "total_tokens": 300, + "cost": 0.0005, + } + + result = config._transform_usage(usage_data) + + assert result["cost"] == 0.0005 + + def test_preset_handling(self): + """Preset model names work""" + config = PerplexityResponsesConfig() + + data = config.transform_responses_api_request( + model="preset/pro-search", + input="What is AI?", + response_api_optional_request_params={"temperature": 0.7}, + litellm_params={}, + headers={}, + ) + + assert data["preset"] == "pro-search" + assert data["input"] == "What is AI?" + assert "temperature" in data + + def test_get_complete_url(self): + """Correct endpoint URL""" + config = PerplexityResponsesConfig() + + url = config.get_complete_url(api_base=None, litellm_params={}) + assert url == "https://api.perplexity.ai/v1/responses" + + custom_url = config.get_complete_url( + api_base="https://custom.perplexity.ai", + litellm_params={}, + ) + assert custom_url == "https://custom.perplexity.ai/v1/responses" + + url_with_slash = config.get_complete_url( + api_base="https://api.perplexity.ai/", + litellm_params={}, + ) + assert url_with_slash == "https://api.perplexity.ai/v1/responses" + + def test_perplexity_provider_config_registration(self): + """Test that Perplexity provider returns PerplexityResponsesConfig""" + config = ProviderConfigManager.get_provider_responses_api_config( + model="perplexity/openai/gpt-5.2", + provider=LlmProviders.PERPLEXITY, + ) + + assert config is not None + assert isinstance(config, PerplexityResponsesConfig) + assert config.custom_llm_provider == LlmProviders.PERPLEXITY 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 041cc687b9..700ba86b10 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 @@ -1487,3 +1487,182 @@ async def test_discovery_root_includes_server_name_prefix(): assert response["scopes_supported"] == ["read", "write"] finally: global_mcp_server_manager.registry.clear() + + +@pytest.mark.asyncio +async def test_oauth_callback_redirects_with_state(): + """Test OAuth callback endpoint properly decodes state and redirects to client callback URL.""" + try: + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + callback, + ) + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + # Mock the state decoding + mock_state_data = { + "base_url": "http://localhost:3000/ui/mcp/oauth/callback", + "original_state": "test-uuid-state-123", + "code_challenge": "test_challenge", + "code_challenge_method": "S256", + "client_redirect_uri": "http://localhost:3000/ui/mcp/oauth/callback", + } + + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.decode_state_hash" + ) as mock_decode: + mock_decode.return_value = mock_state_data + + # Call callback endpoint with code and state + response = await callback( + code="test_authorization_code_12345", + state="encrypted_state_value", + ) + + # Should redirect to the client callback URL with code and original state + assert response.status_code == 302 + assert "http://localhost:3000/ui/mcp/oauth/callback" in response.headers["location"] + assert "code=test_authorization_code_12345" in response.headers["location"] + assert "state=test-uuid-state-123" in response.headers["location"] + + # Verify state was decoded + mock_decode.assert_called_once_with("encrypted_state_value") + + +@pytest.mark.asyncio +async def test_oauth_callback_handles_invalid_state(): + """Test OAuth callback returns error page when state decryption fails.""" + try: + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + callback, + ) + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + # Mock state decoding to raise an exception + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.decode_state_hash" + ) as mock_decode: + mock_decode.side_effect = Exception("Failed to decrypt state") + + # Call callback endpoint with invalid state + response = await callback( + code="test_code", + state="invalid_encrypted_state", + ) + + # Should return HTML error page + assert response.status_code == 200 + assert "Authentication incomplete" in response.body.decode() + + +@pytest.mark.asyncio +async def test_oauth_authorize_includes_scopes_from_server_config(): + """Test that authorize endpoint includes scopes from server configuration.""" + try: + from fastapi import Request + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + authorize_with_server, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp import MCPAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + # Create server with specific scopes (e.g., GitLab requires 'ai_workflows') + oauth_server = MCPServer( + server_id="gitlab_server", + name="gitlab", + server_name="gitlab", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + authorization_url="https://gitlab.com/oauth/authorize", + token_url="https://gitlab.com/oauth/token", + client_id="test_client", + scopes=["api", "read_user", "ai_workflows"], # GitLab-specific scopes + ) + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://litellm.example.com/" + mock_request.headers = {} + + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper" + ) as mock_encrypt: + mock_encrypt.return_value = "encrypted_state" + + # Call authorize without explicit scope parameter + response = await authorize_with_server( + request=mock_request, + mcp_server=oauth_server, + client_id="test_client", + redirect_uri="http://localhost:3000/callback", + state="test_state", + code_challenge="test_challenge", + code_challenge_method="S256", + response_type="code", + scope=None, # No scope in request, should use server's scopes + ) + + # Should redirect with scopes from server config + assert response.status_code in (307, 302) + redirect_url = response.headers["location"] + assert "scope=api+read_user+ai_workflows" in redirect_url or "scope=api%20read_user%20ai_workflows" in redirect_url + + +@pytest.mark.asyncio +async def test_oauth_authorize_prefers_request_scope_over_server_config(): + """Test that explicit scope parameter takes precedence over server configuration.""" + try: + from fastapi import Request + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + authorize_with_server, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp import MCPAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + oauth_server = MCPServer( + server_id="test_server", + name="test", + server_name="test", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + authorization_url="https://provider.com/oauth/authorize", + token_url="https://provider.com/oauth/token", + client_id="test_client", + scopes=["default_scope1", "default_scope2"], + ) + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://litellm.example.com/" + mock_request.headers = {} + + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper" + ) as mock_encrypt: + mock_encrypt.return_value = "encrypted_state" + + # Call authorize WITH explicit scope parameter + response = await authorize_with_server( + request=mock_request, + mcp_server=oauth_server, + client_id="test_client", + redirect_uri="http://localhost:3000/callback", + state="test_state", + code_challenge="test_challenge", + code_challenge_method="S256", + response_type="code", + scope="custom_scope1 custom_scope2", # Explicit scope should take precedence + ) + + # Should use the explicit scope, not server config + assert response.status_code in (307, 302) + redirect_url = response.headers["location"] + assert "scope=custom_scope1+custom_scope2" in redirect_url or "scope=custom_scope1%20custom_scope2" in redirect_url + assert "default_scope" not in redirect_url diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 464e523832..c105052479 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -1036,6 +1036,240 @@ class TestMCPServerManager: assert result.status == "healthy" assert result.health_check_error is None + @pytest.mark.asyncio + async def test_health_check_skips_passthrough_auth_with_authorization_header(self): + """Test that health check is skipped for servers with passthrough Authorization header""" + manager = MCPServerManager() + + # Mock server with auth_type=none and Authorization in extra_headers (passthrough auth) + server = MCPServer( + server_id="github-server", + name="github-server", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + authentication_token=None, + url="http://github-server.com", + extra_headers=["Authorization"], # Passthrough auth configured + ) + + manager.get_mcp_server_by_id = MagicMock(return_value=server) + + # _create_mcp_client should not be called (health check should be skipped) + manager._create_mcp_client = AsyncMock() + + # Perform health check + result = await manager.health_check_server("github-server") + + # Verify that client was not created (health check was skipped) + manager._create_mcp_client.assert_not_called() + + # Verify results + assert isinstance(result, LiteLLM_MCPServerTable) + assert result.server_id == "github-server" + assert result.status == "unknown" + assert result.health_check_error is None + assert result.last_health_check is not None + + @pytest.mark.asyncio + async def test_health_check_skips_passthrough_auth_with_api_key_header(self): + """Test that health check is skipped for servers with passthrough x-api-key header""" + manager = MCPServerManager() + + # Mock server with auth_type=none and x-api-key in extra_headers + server = MCPServer( + server_id="sourcegraph-server", + name="sourcegraph-server", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + authentication_token=None, + url="http://sourcegraph-server.com", + extra_headers=["x-api-key"], # Passthrough auth configured + ) + + manager.get_mcp_server_by_id = MagicMock(return_value=server) + + # _create_mcp_client should not be called + manager._create_mcp_client = AsyncMock() + + # Perform health check + result = await manager.health_check_server("sourcegraph-server") + + # Verify that client was not created (health check was skipped) + manager._create_mcp_client.assert_not_called() + + # Verify results + assert isinstance(result, LiteLLM_MCPServerTable) + assert result.server_id == "sourcegraph-server" + assert result.status == "unknown" + assert result.health_check_error is None + assert result.last_health_check is not None + + @pytest.mark.asyncio + async def test_health_check_runs_when_no_passthrough_auth(self): + """Test that health check runs normally for servers with auth_type=none but no passthrough headers""" + manager = MCPServerManager() + + # Mock server with auth_type=none but no extra_headers (no passthrough auth) + server = MCPServer( + server_id="public-server", + name="public-server", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + authentication_token=None, + url="http://public-server.com", + extra_headers=None, # No passthrough auth + ) + + manager.get_mcp_server_by_id = MagicMock(return_value=server) + + # Mock successful client + mock_client = AsyncMock() + mock_client.run_with_session = AsyncMock(return_value="ok") + manager._create_mcp_client = AsyncMock(return_value=mock_client) + + # Perform health check + result = await manager.health_check_server("public-server") + + # Verify that client WAS created (health check should run) + manager._create_mcp_client.assert_called_once() + + # Verify results + assert isinstance(result, LiteLLM_MCPServerTable) + assert result.server_id == "public-server" + assert result.status == "healthy" + assert result.health_check_error is None + assert result.last_health_check is not None + + @pytest.mark.asyncio + async def test_health_check_runs_when_extra_headers_no_auth(self): + """Test that health check runs when extra_headers exist but don't include auth headers""" + manager = MCPServerManager() + + # Mock server with extra_headers but no auth-related headers + server = MCPServer( + server_id="custom-server", + name="custom-server", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + authentication_token=None, + url="http://custom-server.com", + extra_headers=["X-Custom-Header", "X-Request-ID"], # Non-auth headers + ) + + manager.get_mcp_server_by_id = MagicMock(return_value=server) + + # Mock successful client + mock_client = AsyncMock() + mock_client.run_with_session = AsyncMock(return_value="ok") + manager._create_mcp_client = AsyncMock(return_value=mock_client) + + # Perform health check + result = await manager.health_check_server("custom-server") + + # Verify that client WAS created (health check should run) + manager._create_mcp_client.assert_called_once() + + # Verify results + assert isinstance(result, LiteLLM_MCPServerTable) + assert result.server_id == "custom-server" + assert result.status == "healthy" + assert result.health_check_error is None + + @pytest.mark.asyncio + async def test_requires_per_user_auth_property_oauth2(self): + """Test that requires_per_user_auth returns True for OAuth2 without client credentials""" + # OAuth2 without client credentials + server = MCPServer( + server_id="oauth-server", + name="oauth-server", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + url="http://oauth-server.com", + client_id=None, + client_secret=None, + token_url=None, + ) + assert server.requires_per_user_auth is True + assert server.needs_user_oauth_token is True + + @pytest.mark.asyncio + async def test_requires_per_user_auth_property_oauth2_with_client_creds(self): + """Test that requires_per_user_auth returns False for OAuth2 with client credentials""" + # OAuth2 with client credentials + server = MCPServer( + server_id="oauth-server", + name="oauth-server", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + url="http://oauth-server.com", + client_id="client-id", + client_secret="client-secret", + token_url="http://oauth-server.com/token", + ) + assert server.requires_per_user_auth is False + assert server.has_client_credentials is True + + @pytest.mark.asyncio + async def test_requires_per_user_auth_property_passthrough_auth(self): + """Test that requires_per_user_auth returns True for passthrough auth (auth_type=none + Authorization header)""" + # Passthrough auth with Authorization header + server = MCPServer( + server_id="github-server", + name="github-server", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + url="http://github-server.com", + extra_headers=["Authorization"], + ) + assert server.requires_per_user_auth is True + + # Passthrough auth with x-api-key header + server2 = MCPServer( + server_id="sourcegraph-server", + name="sourcegraph-server", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + url="http://sourcegraph-server.com", + extra_headers=["x-api-key"], + ) + assert server2.requires_per_user_auth is True + + # Passthrough auth with api-key header (case insensitive) + server3 = MCPServer( + server_id="api-server", + name="api-server", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + url="http://api-server.com", + extra_headers=["API-Key"], + ) + assert server3.requires_per_user_auth is True + + @pytest.mark.asyncio + async def test_requires_per_user_auth_property_no_passthrough(self): + """Test that requires_per_user_auth returns False when no passthrough auth is configured""" + # auth_type=none but no extra_headers + server = MCPServer( + server_id="public-server", + name="public-server", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + url="http://public-server.com", + extra_headers=None, + ) + assert server.requires_per_user_auth is False + + # auth_type=none with non-auth extra_headers + server2 = MCPServer( + server_id="custom-server", + name="custom-server", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + url="http://custom-server.com", + extra_headers=["X-Custom-Header", "X-Request-ID"], + ) + assert server2.requires_per_user_auth is False + @pytest.mark.asyncio async def test_register_openapi_tools_includes_static_headers(self, tmp_path): """Ensure OpenAPI-to-MCP tool calls include server.static_headers (Issue #19341).""" diff --git a/tests/test_litellm/proxy/test_aiohttp_cleanup_closed.py b/tests/test_litellm/proxy/test_aiohttp_cleanup_closed.py new file mode 100644 index 0000000000..f16b687a24 --- /dev/null +++ b/tests/test_litellm/proxy/test_aiohttp_cleanup_closed.py @@ -0,0 +1,34 @@ +import asyncio +from unittest.mock import MagicMock, patch + + +def test_initialize_shared_aiohttp_session_sets_enable_cleanup_closed_when_needed( + monkeypatch, +): + from litellm.proxy import proxy_server as proxy_server_module + + connector_mock = MagicMock(name="connector") + session_mock = MagicMock(name="session") + monkeypatch.setattr(proxy_server_module, "AIOHTTP_NEEDS_CLEANUP_CLOSED", True) + + with patch("aiohttp.TCPConnector", return_value=connector_mock) as mock_tcp_connector: + with patch("aiohttp.ClientSession", return_value=session_mock): + asyncio.run(proxy_server_module._initialize_shared_aiohttp_session()) + + assert mock_tcp_connector.call_args.kwargs["enable_cleanup_closed"] is True + + +def test_initialize_shared_aiohttp_session_omits_enable_cleanup_closed_when_not_needed( + monkeypatch, +): + from litellm.proxy import proxy_server as proxy_server_module + + connector_mock = MagicMock(name="connector") + session_mock = MagicMock(name="session") + monkeypatch.setattr(proxy_server_module, "AIOHTTP_NEEDS_CLEANUP_CLOSED", False) + + with patch("aiohttp.TCPConnector", return_value=connector_mock) as mock_tcp_connector: + with patch("aiohttp.ClientSession", return_value=session_mock): + asyncio.run(proxy_server_module._initialize_shared_aiohttp_session()) + + assert "enable_cleanup_closed" not in mock_tcp_connector.call_args.kwargs diff --git a/tests/test_litellm/test_video_generation.py b/tests/test_litellm/test_video_generation.py index 47a60b09af..661cdd8709 100644 --- a/tests/test_litellm/test_video_generation.py +++ b/tests/test_litellm/test_video_generation.py @@ -243,6 +243,77 @@ class TestVideoGeneration: custom_llm_provider="openai" ) + def test_video_generation_cost_with_custom_model_info(self): + """Test that custom model_info pricing is applied for video generation. + + When a deployment has custom pricing via model_info, it should be used + instead of looking up the global litellm.model_cost map. + + Related: https://github.com/BerriAI/litellm/issues/21907 + """ + model_info = { + "output_cost_per_video_per_second": 0.05, + } + cost = default_video_cost_calculator( + model="my-custom-video-model", + duration_seconds=10.0, + model_info=model_info, + ) + assert cost == 0.5 + + def test_video_generation_cost_custom_model_info_fallback_to_per_second(self): + """Test that output_cost_per_second is used as fallback when + output_cost_per_video_per_second is not set in custom model_info. + + Related: https://github.com/BerriAI/litellm/issues/21907 + """ + model_info = { + "output_cost_per_second": 0.10, + } + cost = default_video_cost_calculator( + model="my-custom-video-model", + duration_seconds=5.0, + model_info=model_info, + ) + assert cost == 0.5 + + def test_video_generation_cost_custom_pricing_through_completion_cost(self): + """Test that custom video pricing flows through completion_cost via litellm_logging_obj. + + This tests the full cost calculation path: completion_cost extracts model_info + from litellm_logging_obj.litellm_params.metadata.model_info and passes it to + the video cost calculator. + + Related: https://github.com/BerriAI/litellm/issues/21907 + """ + from litellm.cost_calculator import completion_cost + + # Create mock response with usage containing duration_seconds + mock_response = MagicMock() + mock_response.usage = MagicMock() + mock_response.usage.duration_seconds = 10.0 + type(mock_response)._hidden_params = {} + + # Create mock litellm_logging_obj with custom pricing + mock_logging_obj = MagicMock() + mock_logging_obj.litellm_params = { + "metadata": { + "model_info": { + "output_cost_per_video_per_second": 0.05, + } + } + } + + cost = completion_cost( + completion_response=mock_response, + model="openai/hunyuanvideo", + call_type="create_video", + custom_llm_provider="openai", + custom_pricing=True, + litellm_logging_obj=mock_logging_obj, + ) + assert cost == 0.5 + def test_video_generation_with_files(self): """Test video generation with file uploads.""" config = OpenAIVideoConfig() diff --git a/ui/litellm-dashboard/src/app/mcp/oauth/callback/page.tsx b/ui/litellm-dashboard/src/app/mcp/oauth/callback/page.tsx index f005eab414..73ff8c51ba 100644 --- a/ui/litellm-dashboard/src/app/mcp/oauth/callback/page.tsx +++ b/ui/litellm-dashboard/src/app/mcp/oauth/callback/page.tsx @@ -41,13 +41,16 @@ const McpOAuthCallbackContent = () => { } try { + // Store in both sessionStorage and localStorage for redundancy window.sessionStorage.setItem(RESULT_STORAGE_KEY, JSON.stringify(payload)); + window.localStorage.setItem(RESULT_STORAGE_KEY, JSON.stringify(payload)); } catch (err) { - console.error("Failed to persist OAuth callback payload", err); + // Silently ignore storage errors } - const returnUrl = window.sessionStorage.getItem(RETURN_URL_STORAGE_KEY); - console.info("[MCP OAuth callback] returnUrl", returnUrl); + // Check both sessionStorage and localStorage for return URL + const returnUrl = window.sessionStorage.getItem(RETURN_URL_STORAGE_KEY) || + window.localStorage.getItem(RETURN_URL_STORAGE_KEY); const destination = returnUrl || resolveDefaultRedirect(); window.location.replace(destination); }, [payload]); 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 fb9efffb9b..bd7e3f9034 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 @@ -129,6 +129,21 @@ const CreateMCPServer: React.FC = ({ }, onTokenReceived: (token) => { setOauthAccessToken(token?.access_token ?? null); + + if (token?.access_token) { + const credentials = { + access_token: token.access_token, + ...(token.refresh_token && { refresh_token: token.refresh_token }), + ...(token.expires_in && { expires_in: token.expires_in }), + ...(token.scope && { scope: token.scope }), + }; + + form.setFieldsValue({ credentials }); + + NotificationsManager.success( + "OAuth authorization successful! Please click 'Create MCP Server' to save the configuration." + ); + } }, onBeforeRedirect: persistCreateUiState, }); diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_columns.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_columns.tsx index 78e9c6f465..d49949e7e0 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_columns.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_columns.tsx @@ -47,7 +47,13 @@ export const mcpServerColumns = ( { accessorKey: "transport", header: "Transport", - cell: ({ getValue }) => {((getValue() as string) || "http").toUpperCase()}, + cell: ({ row }) => { + const transport = row.original.transport || "http"; + const specPath = row.original.spec_path; + // If server has spec_path, display as "OPENAPI" instead of the raw transport type + const displayTransport = specPath && transport !== "stdio" ? "OPENAPI" : transport; + return {displayTransport.toUpperCase()}; + }, }, { accessorKey: "auth_type", diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx index 88f8a737af..4209d8bf11 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx @@ -117,6 +117,21 @@ const MCPServerEdit: React.FC = ({ }, onTokenReceived: (token) => { setOauthAccessToken(token?.access_token ?? null); + + if (token?.access_token) { + const credentials = { + access_token: token.access_token, + ...(token.refresh_token && { refresh_token: token.refresh_token }), + ...(token.expires_in && { expires_in: token.expires_in }), + ...(token.scope && { scope: token.scope }), + }; + + form.setFieldsValue({ credentials }); + + NotificationsManager.success( + "OAuth authorization successful! Please click 'Update MCP Server' to save the credentials." + ); + } }, onBeforeRedirect: persistEditUiState, }); @@ -144,9 +159,9 @@ const MCPServerEdit: React.FC = ({ }, [mcpServer.env]); - // If server has spec_path and no url, show it as "openapi" transport in the UI + // If server has spec_path, show it as "openapi" transport in the UI const effectiveTransport = React.useMemo(() => { - if (mcpServer.spec_path && !mcpServer.url && mcpServer.transport !== "stdio") { + if (mcpServer.spec_path && mcpServer.transport !== "stdio") { return TRANSPORT.OPENAPI; } return mcpServer.transport; @@ -234,8 +249,13 @@ const MCPServerEdit: React.FC = ({ } }, [mcpServer]); - // Fetch tools when component mounts + // Fetch tools when component mounts or when OAuth token is received + // But only if the server has been properly saved (has a permanent server_id) useEffect(() => { + // Don't fetch if server hasn't been saved yet (no permanent server_id) + if (!mcpServer.server_id || mcpServer.server_id.trim() === "") { + return; + } fetchTools(); }, [mcpServer, accessToken, oauthAccessToken]); diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_view.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_view.tsx index 635c787f30..cb7768ab21 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_view.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_view.tsx @@ -131,7 +131,7 @@ export const MCPServerView: React.FC = ({ Transport
- {handleTransport(mcpServer.transport ?? undefined)} + {handleTransport(mcpServer.transport ?? undefined, mcpServer.spec_path ?? undefined).toUpperCase()}
@@ -171,6 +171,7 @@ export const MCPServerView: React.FC = ({ userRole={userRole} userID={userID} serverAlias={mcpServer.alias} + extraHeaders={mcpServer.extra_headers} /> @@ -220,7 +221,7 @@ export const MCPServerView: React.FC = ({
Transport -
{handleTransport(mcpServer.transport)}
+
{handleTransport(mcpServer.transport, mcpServer.spec_path).toUpperCase()}
Extra Headers diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_tools.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_tools.tsx index 1cc505ec02..3a572dec89 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_tools.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_tools.tsx @@ -1,12 +1,12 @@ import React, { useState } from "react"; import { useQuery, useMutation } from "@tanstack/react-query"; import { ToolTestPanel } from "./ToolTestPanel"; -import { MCPTool, MCPToolsViewerProps, MCPContent, CallMCPToolResponse } from "./types"; +import { MCPTool, MCPToolsViewerProps, MCPContent, CallMCPToolResponse, AUTH_TYPE } from "./types"; import { listMCPTools, callMCPTool } from "../networking"; import { Card, Title, Text } from "@tremor/react"; -import { RobotOutlined, ToolOutlined, SearchOutlined } from "@ant-design/icons"; -import { Input } from "antd"; +import { RobotOutlined, ToolOutlined, SearchOutlined, LockOutlined, KeyOutlined } from "@ant-design/icons"; +import { Input, Alert, Button as AntdButton } from "antd"; const MCPToolsViewer = ({ serverId, @@ -14,23 +14,50 @@ const MCPToolsViewer = ({ auth_type, userRole, userID, - serverAlias, // Add serverAlias prop + serverAlias, + extraHeaders, }: MCPToolsViewerProps) => { const [selectedTool, setSelectedTool] = useState(null); const [toolResult, setToolResult] = useState(null); const [toolError, setToolError] = useState(null); const [toolSearchTerm, setToolSearchTerm] = useState(""); + + // State for passthrough headers + const [passthroughHeaders, setPassthroughHeaders] = useState>({}); + const [showHeaderInput, setShowHeaderInput] = useState(false); + + // Check if this server has extra headers configured + const hasExtraHeaders = extraHeaders && extraHeaders.length > 0; + + // Build custom headers for MCP server requests + const buildCustomHeaders = () => { + if (!serverAlias || !hasExtraHeaders) return undefined; + + const customHeaders: Record = {}; + + // Add passthrough headers with server-specific prefix + Object.entries(passthroughHeaders).forEach(([headerName, headerValue]) => { + if (headerValue && headerValue.trim()) { + // Format: x-mcp-{alias}-{header_name} + const mcpHeaderName = `x-mcp-${serverAlias}-${headerName.toLowerCase()}`; + customHeaders[mcpHeaderName] = headerValue; + } + }); + + return Object.keys(customHeaders).length > 0 ? customHeaders : undefined; + }; // Query to fetch MCP tools const { data: mcpToolsResponse, isLoading: isLoadingTools, error: mcpToolsError, + refetch: refetchTools, } = useQuery({ - queryKey: ["mcpTools", serverId], + queryKey: ["mcpTools", serverId, passthroughHeaders], queryFn: () => { if (!accessToken) throw new Error("Access Token required"); - return listMCPTools(accessToken, serverId); + return listMCPTools(accessToken, serverId, buildCustomHeaders()); }, enabled: !!accessToken, staleTime: 30000, // Consider data fresh for 30 seconds @@ -42,7 +69,13 @@ const MCPToolsViewer = ({ if (!accessToken) throw new Error("Access Token required"); try { - const result: CallMCPToolResponse = await callMCPTool(accessToken, serverId, args.tool.name, args.arguments); + const result: CallMCPToolResponse = await callMCPTool( + accessToken, + serverId, + args.tool.name, + args.arguments, + { customHeaders: buildCustomHeaders() } + ); return result; } catch (error) { throw error; @@ -79,6 +112,80 @@ const MCPToolsViewer = ({ MCP Tools
+ {/* Extra Headers Input Section */} + {hasExtraHeaders && ( +
+
+
+ + + Additional Headers + +
+ setShowHeaderInput(!showHeaderInput)} + className="text-blue-700 p-0 h-auto" + > + {showHeaderInput ? "Hide" : "Configure"} + +
+ + {!showHeaderInput && Object.keys(passthroughHeaders).length === 0 && ( + + This server requires additional headers. Click "Configure" to provide values. + + )} + + {showHeaderInput && ( +
+ {extraHeaders?.map((headerName) => ( +
+ + { + setPassthroughHeaders({ + ...passthroughHeaders, + [headerName]: e.target.value, + }); + }} + prefix={} + className="rounded" + /> +
+ ))} + { + refetchTools(); + setShowHeaderInput(false); + }} + disabled={Object.values(passthroughHeaders).every(v => !v || !v.trim())} + className="w-full mt-2" + > + Load Tools + +
+ )} + + {!showHeaderInput && Object.keys(passthroughHeaders).length > 0 && ( +
+ + + {Object.keys(passthroughHeaders).length} header(s) configured + +
+ )} +
+ )} + {/* Tool Selection - Show tools first */}
diff --git a/ui/litellm-dashboard/src/components/mcp_tools/types.tsx b/ui/litellm-dashboard/src/components/mcp_tools/types.tsx index 6856322b53..1fa447c0e6 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/types.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/types.tsx @@ -25,11 +25,16 @@ export const TRANSPORT = { OPENAPI: "openapi", }; -export const handleTransport = (transport?: string | null): string => { +export const handleTransport = (transport?: string | null, specPath?: string | null): string => { if (transport === null || transport === undefined) { return TRANSPORT.SSE; } + // If server has spec_path, display as "openapi" instead of the raw transport type + if (specPath && transport !== TRANSPORT.STDIO) { + return TRANSPORT.OPENAPI; + } + return transport; }; @@ -132,6 +137,7 @@ export interface MCPToolsViewerProps { userRole: string | null; userID: string | null; serverAlias?: string | null; + extraHeaders?: string[] | null; } export interface MCPServer { diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 29bb7e4352..6ffd744cce 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -7147,7 +7147,11 @@ export const testSearchToolConnection = async (accessToken: string, litellmParam } }; -export const listMCPTools = async (accessToken: string, serverId: string) => { +export const listMCPTools = async ( + accessToken: string, + serverId: string, + customHeaders?: Record +) => { try { // Construct base URL let url = proxyBaseUrl @@ -7159,6 +7163,7 @@ export const listMCPTools = async (accessToken: string, serverId: string) => { const headers: Record = { [globalLitellmHeaderName]: `Bearer ${accessToken}`, "Content-Type": "application/json", + ...customHeaders, // Merge custom headers for passthrough auth }; const response = await fetch(url, { @@ -7194,6 +7199,7 @@ export const listMCPTools = async (accessToken: string, serverId: string) => { export interface CallMCPToolOptions { guardrails?: string[]; + customHeaders?: Record; } export const callMCPTool = async ( @@ -7212,6 +7218,7 @@ export const callMCPTool = async ( const headers: Record = { [globalLitellmHeaderName]: `Bearer ${accessToken}`, "Content-Type": "application/json", + ...(options?.customHeaders || {}), // Merge custom headers for passthrough auth }; const body: Record = { diff --git a/ui/litellm-dashboard/src/hooks/useMcpOAuthFlow.tsx b/ui/litellm-dashboard/src/hooks/useMcpOAuthFlow.tsx index a62d8baa75..f914e42f04 100644 --- a/ui/litellm-dashboard/src/hooks/useMcpOAuthFlow.tsx +++ b/ui/litellm-dashboard/src/hooks/useMcpOAuthFlow.tsx @@ -1,6 +1,6 @@ "use client"; -import { useCallback, useEffect, useState } from "react"; +import { useCallback, useEffect, useRef, useState } from "react"; import NotificationsManager from "@/components/molecules/notifications_manager"; import { buildMcpOAuthAuthorizeUrl, @@ -61,6 +61,7 @@ export const useMcpOAuthFlow = ({ const [status, setStatus] = useState("idle"); const [error, setError] = useState(null); const [tokenResponse, setTokenResponse] = useState | null>(null); + const processingRef = useRef(false); const FLOW_STATE_KEY = "litellm-mcp-oauth-flow-state"; const RESULT_KEY = "litellm-mcp-oauth-result"; @@ -75,6 +76,28 @@ export const useMcpOAuthFlow = ({ redirectUri: string; }; + const setStorageItem = (key: string, value: string) => { + if (typeof window === "undefined") return; + try { + // Store in both sessionStorage and localStorage for redundancy + window.sessionStorage.setItem(key, value); + window.localStorage.setItem(key, value); + } catch (err) { + console.warn(`Failed to set storage item ${key}`, err); + } + }; + + const getStorageItem = (key: string): string | null => { + if (typeof window === "undefined") return null; + try { + // Try sessionStorage first, fall back to localStorage + return window.sessionStorage.getItem(key) || window.localStorage.getItem(key); + } catch (err) { + console.warn(`Failed to get storage item ${key}`, err); + return null; + } + }; + const clearStoredFlow = () => { if (typeof window === "undefined") { return; @@ -83,6 +106,9 @@ export const useMcpOAuthFlow = ({ window.sessionStorage.removeItem(FLOW_STATE_KEY); window.sessionStorage.removeItem(RESULT_KEY); window.sessionStorage.removeItem(RETURN_URL_KEY); + window.localStorage.removeItem(FLOW_STATE_KEY); + window.localStorage.removeItem(RESULT_KEY); + window.localStorage.removeItem(RETURN_URL_KEY); } catch (err) { console.warn("Failed to clear OAuth storage", err); } @@ -187,10 +213,9 @@ export const useMcpOAuthFlow = ({ } try { - window.sessionStorage.setItem(FLOW_STATE_KEY, JSON.stringify(flowState)); - window.sessionStorage.setItem(RETURN_URL_KEY, window.location.href); + setStorageItem(FLOW_STATE_KEY, JSON.stringify(flowState)); + setStorageItem(RETURN_URL_KEY, window.location.href); } catch (storageErr) { - console.error("Unable to persist OAuth state", storageErr); throw new Error("Unable to access browser storage for OAuth. Please enable storage and retry."); } @@ -209,19 +234,28 @@ export const useMcpOAuthFlow = ({ return; } + // Prevent duplicate processing + if (processingRef.current) { + return; + } + let payload: Record | null = null; let flowState: StoredFlowState | null = null; try { - const storedPayload = window.sessionStorage.getItem(RESULT_KEY); + const storedPayload = getStorageItem(RESULT_KEY); if (!storedPayload) { return; } + + // Mark as processing + processingRef.current = true; payload = JSON.parse(storedPayload); - flowState = JSON.parse(window.sessionStorage.getItem(FLOW_STATE_KEY) || "null"); + const storedFlowState = getStorageItem(FLOW_STATE_KEY); + flowState = storedFlowState ? JSON.parse(storedFlowState) : null; } catch (err) { - console.error("Failed to read OAuth session state", err); clearStoredFlow(); + processingRef.current = false; setError("Failed to resume OAuth flow. Please retry."); setStatus("error"); NotificationsManager.error("Failed to resume OAuth flow. Please retry."); @@ -229,14 +263,26 @@ export const useMcpOAuthFlow = ({ } if (!payload) { + processingRef.current = false; return; } - window.sessionStorage.removeItem(RESULT_KEY); + // Clear the result key after reading it + if (typeof window !== "undefined") { + try { + window.sessionStorage.removeItem(RESULT_KEY); + window.localStorage.removeItem(RESULT_KEY); + } catch (err) { + // Silently ignore storage errors + } + } try { if (!flowState || !flowState.state || !flowState.codeVerifier || !flowState.serverId) { - throw new Error("Missing OAuth session state. Please retry."); + throw new Error( + "OAuth session state was lost. This can happen if you have strict browser privacy settings. " + + "Please try again and ensure cookies/storage is enabled." + ); } if (!payload.state || payload.state !== flowState.state) { throw new Error("OAuth state mismatch. Please retry."); @@ -264,31 +310,21 @@ export const useMcpOAuthFlow = ({ setError(null); NotificationsManager.success("OAuth token retrieved successfully"); } catch (err) { - console.error("OAuth flow failed", err); const message = err instanceof Error ? err.message : String(err); setError(message); setStatus("error"); NotificationsManager.error(message); } finally { clearStoredFlow(); + // Reset processing flag after a delay to allow UI updates + setTimeout(() => { + processingRef.current = false; + }, 1000); } }, [onTokenReceived]); useEffect(() => { - let cancelled = false; - - const maybeResume = async () => { - if (cancelled) { - return; - } - await resumeOAuthFlow(); - }; - - maybeResume(); - - return () => { - cancelled = true; - }; + resumeOAuthFlow(); }, [resumeOAuthFlow]); return {