Merge remote-tracking branch 'origin' into litellm_yj_march_18_2026

CLAUDE.md

@@ -140,6 +140,11 @@ LiteLLM is a unified interface for 100+ LLM providers with two main components:
 - **Check index coverage.** For new or modified queries, check `schema.prisma` for a supporting index. Prefer extending an existing index (e.g. `@@index([a])` → `@@index([a, b])`) over adding a new one, unless it's a `@@unique`. Only add indexes for large/frequent queries.
 - **Keep schema files in sync.** Apply schema changes to all `schema.prisma` copies (`schema.prisma`, `litellm/proxy/`, `litellm-proxy-extras/`, `litellm-js/spend-logs/` for SpendLogs) with a migration under `litellm-proxy-extras/litellm_proxy_extras/migrations/`.
 
+### Setup Wizard (`litellm/setup_wizard.py`)
+- The wizard is implemented as a single `SetupWizard` class with `@staticmethod` methods — keep it that way. No module-level functions except `run_setup_wizard()` (the public entrypoint) and pure helpers (color, ANSI).
+- Use `litellm.utils.check_valid_key(model, api_key)` for credential validation — never roll a custom completion call.
+- Do not hardcode provider env-key names or model lists that already exist in the codebase. Add a `test_model` field to each provider entry to drive `check_valid_key`; set it to `None` for providers that can't be validated with a single API key (Azure, Bedrock, Ollama).
+
 ### Enterprise Features
 - Enterprise-specific code in `enterprise/` directory
 - Optional features enabled via environment variables

docs/my-website/docs/observability/langfuse_otel_integration.md

@@ -83,6 +83,9 @@ os.environ["LANGFUSE_OTEL_HOST"] = "https://cloud.langfuse.com"  # EU region
 # Or use self-hosted instance
 # os.environ["LANGFUSE_OTEL_HOST"] = "https://my-langfuse.company.com"
 
+# Optional: Ignore otel context propagation to prevent parent-child relationships with spans from other providers
+# os.environ["OTEL_IGNORE_CONTEXT_PROPAGATION"] = "true"
+
 litellm.callbacks = ["langfuse_otel"]
 ```
 
@@ -124,6 +127,9 @@ export LANGFUSE_PUBLIC_KEY="pk-lf-..."
 export LANGFUSE_SECRET_KEY="sk-lf-..."
 export LANGFUSE_OTEL_HOST="https://us.cloud.langfuse.com"  # Default US region
 # export LANGFUSE_OTEL_HOST="https://otel.my-langfuse.company.com"  # custom OTEL endpoint
+
+# Optional: Ignore otel context propagation to prevent parent-child relationships with spans from other providers
+# export OTEL_IGNORE_CONTEXT_PROPAGATION="true"
 ```
 
 2. Setup config.yaml

docs/my-website/docs/proxy/docker_quick_start.md

@@ -5,11 +5,76 @@ import Image from '@theme/IdealImage';
 # Getting Started Tutorial
 
 End-to-End tutorial for LiteLLM Proxy to:
-- Add an Azure OpenAI model 
-- Make a successful /chat/completion call 
-- Generate a virtual key 
-- Set RPM limit on virtual key 
+- Add an Azure OpenAI model
+- Make a successful /chat/completion call
+- Generate a virtual key
+- Set RPM limit on virtual key
 
+## Quick Install (Recommended for local / beginners)
+
+New to LiteLLM? This is the easiest way to get started locally. One command installs LiteLLM and walks you through setup interactively — no config files to write by hand.
+
+### 1. Install
+
+```bash
+curl -fsSL https://raw.githubusercontent.com/BerriAI/litellm/main/scripts/install.sh | sh
+```
+
+This detects your OS, installs `litellm[proxy]`, and drops you straight into the setup wizard.
+
+### 2. Follow the wizard
+
+```
+$ litellm --setup
+
+  Welcome to LiteLLM
+
+  Choose your LLM providers
+  ○ 1. OpenAI        GPT-4o, GPT-4o-mini, o1
+  ○ 2. Anthropic     Claude Opus, Sonnet, Haiku
+  ○ 3. Azure OpenAI  GPT-4o via Azure
+  ○ 4. Google Gemini Gemini 2.0 Flash, 1.5 Pro
+  ○ 5. AWS Bedrock   Claude, Llama via AWS
+  ○ 6. Ollama        Local models
+
+  ❯ Provider(s): 1,2
+
+  ❯ OpenAI API key: sk-...
+  ❯ Anthropic API key: sk-ant-...
+
+  ❯ Port [4000]:
+  ❯ Master key [auto-generate]:
+
+  ✔ Config saved → ./litellm_config.yaml
+
+  ❯ Start the proxy now? (Y/n):
+```
+
+The wizard walks you through:
+1. Pick your LLM providers (OpenAI, Anthropic, Azure, Bedrock, Gemini, Ollama)
+2. Enter API keys for each provider
+3. Set a port and master key (or accept the defaults)
+4. Config is saved to `./litellm_config.yaml` and the proxy starts immediately
+
+### 3. Make a call
+
+Your proxy is running on `http://0.0.0.0:4000`. Test it:
+
+```bash
+curl -X POST 'http://0.0.0.0:4000/chat/completions' \
+-H 'Content-Type: application/json' \
+-H 'Authorization: Bearer <your-master-key>' \
+-d '{
+    "model": "gpt-4o",
+    "messages": [{"role": "user", "content": "Hello!"}]
+}'
+```
+
+:::tip Already have pip installed?
+You can skip the curl install and run `litellm --setup` directly after `pip install 'litellm[proxy]'`.
+:::
+
+---
 
 ## Pre-Requisites 
 

docs/my-website/sidebars.js

@@ -672,6 +672,7 @@ const sidebars = {
             "mcp_control",
             "mcp_cost",
             "mcp_guardrail",
+            "mcp_zero_trust",
             "mcp_troubleshoot",
           ]
         },

litellm-proxy-extras/litellm_proxy_extras/migrations/20260318140652_add_index_to_team_table/migration.sql

@@ -0,0 +1,9 @@
+-- CreateIndex
+CREATE INDEX "LiteLLM_TeamTable_organization_id_idx" ON "LiteLLM_TeamTable"("organization_id");
+
+-- CreateIndex
+CREATE INDEX "LiteLLM_TeamTable_team_alias_idx" ON "LiteLLM_TeamTable"("team_alias");
+
+-- CreateIndex
+CREATE INDEX "LiteLLM_TeamTable_created_at_idx" ON "LiteLLM_TeamTable"("created_at");
+

litellm-proxy-extras/litellm_proxy_extras/schema.prisma

@@ -146,6 +146,10 @@ model LiteLLM_TeamTable {
     litellm_model_table LiteLLM_ModelTable? @relation(fields: [model_id], references: [id])
     object_permission LiteLLM_ObjectPermissionTable?   @relation(fields: [object_permission_id], references: [object_permission_id])
     projects LiteLLM_ProjectTable[]
+
+    @@index([organization_id])
+    @@index([team_alias])
+    @@index([created_at])
 }
 
 // Projects sit between teams and keys for use-case management

litellm-proxy-extras/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "litellm-proxy-extras"
-version = "0.4.57"
+version = "0.4.58"
 description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
 authors = ["BerriAI"]
 readme = "README.md"
@@ -22,7 +22,7 @@ requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"
 
 [tool.commitizen]
-version = "0.4.57"
+version = "0.4.58"
 version_files = [
     "pyproject.toml:version",
     "../requirements.txt:litellm-proxy-extras==",

litellm/__init__.py

@@ -1465,9 +1465,15 @@ if TYPE_CHECKING:
     from .llms.petals.completion.transformation import PetalsConfig as PetalsConfig
     from .llms.ollama.chat.transformation import OllamaChatConfig as OllamaChatConfig
     from .llms.ollama.completion.transformation import OllamaConfig as OllamaConfig
-    from .llms.sagemaker.completion.transformation import SagemakerConfig as SagemakerConfig
-    from .llms.sagemaker.chat.transformation import SagemakerChatConfig as SagemakerChatConfig
-    from .llms.sagemaker.nova.transformation import SagemakerNovaConfig as SagemakerNovaConfig
+    from .llms.sagemaker.completion.transformation import (
+        SagemakerConfig as SagemakerConfig,
+    )
+    from .llms.sagemaker.chat.transformation import (
+        SagemakerChatConfig as SagemakerChatConfig,
+    )
+    from .llms.sagemaker.nova.transformation import (
+        SagemakerNovaConfig as SagemakerNovaConfig,
+    )
     from .llms.cohere.chat.transformation import CohereChatConfig as CohereChatConfig
     from .llms.anthropic.experimental_pass_through.messages.transformation import (
         AnthropicMessagesConfig as AnthropicMessagesConfig,

litellm/_logging.py

@@ -17,7 +17,9 @@ if set_verbose is True:
         "`litellm.set_verbose` is deprecated. Please set `os.environ['LITELLM_LOG'] = 'DEBUG'` for debug logs."
     )
 
-_ENABLE_SECRET_REDACTION = os.getenv("LITELLM_DISABLE_REDACT_SECRETS", "").lower() != "true"
+_ENABLE_SECRET_REDACTION = (
+    os.getenv("LITELLM_DISABLE_REDACT_SECRETS", "").lower() != "true"
+)
 
 _REDACTED = "REDACTED"
 
@@ -199,7 +201,9 @@ class JsonFormatter(Formatter):
                 json_record[key] = value
 
         if record.exc_info:
-            json_record["stacktrace"] = record.exc_text or self.formatException(record.exc_info)
+            json_record["stacktrace"] = record.exc_text or self.formatException(
+                record.exc_info
+            )
 
         return safe_dumps(json_record)
 

litellm/cost_calculator.py

@@ -1189,7 +1189,9 @@ def completion_cost(  # noqa: PLR0915
                         and _usage["prompt_tokens_details"] != {}
                         and _usage["prompt_tokens_details"]
                     ):
-                        prompt_tokens_details = _usage.get("prompt_tokens_details") or {}
+                        prompt_tokens_details = (
+                            _usage.get("prompt_tokens_details") or {}
+                        )
                         cache_read_input_tokens = prompt_tokens_details.get(
                             "cached_tokens", 0
                         )
@@ -1515,7 +1517,9 @@ def completion_cost(  # noqa: PLR0915
                 if custom_llm_provider == "azure_ai":
                     model_for_additional_costs = request_model_for_cost
                     if completion_response is not None:
-                        hidden_params = getattr(completion_response, "_hidden_params", None) or {}
+                        hidden_params = (
+                            getattr(completion_response, "_hidden_params", None) or {}
+                        )
                         hidden_model = hidden_params.get("model") or hidden_params.get(
                             "litellm_model_name"
                         )

litellm/integrations/focus/destinations/factory.py

@@ -59,17 +59,14 @@ class FocusDestinationFactory:
             return {k: v for k, v in resolved.items() if v is not None}
         if provider == "vantage":
             resolved = {
-                "api_key": overrides.get("api_key")
-                or os.getenv("VANTAGE_API_KEY"),
+                "api_key": overrides.get("api_key") or os.getenv("VANTAGE_API_KEY"),
                 "integration_token": overrides.get("integration_token")
                 or os.getenv("VANTAGE_INTEGRATION_TOKEN"),
                 "base_url": overrides.get("base_url")
                 or os.getenv("VANTAGE_BASE_URL", "https://api.vantage.sh"),
             }
             if not resolved.get("api_key"):
-                raise ValueError(
-                    "VANTAGE_API_KEY must be provided for Vantage exports"
-                )
+                raise ValueError("VANTAGE_API_KEY must be provided for Vantage exports")
             if not resolved.get("integration_token"):
                 raise ValueError(
                     "VANTAGE_INTEGRATION_TOKEN must be provided for Vantage exports"

litellm/integrations/langfuse/langfuse_prompt_management.py

@@ -340,9 +340,9 @@ class LangfusePromptManagement(LangFuseLogger, PromptManagementBase, CustomLogge
             )
             status_message = str(kwargs.get("exception", "Unknown error"))
             if standard_logging_object is not None:
-                status_message = standard_logging_object.get(
-                    "error_str", None
-                ) or status_message
+                status_message = (
+                    standard_logging_object.get("error_str", None) or status_message
+                )
             langfuse_logger_to_use.log_event_on_langfuse(
                 start_time=start_time,
                 end_time=end_time,

litellm/integrations/opentelemetry.py

@@ -11,7 +11,7 @@ from litellm.integrations._types.open_inference import (
 )
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
-from litellm.secret_managers.main import get_secret_bool
+from litellm.secret_managers.main import get_secret_bool, str_to_bool
 from litellm.types.services import ServiceLoggerPayload
 from litellm.types.utils import (
     ChatCompletionMessageToolCall,
@@ -68,6 +68,7 @@ class OpenTelemetryConfig:
     service_name: Optional[str] = None
     deployment_environment: Optional[str] = None
     model_id: Optional[str] = None
+    ignore_context_propagation: Optional[bool] = None
 
     def __post_init__(self) -> None:
         # If endpoint is specified but exporter is still the default "console",
@@ -89,6 +90,10 @@ class OpenTelemetryConfig:
             )
         if not self.model_id:
             self.model_id = os.getenv("OTEL_MODEL_ID", self.service_name)
+        if self.ignore_context_propagation is None:
+            self.ignore_context_propagation = str_to_bool(
+                os.getenv("OTEL_IGNORE_CONTEXT_PROPAGATION")
+            )
 
     @classmethod
     def from_env(cls):
@@ -710,12 +715,7 @@ class OpenTelemetry(CustomLogger):
         )
         ctx, parent_span = self._get_span_context(kwargs)
 
-        # CRITICAL FIX: For langfuse_otel, ALWAYS create primary spans
-        # Don't use parent spans from other providers as they cause trace corruption
-        is_langfuse_otel = (
-            hasattr(self, "callback_name") and self.callback_name == "langfuse_otel"
-        )
-        if is_langfuse_otel:
+        if self.config.ignore_context_propagation:
             parent_span = None  # Ignore parent spans from other providers
             ctx = None
 
@@ -1256,12 +1256,7 @@ class OpenTelemetry(CustomLogger):
         )
         _parent_context, parent_otel_span = self._get_span_context(kwargs)
 
-        # CRITICAL FIX: For langfuse_otel, ALWAYS create primary spans
-        # Don't use parent spans from other providers as they cause trace corruption
-        is_langfuse_otel = (
-            hasattr(self, "callback_name") and self.callback_name == "langfuse_otel"
-        )
-        if is_langfuse_otel:
+        if self.config.ignore_context_propagation:
             parent_otel_span = None  # Ignore parent spans from other providers
             _parent_context = None
 

litellm/integrations/vantage/vantage_logger.py

@@ -83,7 +83,9 @@ class VantageLogger(FocusLogger):
 
         verbose_logger.debug(
             "VantageLogger initialized (integration_token=%s)",
-            resolved_token[:4] + "***" if resolved_token and len(resolved_token) > 4 else "***",
+            resolved_token[:4] + "***"
+            if resolved_token and len(resolved_token) > 4
+            else "***",
         )
 
     async def initialize_focus_export_job(self) -> None:
@@ -128,9 +130,7 @@ class VantageLogger(FocusLogger):
             callback_type=VantageLogger
         )
         if not vantage_loggers:
-            verbose_logger.debug(
-                "No Vantage logger registered; skipping scheduler"
-            )
+            verbose_logger.debug("No Vantage logger registered; skipping scheduler")
             return
 
         vantage_logger = cast(VantageLogger, vantage_loggers[0])

litellm/litellm_core_utils/default_encoding.py

@@ -26,7 +26,9 @@ if custom_cache_dir:
 else:
     cache_dir = filename
 
-os.environ["TIKTOKEN_CACHE_DIR"] = cache_dir  # use local copy of tiktoken b/c of - https://github.com/BerriAI/litellm/issues/1071
+os.environ[
+    "TIKTOKEN_CACHE_DIR"
+] = cache_dir  # use local copy of tiktoken b/c of - https://github.com/BerriAI/litellm/issues/1071
 
 import tiktoken
 import time
@@ -48,4 +50,3 @@ for attempt in range(_max_retries):
         # Exponential backoff with jitter to reduce collision probability
         delay = _retry_delay * (2**attempt) + random.uniform(0, 0.1)
         time.sleep(delay)
-

litellm/litellm_core_utils/litellm_logging.py

@@ -352,9 +352,9 @@ class Logging(LiteLLMLoggingBaseClass):
         )
         self.function_id = function_id
         self.streaming_chunks: List[Any] = []  # for generating complete stream response
-        self.sync_streaming_chunks: List[Any] = (
-            []
-        )  # for generating complete stream response
+        self.sync_streaming_chunks: List[
+            Any
+        ] = []  # for generating complete stream response
         self.log_raw_request_response = log_raw_request_response
 
         # Initialize dynamic callbacks
@@ -782,9 +782,9 @@ class Logging(LiteLLMLoggingBaseClass):
                         prompt_spec=prompt_spec,
                         dynamic_callback_params=dynamic_callback_params,
                     ):
-                        self.model_call_details["prompt_integration"] = (
-                            logger.__class__.__name__
-                        )
+                        self.model_call_details[
+                            "prompt_integration"
+                        ] = logger.__class__.__name__
                         return logger
                 except Exception:
                     # If check fails, continue to next logger
@@ -852,9 +852,9 @@ class Logging(LiteLLMLoggingBaseClass):
         if anthropic_cache_control_logger := AnthropicCacheControlHook.get_custom_logger_for_anthropic_cache_control_hook(
             non_default_params
         ):
-            self.model_call_details["prompt_integration"] = (
-                anthropic_cache_control_logger.__class__.__name__
-            )
+            self.model_call_details[
+                "prompt_integration"
+            ] = anthropic_cache_control_logger.__class__.__name__
             return anthropic_cache_control_logger
 
         #########################################################
@@ -866,9 +866,9 @@ class Logging(LiteLLMLoggingBaseClass):
                 internal_usage_cache=None,
                 llm_router=None,
             )
-            self.model_call_details["prompt_integration"] = (
-                vector_store_custom_logger.__class__.__name__
-            )
+            self.model_call_details[
+                "prompt_integration"
+            ] = vector_store_custom_logger.__class__.__name__
             # Add to global callbacks so post-call hooks are invoked
             if (
                 vector_store_custom_logger
@@ -928,9 +928,9 @@ class Logging(LiteLLMLoggingBaseClass):
             model
         ):  # if model name was changes pre-call, overwrite the initial model call name with the new one
             self.model_call_details["model"] = model
-        self.model_call_details["litellm_params"]["api_base"] = (
-            self._get_masked_api_base(additional_args.get("api_base", ""))
-        )
+        self.model_call_details["litellm_params"][
+            "api_base"
+        ] = self._get_masked_api_base(additional_args.get("api_base", ""))
 
     def pre_call(self, input, api_key, model=None, additional_args={}):  # noqa: PLR0915
         # Log the exact input to the LLM API
@@ -959,10 +959,10 @@ class Logging(LiteLLMLoggingBaseClass):
                 try:
                     # [Non-blocking Extra Debug Information in metadata]
                     if turn_off_message_logging is True:
-                        _metadata["raw_request"] = (
-                            "redacted by litellm. \
+                        _metadata[
+                            "raw_request"
+                        ] = "redacted by litellm. \
                             'litellm.turn_off_message_logging=True'"
-                        )
                     else:
                         curl_command = self._get_request_curl_command(
                             api_base=additional_args.get("api_base", ""),
@@ -973,34 +973,34 @@ class Logging(LiteLLMLoggingBaseClass):
 
                         _metadata["raw_request"] = str(curl_command)
                         # split up, so it's easier to parse in the UI
-                        self.model_call_details["raw_request_typed_dict"] = (
-                            RawRequestTypedDict(
-                                raw_request_api_base=str(
-                                    additional_args.get("api_base") or ""
-                                ),
-                                raw_request_body=self._get_raw_request_body(
-                                    additional_args.get("complete_input_dict", {})
-                                ),
-                                # NOTE: setting ignore_sensitive_headers to True will cause
-                                # the Authorization header to be leaked when calls to the health
-                                # endpoint are made and fail.
-                                raw_request_headers=self._get_masked_headers(
-                                    additional_args.get("headers", {}) or {},
-                                ),
-                                error=None,
-                            )
+                        self.model_call_details[
+                            "raw_request_typed_dict"
+                        ] = RawRequestTypedDict(
+                            raw_request_api_base=str(
+                                additional_args.get("api_base") or ""
+                            ),
+                            raw_request_body=self._get_raw_request_body(
+                                additional_args.get("complete_input_dict", {})
+                            ),
+                            # NOTE: setting ignore_sensitive_headers to True will cause
+                            # the Authorization header to be leaked when calls to the health
+                            # endpoint are made and fail.
+                            raw_request_headers=self._get_masked_headers(
+                                additional_args.get("headers", {}) or {},
+                            ),
+                            error=None,
                         )
                 except Exception as e:
-                    self.model_call_details["raw_request_typed_dict"] = (
-                        RawRequestTypedDict(
-                            error=str(e),
-                        )
+                    self.model_call_details[
+                        "raw_request_typed_dict"
+                    ] = RawRequestTypedDict(
+                        error=str(e),
                     )
-                    _metadata["raw_request"] = (
-                        "Unable to Log \
+                    _metadata[
+                        "raw_request"
+                    ] = "Unable to Log \
                         raw request: {}".format(
-                            str(e)
-                        )
+                        str(e)
                     )
             if getattr(self, "logger_fn", None) and callable(self.logger_fn):
                 try:
@@ -1301,13 +1301,13 @@ class Logging(LiteLLMLoggingBaseClass):
         for callback in callbacks:
             try:
                 if isinstance(callback, CustomLogger):
-                    response: Optional[MCPPostCallResponseObject] = (
-                        await callback.async_post_mcp_tool_call_hook(
-                            kwargs=kwargs,
-                            response_obj=post_mcp_tool_call_response_obj,
-                            start_time=start_time,
-                            end_time=end_time,
-                        )
+                    response: Optional[
+                        MCPPostCallResponseObject
+                    ] = await callback.async_post_mcp_tool_call_hook(
+                        kwargs=kwargs,
+                        response_obj=post_mcp_tool_call_response_obj,
+                        start_time=start_time,
+                        end_time=end_time,
                     )
                     ######################################################################
                     # if any of the callbacks modify the response, use the modified response
@@ -1502,9 +1502,9 @@ class Logging(LiteLLMLoggingBaseClass):
             verbose_logger.debug(
                 f"response_cost_failure_debug_information: {debug_info}"
             )
-            self.model_call_details["response_cost_failure_debug_information"] = (
-                debug_info
-            )
+            self.model_call_details[
+                "response_cost_failure_debug_information"
+            ] = debug_info
             return None
 
         try:
@@ -1530,9 +1530,9 @@ class Logging(LiteLLMLoggingBaseClass):
             verbose_logger.debug(
                 f"response_cost_failure_debug_information: {debug_info}"
             )
-            self.model_call_details["response_cost_failure_debug_information"] = (
-                debug_info
-            )
+            self.model_call_details[
+                "response_cost_failure_debug_information"
+            ] = debug_info
 
         return None
 
@@ -1688,9 +1688,9 @@ class Logging(LiteLLMLoggingBaseClass):
                 result=logging_result
             )
 
-        self.model_call_details["standard_logging_object"] = (
-            self._build_standard_logging_payload(logging_result, start_time, end_time)
-        )
+        self.model_call_details[
+            "standard_logging_object"
+        ] = self._build_standard_logging_payload(logging_result, start_time, end_time)
 
         if (
             standard_logging_payload := self.model_call_details.get(
@@ -1768,9 +1768,9 @@ class Logging(LiteLLMLoggingBaseClass):
                 end_time = datetime.datetime.now()
             if self.completion_start_time is None:
                 self.completion_start_time = end_time
-                self.model_call_details["completion_start_time"] = (
-                    self.completion_start_time
-                )
+                self.model_call_details[
+                    "completion_start_time"
+                ] = self.completion_start_time
 
             self.model_call_details["log_event_type"] = "successful_api_call"
             self.model_call_details["end_time"] = end_time
@@ -1807,10 +1807,10 @@ class Logging(LiteLLMLoggingBaseClass):
                         end_time=end_time,
                     )
                 elif isinstance(result, dict) or isinstance(result, list):
-                    self.model_call_details["standard_logging_object"] = (
-                        self._build_standard_logging_payload(
-                            result, start_time, end_time
-                        )
+                    self.model_call_details[
+                        "standard_logging_object"
+                    ] = self._build_standard_logging_payload(
+                        result, start_time, end_time
                     )
                     if (
                         standard_logging_payload := self.model_call_details.get(
@@ -1819,9 +1819,9 @@ class Logging(LiteLLMLoggingBaseClass):
                     ) is not None:
                         emit_standard_logging_payload(standard_logging_payload)
             elif standard_logging_object is not None:
-                self.model_call_details["standard_logging_object"] = (
-                    standard_logging_object
-                )
+                self.model_call_details[
+                    "standard_logging_object"
+                ] = standard_logging_object
             else:
                 self.model_call_details["response_cost"] = None
 
@@ -1979,17 +1979,17 @@ class Logging(LiteLLMLoggingBaseClass):
                 verbose_logger.debug(
                     "Logging Details LiteLLM-Success Call streaming complete"
                 )
-                self.model_call_details["complete_streaming_response"] = (
-                    complete_streaming_response
-                )
-                self.model_call_details["response_cost"] = (
-                    self._response_cost_calculator(result=complete_streaming_response)
-                )
+                self.model_call_details[
+                    "complete_streaming_response"
+                ] = complete_streaming_response
+                self.model_call_details[
+                    "response_cost"
+                ] = self._response_cost_calculator(result=complete_streaming_response)
                 ## STANDARDIZED LOGGING PAYLOAD
-                self.model_call_details["standard_logging_object"] = (
-                    self._build_standard_logging_payload(
-                        complete_streaming_response, start_time, end_time
-                    )
+                self.model_call_details[
+                    "standard_logging_object"
+                ] = self._build_standard_logging_payload(
+                    complete_streaming_response, start_time, end_time
                 )
                 if (
                     standard_logging_payload := self.model_call_details.get(
@@ -2323,10 +2323,10 @@ class Logging(LiteLLMLoggingBaseClass):
                             )
                         else:
                             if self.stream and complete_streaming_response:
-                                self.model_call_details["complete_response"] = (
-                                    self.model_call_details.get(
-                                        "complete_streaming_response", {}
-                                    )
+                                self.model_call_details[
+                                    "complete_response"
+                                ] = self.model_call_details.get(
+                                    "complete_streaming_response", {}
                                 )
                                 result = self.model_call_details["complete_response"]
                             openMeterLogger.log_success_event(
@@ -2350,10 +2350,10 @@ class Logging(LiteLLMLoggingBaseClass):
                             )
                         else:
                             if self.stream and complete_streaming_response:
-                                self.model_call_details["complete_response"] = (
-                                    self.model_call_details.get(
-                                        "complete_streaming_response", {}
-                                    )
+                                self.model_call_details[
+                                    "complete_response"
+                                ] = self.model_call_details.get(
+                                    "complete_streaming_response", {}
                                 )
                                 result = self.model_call_details["complete_response"]
 
@@ -2492,9 +2492,9 @@ class Logging(LiteLLMLoggingBaseClass):
         if complete_streaming_response is not None:
             print_verbose("Async success callbacks: Got a complete streaming response")
 
-            self.model_call_details["async_complete_streaming_response"] = (
-                complete_streaming_response
-            )
+            self.model_call_details[
+                "async_complete_streaming_response"
+            ] = complete_streaming_response
 
             try:
                 if self.model_call_details.get("cache_hit", False) is True:
@@ -2505,10 +2505,10 @@ class Logging(LiteLLMLoggingBaseClass):
                         model_call_details=self.model_call_details
                     )
                     # base_model defaults to None if not set on model_info
-                    self.model_call_details["response_cost"] = (
-                        self._response_cost_calculator(
-                            result=complete_streaming_response
-                        )
+                    self.model_call_details[
+                        "response_cost"
+                    ] = self._response_cost_calculator(
+                        result=complete_streaming_response
                     )
 
                 verbose_logger.debug(
@@ -2521,10 +2521,10 @@ class Logging(LiteLLMLoggingBaseClass):
                 self.model_call_details["response_cost"] = None
 
             ## STANDARDIZED LOGGING PAYLOAD
-            self.model_call_details["standard_logging_object"] = (
-                self._build_standard_logging_payload(
-                    complete_streaming_response, start_time, end_time
-                )
+            self.model_call_details[
+                "standard_logging_object"
+            ] = self._build_standard_logging_payload(
+                complete_streaming_response, start_time, end_time
             )
 
             # print standard logging payload
@@ -2551,9 +2551,9 @@ class Logging(LiteLLMLoggingBaseClass):
             # _success_handler_helper_fn
             if self.model_call_details.get("standard_logging_object") is None:
                 ## STANDARDIZED LOGGING PAYLOAD
-                self.model_call_details["standard_logging_object"] = (
-                    self._build_standard_logging_payload(result, start_time, end_time)
-                )
+                self.model_call_details[
+                    "standard_logging_object"
+                ] = self._build_standard_logging_payload(result, start_time, end_time)
 
                 # print standard logging payload
                 if (
@@ -2796,18 +2796,18 @@ class Logging(LiteLLMLoggingBaseClass):
 
         ## STANDARDIZED LOGGING PAYLOAD
 
-        self.model_call_details["standard_logging_object"] = (
-            get_standard_logging_object_payload(
-                kwargs=self.model_call_details,
-                init_response_obj={},
-                start_time=start_time,
-                end_time=end_time,
-                logging_obj=self,
-                status="failure",
-                error_str=str(exception),
-                original_exception=exception,
-                standard_built_in_tools_params=self.standard_built_in_tools_params,
-            )
+        self.model_call_details[
+            "standard_logging_object"
+        ] = get_standard_logging_object_payload(
+            kwargs=self.model_call_details,
+            init_response_obj={},
+            start_time=start_time,
+            end_time=end_time,
+            logging_obj=self,
+            status="failure",
+            error_str=str(exception),
+            original_exception=exception,
+            standard_built_in_tools_params=self.standard_built_in_tools_params,
         )
         return start_time, end_time
 
@@ -3774,9 +3774,9 @@ def _init_custom_logger_compatible_class(  # noqa: PLR0915
                 service_name=arize_config.project_name,
             )
 
-            os.environ["OTEL_EXPORTER_OTLP_TRACES_HEADERS"] = (
-                f"space_id={arize_config.space_key or arize_config.space_id},api_key={arize_config.api_key}"
-            )
+            os.environ[
+                "OTEL_EXPORTER_OTLP_TRACES_HEADERS"
+            ] = f"space_id={arize_config.space_key or arize_config.space_id},api_key={arize_config.api_key}"
             for callback in _in_memory_loggers:
                 if (
                     isinstance(callback, ArizeLogger)
@@ -3802,13 +3802,13 @@ def _init_custom_logger_compatible_class(  # noqa: PLR0915
                 existing_attrs = os.environ.get("OTEL_RESOURCE_ATTRIBUTES", "")
                 # Add openinference.project.name attribute
                 if existing_attrs:
-                    os.environ["OTEL_RESOURCE_ATTRIBUTES"] = (
-                        f"{existing_attrs},openinference.project.name={arize_phoenix_config.project_name}"
-                    )
+                    os.environ[
+                        "OTEL_RESOURCE_ATTRIBUTES"
+                    ] = f"{existing_attrs},openinference.project.name={arize_phoenix_config.project_name}"
                 else:
-                    os.environ["OTEL_RESOURCE_ATTRIBUTES"] = (
-                        f"openinference.project.name={arize_phoenix_config.project_name}"
-                    )
+                    os.environ[
+                        "OTEL_RESOURCE_ATTRIBUTES"
+                    ] = f"openinference.project.name={arize_phoenix_config.project_name}"
 
             # Set Phoenix project name from environment variable
             phoenix_project_name = os.environ.get("PHOENIX_PROJECT_NAME", None)
@@ -3816,19 +3816,19 @@ def _init_custom_logger_compatible_class(  # noqa: PLR0915
                 existing_attrs = os.environ.get("OTEL_RESOURCE_ATTRIBUTES", "")
                 # Add openinference.project.name attribute
                 if existing_attrs:
-                    os.environ["OTEL_RESOURCE_ATTRIBUTES"] = (
-                        f"{existing_attrs},openinference.project.name={phoenix_project_name}"
-                    )
+                    os.environ[
+                        "OTEL_RESOURCE_ATTRIBUTES"
+                    ] = f"{existing_attrs},openinference.project.name={phoenix_project_name}"
                 else:
-                    os.environ["OTEL_RESOURCE_ATTRIBUTES"] = (
-                        f"openinference.project.name={phoenix_project_name}"
-                    )
+                    os.environ[
+                        "OTEL_RESOURCE_ATTRIBUTES"
+                    ] = f"openinference.project.name={phoenix_project_name}"
 
             # auth can be disabled on local deployments of arize phoenix
             if arize_phoenix_config.otlp_auth_headers is not None:
-                os.environ["OTEL_EXPORTER_OTLP_TRACES_HEADERS"] = (
-                    arize_phoenix_config.otlp_auth_headers
-                )
+                os.environ[
+                    "OTEL_EXPORTER_OTLP_TRACES_HEADERS"
+                ] = arize_phoenix_config.otlp_auth_headers
 
             for callback in _in_memory_loggers:
                 if (
@@ -3907,7 +3907,9 @@ def _init_custom_logger_compatible_class(  # noqa: PLR0915
             from litellm.integrations.focus.focus_logger import FocusLogger
 
             for callback in _in_memory_loggers:
-                if type(callback) is FocusLogger:  # exact match; exclude subclasses like VantageLogger
+                if (
+                    type(callback) is FocusLogger
+                ):  # exact match; exclude subclasses like VantageLogger
                     return callback  # type: ignore
             focus_logger = FocusLogger()
             _in_memory_loggers.append(focus_logger)
@@ -4013,9 +4015,9 @@ def _init_custom_logger_compatible_class(  # noqa: PLR0915
                 exporter="otlp_http",
                 endpoint="https://langtrace.ai/api/trace",
             )
-            os.environ["OTEL_EXPORTER_OTLP_TRACES_HEADERS"] = (
-                f"api_key={os.getenv('LANGTRACE_API_KEY')}"
-            )
+            os.environ[
+                "OTEL_EXPORTER_OTLP_TRACES_HEADERS"
+            ] = f"api_key={os.getenv('LANGTRACE_API_KEY')}"
             for callback in _in_memory_loggers:
                 if (
                     isinstance(callback, OpenTelemetry)
@@ -4289,7 +4291,9 @@ def get_custom_logger_compatible_class(  # noqa: PLR0915
             from litellm.integrations.focus.focus_logger import FocusLogger
 
             for callback in _in_memory_loggers:
-                if type(callback) is FocusLogger:  # exact match; exclude subclasses like VantageLogger
+                if (
+                    type(callback) is FocusLogger
+                ):  # exact match; exclude subclasses like VantageLogger
                     return callback
         elif logging_integration == "vantage":
             from litellm.integrations.vantage.vantage_logger import VantageLogger
@@ -4937,10 +4941,10 @@ class StandardLoggingPayloadSetup:
             for key in StandardLoggingHiddenParams.__annotations__.keys():
                 if key in hidden_params:
                     if key == "additional_headers":
-                        clean_hidden_params["additional_headers"] = (
-                            StandardLoggingPayloadSetup.get_additional_headers(
-                                hidden_params[key]
-                            )
+                        clean_hidden_params[
+                            "additional_headers"
+                        ] = StandardLoggingPayloadSetup.get_additional_headers(
+                            hidden_params[key]
                         )
                     else:
                         clean_hidden_params[key] = hidden_params[key]  # type: ignore
@@ -5579,9 +5583,9 @@ def scrub_sensitive_keys_in_metadata(litellm_params: Optional[dict]):
     ):
         for k, v in metadata["user_api_key_metadata"].items():
             if k == "logging":  # prevent logging user logging keys
-                cleaned_user_api_key_metadata[k] = (
-                    "scrubbed_by_litellm_for_sensitive_keys"
-                )
+                cleaned_user_api_key_metadata[
+                    k
+                ] = "scrubbed_by_litellm_for_sensitive_keys"
             else:
                 cleaned_user_api_key_metadata[k] = v
 

litellm/litellm_core_utils/prompt_templates/factory.py

@@ -2442,7 +2442,9 @@ def anthropic_messages_pt(  # noqa: PLR0915
                             _document_content_element = cast(
                                 AnthropicMessagesDocumentParam,
                                 add_cache_control_to_content(
-                                    anthropic_content_element=cast(AnthropicMessagesDocumentParam, m),
+                                    anthropic_content_element=cast(
+                                        AnthropicMessagesDocumentParam, m
+                                    ),
                                     original_content_element=dict(m),
                                 ),
                             )
@@ -2454,10 +2456,18 @@ def anthropic_messages_pt(  # noqa: PLR0915
                                 )
                             )
                             _file_content_element = add_cache_control_to_content(
-                                anthropic_content_element=cast(AnthropicMessagesDocumentParam, _file_content_element),
+                                anthropic_content_element=cast(
+                                    AnthropicMessagesDocumentParam,
+                                    _file_content_element,
+                                ),
                                 original_content_element=dict(m),
                             )
-                            user_content.append(cast(AnthropicMessagesDocumentParam,_file_content_element))
+                            user_content.append(
+                                cast(
+                                    AnthropicMessagesDocumentParam,
+                                    _file_content_element,
+                                )
+                            )
                 elif isinstance(user_message_types_block["content"], str):
                     _anthropic_content_text_element: AnthropicMessagesTextParam = {
                         "type": "text",

litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py

@@ -780,7 +780,7 @@ class LiteLLMAnthropicMessagesAdapter:
                 # Keep Anthropic-native tools in their original format
                 new_tools.append(tool)  # type: ignore[arg-type]
                 continue
-            
+
             original_name = tool["name"]
             truncated_name = truncate_tool_name(original_name)
 

litellm/llms/base_llm/videos/transformation.py

@@ -336,9 +336,7 @@ class BaseVideoConfig(ABC):
         Returns:
             Tuple[str, Dict]: (url, data) for the POST request
         """
-        raise NotImplementedError(
-            "video edit is not supported for this provider"
-        )
+        raise NotImplementedError("video edit is not supported for this provider")
 
     def transform_video_edit_response(
         self,
@@ -346,9 +344,7 @@ class BaseVideoConfig(ABC):
         logging_obj: LiteLLMLoggingObj,
         custom_llm_provider: Optional[str] = None,
     ) -> VideoObject:
-        raise NotImplementedError(
-            "video edit is not supported for this provider"
-        )
+        raise NotImplementedError("video edit is not supported for this provider")
 
     def transform_video_extension_request(
         self,
@@ -366,9 +362,7 @@ class BaseVideoConfig(ABC):
         Returns:
             Tuple[str, Dict]: (url, data) for the POST request
         """
-        raise NotImplementedError(
-            "video extension is not supported for this provider"
-        )
+        raise NotImplementedError("video extension is not supported for this provider")
 
     def transform_video_extension_response(
         self,
@@ -376,9 +370,7 @@ class BaseVideoConfig(ABC):
         logging_obj: LiteLLMLoggingObj,
         custom_llm_provider: Optional[str] = None,
     ) -> VideoObject:
-        raise NotImplementedError(
-            "video extension is not supported for this provider"
-        )
+        raise NotImplementedError("video extension is not supported for this provider")
 
     def get_error_class(
         self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers]

litellm/llms/custom_httpx/llm_http_handler.py

@@ -6162,7 +6162,10 @@ class BaseLLMHTTPHandler:
             litellm_params=dict(litellm_params),
         )
 
-        url, files_list = video_provider_config.transform_video_create_character_request(
+        (
+            url,
+            files_list,
+        ) = video_provider_config.transform_video_create_character_request(
             name=name,
             video=video,
             api_base=api_base,
@@ -6230,7 +6233,10 @@ class BaseLLMHTTPHandler:
             litellm_params=dict(litellm_params),
         )
 
-        url, files_list = video_provider_config.transform_video_create_character_request(
+        (
+            url,
+            files_list,
+        ) = video_provider_config.transform_video_create_character_request(
             name=name,
             video=video,
             api_base=api_base,
@@ -6324,11 +6330,7 @@ class BaseLLMHTTPHandler:
         )
 
         try:
-            response = sync_httpx_client.get(
-                url=url,
-                headers=headers,
-                params=params
-            )
+            response = sync_httpx_client.get(url=url, headers=headers, params=params)
             response.raise_for_status()
             return video_provider_config.transform_video_get_character_response(
                 raw_response=response,
@@ -6386,9 +6388,7 @@ class BaseLLMHTTPHandler:
 
         try:
             response = await async_httpx_client.get(
-                url=url,
-                headers=headers,
-                params=params
+                url=url, headers=headers, params=params
             )
             response.raise_for_status()
             return video_provider_config.transform_video_get_character_response(

litellm/llms/gemini/videos/transformation.py

@@ -525,28 +525,47 @@ class GeminiVideoConfig(BaseVideoConfig):
         """Video delete is not supported."""
         raise NotImplementedError("Video delete is not supported by Google Veo.")
 
-    def transform_video_create_character_request(self, name, video, api_base, litellm_params, headers):
+    def transform_video_create_character_request(
+        self, name, video, api_base, litellm_params, headers
+    ):
         raise NotImplementedError("video create character is not supported for Gemini")
 
     def transform_video_create_character_response(self, raw_response, logging_obj):
         raise NotImplementedError("video create character is not supported for Gemini")
 
-    def transform_video_get_character_request(self, character_id, api_base, litellm_params, headers):
+    def transform_video_get_character_request(
+        self, character_id, api_base, litellm_params, headers
+    ):
         raise NotImplementedError("video get character is not supported for Gemini")
 
     def transform_video_get_character_response(self, raw_response, logging_obj):
         raise NotImplementedError("video get character is not supported for Gemini")
 
-    def transform_video_edit_request(self, prompt, video_id, api_base, litellm_params, headers, extra_body=None):
+    def transform_video_edit_request(
+        self, prompt, video_id, api_base, litellm_params, headers, extra_body=None
+    ):
         raise NotImplementedError("video edit is not supported for Gemini")
 
-    def transform_video_edit_response(self, raw_response, logging_obj, custom_llm_provider=None):
+    def transform_video_edit_response(
+        self, raw_response, logging_obj, custom_llm_provider=None
+    ):
         raise NotImplementedError("video edit is not supported for Gemini")
 
-    def transform_video_extension_request(self, prompt, video_id, seconds, api_base, litellm_params, headers, extra_body=None):
+    def transform_video_extension_request(
+        self,
+        prompt,
+        video_id,
+        seconds,
+        api_base,
+        litellm_params,
+        headers,
+        extra_body=None,
+    ):
         raise NotImplementedError("video extension is not supported for Gemini")
 
-    def transform_video_extension_response(self, raw_response, logging_obj, custom_llm_provider=None):
+    def transform_video_extension_response(
+        self, raw_response, logging_obj, custom_llm_provider=None
+    ):
         raise NotImplementedError("video extension is not supported for Gemini")
 
     def get_error_class(

litellm/llms/moonshot/chat/transformation.py

@@ -19,7 +19,8 @@ class MoonshotChatConfig(OpenAIGPTConfig):
     @overload
     def _transform_messages(
         self, messages: List[AllMessageValues], model: str, is_async: Literal[True]
-    ) -> Coroutine[Any, Any, List[AllMessageValues]]: ...
+    ) -> Coroutine[Any, Any, List[AllMessageValues]]:
+        ...
 
     @overload
     def _transform_messages(
@@ -27,7 +28,8 @@ class MoonshotChatConfig(OpenAIGPTConfig):
         messages: List[AllMessageValues],
         model: str,
         is_async: Literal[False] = False,
-    ) -> List[AllMessageValues]: ...
+    ) -> List[AllMessageValues]:
+        ...
 
     def _transform_messages(
         self, messages: List[AllMessageValues], model: str, is_async: bool = False
@@ -53,9 +55,13 @@ class MoonshotChatConfig(OpenAIGPTConfig):
             messages = handle_messages_with_content_list_to_str_conversion(messages)
 
         if is_async:
-            return super()._transform_messages(messages=messages, model=model, is_async=True)
+            return super()._transform_messages(
+                messages=messages, model=model, is_async=True
+            )
         else:
-            return super()._transform_messages(messages=messages, model=model, is_async=False)
+            return super()._transform_messages(
+                messages=messages, model=model, is_async=False
+            )
 
     def _get_openai_compatible_provider_info(
         self, api_base: Optional[str], api_key: Optional[str]
@@ -141,7 +147,9 @@ class MoonshotChatConfig(OpenAIGPTConfig):
                 optional_params["temperature"] = 0.3
         return optional_params
 
-    def fill_reasoning_content(self, messages: List[AllMessageValues]) -> List[AllMessageValues]:
+    def fill_reasoning_content(
+        self, messages: List[AllMessageValues]
+    ) -> List[AllMessageValues]:
         """
         Moonshot reasoning models require `reasoning_content` on every assistant
         message that contains tool_calls (multi-turn tool-calling flows).

litellm/llms/runwayml/videos/transformation.py

@@ -592,28 +592,51 @@ class RunwayMLVideoConfig(BaseVideoConfig):
 
         return video_obj
 
-    def transform_video_create_character_request(self, name, video, api_base, litellm_params, headers):
-        raise NotImplementedError("video create character is not supported for RunwayML")
+    def transform_video_create_character_request(
+        self, name, video, api_base, litellm_params, headers
+    ):
+        raise NotImplementedError(
+            "video create character is not supported for RunwayML"
+        )
 
     def transform_video_create_character_response(self, raw_response, logging_obj):
-        raise NotImplementedError("video create character is not supported for RunwayML")
+        raise NotImplementedError(
+            "video create character is not supported for RunwayML"
+        )
 
-    def transform_video_get_character_request(self, character_id, api_base, litellm_params, headers):
+    def transform_video_get_character_request(
+        self, character_id, api_base, litellm_params, headers
+    ):
         raise NotImplementedError("video get character is not supported for RunwayML")
 
     def transform_video_get_character_response(self, raw_response, logging_obj):
         raise NotImplementedError("video get character is not supported for RunwayML")
 
-    def transform_video_edit_request(self, prompt, video_id, api_base, litellm_params, headers, extra_body=None):
+    def transform_video_edit_request(
+        self, prompt, video_id, api_base, litellm_params, headers, extra_body=None
+    ):
         raise NotImplementedError("video edit is not supported for RunwayML")
 
-    def transform_video_edit_response(self, raw_response, logging_obj, custom_llm_provider=None):
+    def transform_video_edit_response(
+        self, raw_response, logging_obj, custom_llm_provider=None
+    ):
         raise NotImplementedError("video edit is not supported for RunwayML")
 
-    def transform_video_extension_request(self, prompt, video_id, seconds, api_base, litellm_params, headers, extra_body=None):
+    def transform_video_extension_request(
+        self,
+        prompt,
+        video_id,
+        seconds,
+        api_base,
+        litellm_params,
+        headers,
+        extra_body=None,
+    ):
         raise NotImplementedError("video extension is not supported for RunwayML")
 
-    def transform_video_extension_response(self, raw_response, logging_obj, custom_llm_provider=None):
+    def transform_video_extension_response(
+        self, raw_response, logging_obj, custom_llm_provider=None
+    ):
         raise NotImplementedError("video extension is not supported for RunwayML")
 
     def get_error_class(

litellm/llms/sagemaker/chat/transformation.py

@@ -184,9 +184,7 @@ class SagemakerChatConfig(OpenAIGPTConfig, BaseAWSLLM):
                 llm_provider = LlmProviders(custom_llm_provider)
             except ValueError:
                 llm_provider = LlmProviders.SAGEMAKER_CHAT
-            client = get_async_httpx_client(
-                llm_provider=llm_provider, params={}
-            )
+            client = get_async_httpx_client(llm_provider=llm_provider, params={})
 
         try:
             response = await client.post(

litellm/llms/vertex_ai/batches/transformation.py

@@ -142,8 +142,8 @@ class VertexAIBatchTransformation:
         Gets the output file id from the Vertex AI Batch response
         """
 
-        output_file_id: str = (
-            response.get("outputInfo", OutputInfo()).get("gcsOutputDirectory", "")
+        output_file_id: str = response.get("outputInfo", OutputInfo()).get(
+            "gcsOutputDirectory", ""
         )
         if output_file_id:
             output_file_id = output_file_id.rstrip("/") + "/predictions.jsonl"

litellm/llms/vertex_ai/videos/transformation.py

@@ -624,28 +624,51 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase):
         """Video delete is not supported."""
         raise NotImplementedError("Video delete is not supported by Vertex AI Veo.")
 
-    def transform_video_create_character_request(self, name, video, api_base, litellm_params, headers):
-        raise NotImplementedError("video create character is not supported for Vertex AI")
+    def transform_video_create_character_request(
+        self, name, video, api_base, litellm_params, headers
+    ):
+        raise NotImplementedError(
+            "video create character is not supported for Vertex AI"
+        )
 
     def transform_video_create_character_response(self, raw_response, logging_obj):
-        raise NotImplementedError("video create character is not supported for Vertex AI")
+        raise NotImplementedError(
+            "video create character is not supported for Vertex AI"
+        )
 
-    def transform_video_get_character_request(self, character_id, api_base, litellm_params, headers):
+    def transform_video_get_character_request(
+        self, character_id, api_base, litellm_params, headers
+    ):
         raise NotImplementedError("video get character is not supported for Vertex AI")
 
     def transform_video_get_character_response(self, raw_response, logging_obj):
         raise NotImplementedError("video get character is not supported for Vertex AI")
 
-    def transform_video_edit_request(self, prompt, video_id, api_base, litellm_params, headers, extra_body=None):
+    def transform_video_edit_request(
+        self, prompt, video_id, api_base, litellm_params, headers, extra_body=None
+    ):
         raise NotImplementedError("video edit is not supported for Vertex AI")
 
-    def transform_video_edit_response(self, raw_response, logging_obj, custom_llm_provider=None):
+    def transform_video_edit_response(
+        self, raw_response, logging_obj, custom_llm_provider=None
+    ):
         raise NotImplementedError("video edit is not supported for Vertex AI")
 
-    def transform_video_extension_request(self, prompt, video_id, seconds, api_base, litellm_params, headers, extra_body=None):
+    def transform_video_extension_request(
+        self,
+        prompt,
+        video_id,
+        seconds,
+        api_base,
+        litellm_params,
+        headers,
+        extra_body=None,
+    ):
         raise NotImplementedError("video extension is not supported for Vertex AI")
 
-    def transform_video_extension_response(self, raw_response, logging_obj, custom_llm_provider=None):
+    def transform_video_extension_response(
+        self, raw_response, logging_obj, custom_llm_provider=None
+    ):
         raise NotImplementedError("video extension is not supported for Vertex AI")
 
     def get_error_class(

litellm/main.py

@@ -7533,9 +7533,7 @@ def stream_chunk_builder(  # noqa: PLR0915
             # the final chunk.
             all_annotations: list = []
             for ac in annotation_chunks:
-                all_annotations.extend(
-                    ac["choices"][0]["delta"]["annotations"]
-                )
+                all_annotations.extend(ac["choices"][0]["delta"]["annotations"])
             response["choices"][0]["message"]["annotations"] = all_annotations
 
         audio_chunks = [

litellm/model_prices_and_context_window_backup.json

@@ -32354,6 +32354,53 @@
         "supports_vision": true,
         "supports_web_search": true
     },
+    "xai/grok-4.20-multi-agent-beta-0309": {
+        "cache_read_input_token_cost": 2e-07,
+        "input_cost_per_token": 2e-06,
+        "litellm_provider": "xai",
+        "max_input_tokens": 2000000,
+        "max_output_tokens": 2000000,
+        "max_tokens": 2000000,
+        "mode": "chat",
+        "output_cost_per_token": 6e-06,
+        "source": "https://docs.x.ai/docs/models",
+        "supports_function_calling": true,
+        "supports_reasoning": true,
+        "supports_tool_choice": true,
+        "supports_vision": true,
+        "supports_web_search": true
+    },
+    "xai/grok-4.20-beta-0309-reasoning": {
+        "cache_read_input_token_cost": 2e-07,
+        "input_cost_per_token": 2e-06,
+        "litellm_provider": "xai",
+        "max_input_tokens": 2000000,
+        "max_output_tokens": 2000000,
+        "max_tokens": 2000000,
+        "mode": "chat",
+        "output_cost_per_token": 6e-06,
+        "source": "https://docs.x.ai/docs/models",
+        "supports_function_calling": true,
+        "supports_reasoning": true,
+        "supports_tool_choice": true,
+        "supports_vision": true,
+        "supports_web_search": true
+    },
+    "xai/grok-4.20-beta-0309-non-reasoning": {
+        "cache_read_input_token_cost": 2e-07,
+        "input_cost_per_token": 2e-06,
+        "litellm_provider": "xai",
+        "max_input_tokens": 2000000,
+        "max_output_tokens": 2000000,
+        "max_tokens": 2000000,
+        "mode": "chat",
+        "output_cost_per_token": 6e-06,
+        "source": "https://docs.x.ai/docs/models",
+        "supports_function_calling": true,
+        "supports_tool_choice": true,
+        "supports_vision": true,
+        "supports_web_search": true
+    },
     "xai/grok-beta": {
         "input_cost_per_token": 5e-06,
         "litellm_provider": "xai",

litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py

@@ -677,7 +677,60 @@ async def oauth_authorization_server_mcp(
 # Alias for standard OpenID discovery
 @router.get("/.well-known/openid-configuration")
 async def openid_configuration(request: Request):
-    return await oauth_authorization_server_mcp(request)
+    response = await oauth_authorization_server_mcp(request)
+
+    # If MCPJWTSigner is active, augment the discovery doc with JWKS fields so
+    # MCP servers and gateways (e.g. AWS Bedrock AgentCore Gateway) can resolve
+    # the signing keys and verify liteLLM-issued tokens.
+    try:
+        from litellm.proxy.guardrails.guardrail_hooks.mcp_jwt_signer.mcp_jwt_signer import (
+            get_mcp_jwt_signer,
+        )
+
+        signer = get_mcp_jwt_signer()
+        if signer is not None:
+            request_base_url = get_request_base_url(request)
+            if isinstance(response, dict):
+                response = {
+                    **response,
+                    "jwks_uri": f"{request_base_url}/.well-known/jwks.json",
+                    "id_token_signing_alg_values_supported": ["RS256"],
+                }
+    except ImportError:
+        pass
+
+    return response
+
+
+@router.get("/.well-known/jwks.json")
+async def jwks_json(request: Request):
+    """
+    JSON Web Key Set endpoint.
+
+    Returns the RSA public key used by MCPJWTSigner to sign outbound MCP tokens.
+    MCP servers and gateways use this endpoint to verify liteLLM-issued JWTs.
+
+    Returns an empty key set if MCPJWTSigner is not configured.
+    """
+    try:
+        from litellm.proxy.guardrails.guardrail_hooks.mcp_jwt_signer.mcp_jwt_signer import (
+            get_mcp_jwt_signer,
+        )
+
+        signer = get_mcp_jwt_signer()
+        if signer is not None:
+            return JSONResponse(
+                content=signer.get_jwks(),
+                headers={"Cache-Control": f"public, max-age={signer.jwks_max_age}"},
+            )
+    except ImportError:
+        pass
+
+    # No signer active — return empty key set; short cache so activation is picked up quickly.
+    return JSONResponse(
+        content={"keys": []},
+        headers={"Cache-Control": "public, max-age=60"},
+    )
 
 
 # Additional legacy pattern support

litellm/proxy/_experimental/mcp_server/mcp_server_manager.py

@@ -1908,7 +1908,15 @@ class MCPServerManager:
         user_api_key_auth: Optional[UserAPIKeyAuth],
         proxy_logging_obj: ProxyLogging,
         server: MCPServer,
-    ):
+        raw_headers: Optional[Dict[str, str]] = None,
+    ) -> Dict[str, Any]:
+        """
+        Run pre-call checks and guardrail hooks for an MCP tool call.
+
+        Returns a dict that may contain:
+        - "arguments": hook-modified tool arguments (only if changed)
+        - "extra_headers": headers injected by pre_mcp_call guardrail hooks
+        """
         ## check if the tool is allowed or banned for the given server
         if not self.check_allowed_or_banned_tools(name, server):
             raise HTTPException(
@@ -1932,6 +1940,14 @@ class MCPServerManager:
             server=server,
         )
 
+        # Extract incoming Bearer token from raw request headers so
+        # guardrails like MCPJWTSigner can verify + re-sign it (FR-5).
+        normalized_raw = {k.lower(): v for k, v in (raw_headers or {}).items()}
+        incoming_bearer_token: Optional[str] = None
+        auth_hdr = normalized_raw.get("authorization", "")
+        if auth_hdr.lower().startswith("bearer "):
+            incoming_bearer_token = auth_hdr[len("bearer ") :]
+
         pre_hook_kwargs = {
             "name": name,
             "arguments": arguments,
@@ -1957,6 +1973,7 @@ class MCPServerManager:
                 if user_api_key_auth
                 else None
             ),
+            "incoming_bearer_token": incoming_bearer_token,
         }
 
         # Create MCP request object for processing
@@ -1969,6 +1986,7 @@ class MCPServerManager:
             mcp_request_obj, pre_hook_kwargs
         )
 
+        hook_result: Dict[str, Any] = {}
         try:
             # Use standard pre_call_hook
             modified_data = await proxy_logging_obj.pre_call_hook(
@@ -1984,7 +2002,9 @@ class MCPServerManager:
                     )
                 )
                 if modified_kwargs.get("arguments") != arguments:
-                    arguments = modified_kwargs["arguments"]
+                    hook_result["arguments"] = modified_kwargs["arguments"]
+                if modified_kwargs.get("extra_headers"):
+                    hook_result["extra_headers"] = modified_kwargs["extra_headers"]
 
         except (
             BlockedPiiEntityError,
@@ -1995,6 +2015,8 @@ class MCPServerManager:
             verbose_logger.error(f"Guardrail blocked MCP tool call pre call: {str(e)}")
             raise e
 
+        return hook_result
+
     def _create_during_hook_task(
         self,
         name: str,
@@ -2047,6 +2069,7 @@ class MCPServerManager:
         raw_headers: Optional[Dict[str, str]],
         proxy_logging_obj: Optional[ProxyLogging],
         host_progress_callback: Optional[Callable] = None,
+        hook_extra_headers: Optional[Dict[str, str]] = None,
     ) -> CallToolResult:
         """
         Call a regular MCP tool using the MCP client.
@@ -2061,6 +2084,9 @@ class MCPServerManager:
             oauth2_headers: Optional OAuth2 headers
             raw_headers: Optional raw headers from the request
             proxy_logging_obj: Optional ProxyLogging object for hook integration
+            host_progress_callback: Optional callback for progress updates
+            hook_extra_headers: Optional headers injected by pre_mcp_call guardrail
+                hooks. Merged last (highest priority) into outbound request headers.
 
         Returns:
             CallToolResult from the MCP server
@@ -2116,6 +2142,31 @@ class MCPServerManager:
                 extra_headers = {}
             extra_headers.update(mcp_server.static_headers)
 
+        if hook_extra_headers:
+            if extra_headers is None:
+                extra_headers = {}
+            if "Authorization" in hook_extra_headers:
+                if "Authorization" in extra_headers:
+                    verbose_logger.warning(
+                        "MCPServerManager: hook_extra_headers 'Authorization' will overwrite "
+                        "the existing Authorization header from static_headers. "
+                        "The hook JWT will take precedence."
+                    )
+                elif server_auth_header is not None:
+                    # server_auth_header is passed separately to _create_mcp_client as
+                    # auth_value.  Both will reach the upstream server — warn so admins
+                    # know two Authorization credentials are being sent.
+                    verbose_logger.warning(
+                        "MCPServerManager: hook_extra_headers injects 'Authorization' while "
+                        "server '%s' already has a configured authentication_token. "
+                        "Both credentials will be sent; the hook header is in extra_headers "
+                        "and the server token is in auth_value — the upstream server decides "
+                        "which one wins.  Consider unsetting authentication_token if you want "
+                        "the hook JWT to be the sole credential.",
+                        mcp_server.server_name or mcp_server.name,
+                    )
+            extra_headers.update(hook_extra_headers)
+
         stdio_env = self._build_stdio_env(mcp_server, raw_headers)
 
         client = await self._create_mcp_client(
@@ -2201,15 +2252,19 @@ class MCPServerManager:
         # Allow validation and modification of tool calls before execution
         # Using standard pre_call_hook
         #########################################################
+        hook_result: Dict[str, Any] = {}
         if proxy_logging_obj:
-            await self.pre_call_tool_check(
+            hook_result = await self.pre_call_tool_check(
                 name=name,
                 arguments=arguments,
                 server_name=server_name,
                 user_api_key_auth=user_api_key_auth,
                 proxy_logging_obj=proxy_logging_obj,
                 server=mcp_server,
+                raw_headers=raw_headers,
             )
+            if "arguments" in hook_result:
+                arguments = hook_result["arguments"]
 
         # Prepare tasks for during hooks
         tasks = []
@@ -2227,8 +2282,16 @@ class MCPServerManager:
         # For OpenAPI servers, call the tool handler directly instead of via MCP client
         if mcp_server.spec_path:
             verbose_logger.debug(
-                f"Calling OpenAPI tool {name} directly via HTTP handler"
+                "Calling OpenAPI tool %s directly via HTTP handler", name
             )
+            if hook_result.get("extra_headers"):
+                verbose_logger.warning(
+                    "pre_mcp_call hook returned extra_headers for OpenAPI-backed "
+                    "MCP server '%s' — header injection is not supported for "
+                    "OpenAPI servers; headers will be ignored. Use SSE/HTTP "
+                    "transport to enable hook header injection.",
+                    server_name,
+                )
             tasks.append(
                 asyncio.create_task(
                     self._call_openapi_tool_handler(mcp_server, name, arguments)
@@ -2247,6 +2310,7 @@ class MCPServerManager:
                 raw_headers=raw_headers,
                 proxy_logging_obj=proxy_logging_obj,
                 host_progress_callback=host_progress_callback,
+                hook_extra_headers=hook_result.get("extra_headers"),
             )
 
         # For OpenAPI tools, await outside the client context

litellm/proxy/_experimental/mcp_server/rest_endpoints.py

@@ -903,12 +903,12 @@ if MCP_AVAILABLE:
         try:
             client_id, client_secret, scopes = _extract_credentials(request)
 
-            _oauth2_flow: Optional[Literal["client_credentials", "authorization_code"]] = (
-                request.oauth2_flow or (
-                    "client_credentials"
-                    if client_id and client_secret and request.token_url
-                    else None
-                )
+            _oauth2_flow: Optional[
+                Literal["client_credentials", "authorization_code"]
+            ] = request.oauth2_flow or (
+                "client_credentials"
+                if client_id and client_secret and request.token_url
+                else None
             )
             # client_credentials requires token_url to fetch a token; without it the
             # incoming auth header would be dropped with nothing to replace it.

litellm/proxy/_types.py

@@ -2471,6 +2471,9 @@ class UserAPIKeyAuth(
         Any
     ] = None  # Expanded created_by user when expand=user is used
     end_user_object_permission: Optional[LiteLLM_ObjectPermissionTable] = None
+    # Decoded upstream IdP claims (groups, roles, etc.) propagated by JWT auth machinery
+    # and forwarded into outbound tokens by guardrails such as MCPJWTSigner.
+    jwt_claims: Optional[Dict] = None
 
     model_config = ConfigDict(arbitrary_types_allowed=True)
 

litellm/proxy/auth/auth_utils.py

@@ -680,7 +680,7 @@ def get_customer_user_header_from_mapping(user_id_mapping) -> Optional[list]:
 
     if customer_headers_mappings:
         return customer_headers_mappings
-    
+
     return None
 
 
@@ -754,15 +754,11 @@ def get_end_user_id_from_request_body(
                     user_id_str = str(header_value)
                     if user_id_str.strip():
                         return user_id_str
-                    
+
         elif isinstance(custom_header_name_to_check, str):
             for header_name, header_value in request_headers.items():
                 if header_name.lower() == custom_header_name_to_check.lower():
-                    user_id_str = (
-                        str(header_value)
-                        if header_value is not None
-                        else ""
-                    )
+                    user_id_str = str(header_value) if header_value is not None else ""
                     if user_id_str.strip():
                         return user_id_str
 

litellm/proxy/auth/user_api_key_auth.py

@@ -685,6 +685,7 @@ async def _user_api_key_auth_builder(  # noqa: PLR0915
                 do_standard_jwt_auth = True
                 if jwt_handler.litellm_jwtauth.virtual_key_claim_field is not None:
                     # Decode JWT to get claims without running full auth_builder
+                    jwt_claims: Optional[dict]
                     if jwt_handler.litellm_jwtauth.oidc_userinfo_enabled:
                         jwt_claims = await jwt_handler.get_oidc_userinfo(token=api_key)
                     else:
@@ -700,6 +701,7 @@ async def _user_api_key_auth_builder(  # noqa: PLR0915
                     )
                     if valid_token is not None:
                         api_key = valid_token.token or ""
+                        valid_token.jwt_claims = jwt_claims
                         do_standard_jwt_auth = False
                         # Fall through to virtual key checks
 
@@ -729,6 +731,7 @@ async def _user_api_key_auth_builder(  # noqa: PLR0915
                     team_membership: Optional[LiteLLM_TeamMembership] = result.get(
                         "team_membership", None
                     )
+                    jwt_claims = result.get("jwt_claims", None)
 
                     global_proxy_spend = await get_global_proxy_spend(
                         litellm_proxy_admin_name=litellm_proxy_admin_name,
@@ -757,6 +760,7 @@ async def _user_api_key_auth_builder(  # noqa: PLR0915
                             org_id=org_id,
                             end_user_id=end_user_id,
                             parent_otel_span=parent_otel_span,
+                            jwt_claims=jwt_claims,
                         )
 
                     valid_token = UserAPIKeyAuth(
@@ -803,6 +807,7 @@ async def _user_api_key_auth_builder(  # noqa: PLR0915
                         team_metadata=(
                             team_object.metadata if team_object is not None else None
                         ),
+                        jwt_claims=jwt_claims,
                     )
 
                     # Check if model has zero cost - if so, skip all budget checks

litellm/proxy/batches_endpoints/endpoints.py

@@ -537,9 +537,10 @@ async def retrieve_batch(  # noqa: PLR0915
         )
 
         # Fix: bug_feb14_batch_retrieve_returns_raw_input_file_id
-        # Resolve raw provider input_file_id to unified ID.
+        # Resolve raw provider file IDs (input, output, error) to unified IDs.
         if unified_batch_id:
             await resolve_input_file_id_to_unified(response, prisma_client)
+            await resolve_output_file_ids_to_unified(response, prisma_client)
 
         ### ALERTING ###
         asyncio.create_task(

litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/__init__.py

@@ -0,0 +1,84 @@
+"""MCP JWT Signer guardrail — built-in LiteLLM guardrail for zero trust MCP auth."""
+
+from typing import TYPE_CHECKING
+
+from litellm.types.guardrails import SupportedGuardrailIntegrations
+
+from .mcp_jwt_signer import MCPJWTSigner, get_mcp_jwt_signer
+
+if TYPE_CHECKING:
+    from litellm.types.guardrails import Guardrail, LitellmParams
+
+
+def initialize_guardrail(
+    litellm_params: "LitellmParams", guardrail: "Guardrail"
+) -> MCPJWTSigner:
+    import litellm
+
+    guardrail_name = guardrail.get("guardrail_name")
+    if not guardrail_name:
+        raise ValueError("MCPJWTSigner guardrail requires a guardrail_name")
+
+    mode = litellm_params.mode
+    if mode != "pre_mcp_call":
+        raise ValueError(
+            f"MCPJWTSigner guardrail '{guardrail_name}' has mode='{mode}' but must use "
+            "mode='pre_mcp_call'. JWT injection only fires for MCP tool calls."
+        )
+
+    optional_params = getattr(litellm_params, "optional_params", None)
+
+    def _get(key):  # type: ignore[no-untyped-def]
+        if optional_params is not None:
+            v = getattr(optional_params, key, None)
+            if v is not None:
+                return v
+        return getattr(litellm_params, key, None)
+
+    signer = MCPJWTSigner(
+        guardrail_name=guardrail_name,
+        event_hook=litellm_params.mode,
+        default_on=litellm_params.default_on,
+        # Core signing
+        issuer=_get("issuer"),
+        audience=_get("audience"),
+        ttl_seconds=_get("ttl_seconds"),
+        # FR-5: verify + re-sign
+        access_token_discovery_uri=_get("access_token_discovery_uri"),
+        token_introspection_endpoint=_get("token_introspection_endpoint"),
+        verify_issuer=_get("verify_issuer"),
+        verify_audience=_get("verify_audience"),
+        # FR-12: end-user identity mapping
+        end_user_claim_sources=_get("end_user_claim_sources"),
+        # FR-13: claim operations
+        add_claims=_get("add_claims"),
+        set_claims=_get("set_claims"),
+        remove_claims=_get("remove_claims"),
+        # FR-14: two-token model
+        channel_token_audience=_get("channel_token_audience"),
+        channel_token_ttl=_get("channel_token_ttl"),
+        # FR-15: incoming claim validation
+        required_claims=_get("required_claims"),
+        optional_claims=_get("optional_claims"),
+        # FR-9: debug headers
+        debug_headers=_get("debug_headers") or False,
+        # FR-10: configurable scopes
+        allowed_scopes=_get("allowed_scopes"),
+    )
+    litellm.logging_callback_manager.add_litellm_callback(signer)
+    return signer
+
+
+guardrail_initializer_registry = {
+    SupportedGuardrailIntegrations.MCP_JWT_SIGNER.value: initialize_guardrail,
+}
+
+guardrail_class_registry = {
+    SupportedGuardrailIntegrations.MCP_JWT_SIGNER.value: MCPJWTSigner,
+}
+
+__all__ = [
+    "MCPJWTSigner",
+    "initialize_guardrail",
+    "get_mcp_jwt_signer",
+]

litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py

@@ -0,0 +1,891 @@
+"""
+MCPJWTSigner — Built-in LiteLLM guardrail for zero trust MCP authentication.
+
+Signs outbound MCP requests with a LiteLLM-issued RS256 JWT so that MCP servers
+can trust a single signing authority (liteLLM) instead of every upstream IdP.
+
+Usage in config.yaml:
+
+    guardrails:
+      - guardrail_name: "mcp-jwt-signer"
+        litellm_params:
+          guardrail: mcp_jwt_signer
+          mode: "pre_mcp_call"
+          default_on: true
+
+          # Core signing config
+          issuer: "https://my-litellm.example.com"   # optional
+          audience: "mcp"                             # optional
+          ttl_seconds: 300                            # optional
+
+          # FR-5: Verify + re-sign — validate incoming Bearer token before signing
+          access_token_discovery_uri: "https://idp.example.com/.well-known/openid-configuration"
+          token_introspection_endpoint: "https://idp.example.com/introspect"  # opaque tokens
+          verify_issuer: "https://idp.example.com"   # expected iss in incoming JWT
+          verify_audience: "api://my-app"             # expected aud in incoming JWT
+
+          # FR-12: End-user identity mapping — ordered resolution chain
+          # Supported: token:<claim>, litellm:user_id, litellm:email,
+          #            litellm:end_user_id, litellm:team_id
+          end_user_claim_sources:
+            - "token:sub"
+            - "token:email"
+            - "litellm:user_id"
+
+          # FR-13: Claim operations
+          add_claims:             # add if key not already present in the JWT
+            deployment_id: "prod-001"
+          set_claims:             # always set (overrides computed value)
+            env: "production"
+          remove_claims:          # remove from final JWT
+            - "nbf"
+
+          # FR-14: Two-token model — issue a second JWT for the MCP transport channel
+          channel_token_audience: "bedrock-gateway"
+          channel_token_ttl: 60
+
+          # FR-15: Incoming claim validation — enforce required IdP claims
+          required_claims:
+            - "sub"
+            - "email"
+          optional_claims:        # pass through from jwt_claims into outbound JWT
+            - "groups"
+            - "roles"
+
+          # FR-9: Debug headers
+          debug_headers: false    # emit x-litellm-mcp-debug header when true
+
+          # FR-10: Configurable scopes — explicit list replaces auto-generation
+          allowed_scopes:
+            - "mcp:tools/call"
+            - "mcp:tools/list"
+
+MCP servers verify tokens via:
+  GET /.well-known/openid-configuration  → { jwks_uri: ".../.well-known/jwks.json" }
+  GET /.well-known/jwks.json             → RSA public key in JWKS format
+
+Optionally set MCP_JWT_SIGNING_KEY env var (PEM string or file:///path) to use
+your own RSA keypair. If unset, an RSA-2048 keypair is auto-generated at startup.
+"""
+
+import base64
+import hashlib
+import os
+import re
+import time
+from typing import Any, Dict, List, Optional, Union
+
+import jwt
+from cryptography.hazmat.primitives import serialization
+from cryptography.hazmat.primitives.asymmetric import rsa
+from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey
+
+from litellm._logging import verbose_proxy_logger
+from litellm.caching import DualCache
+from litellm.integrations.custom_guardrail import (
+    CustomGuardrail,
+    log_guardrail_information,
+)
+from litellm.proxy._types import UserAPIKeyAuth
+from litellm.types.utils import CallTypesLiteral
+
+# Module-level singleton for the JWKS discovery endpoint to access.
+_mcp_jwt_signer_instance: Optional["MCPJWTSigner"] = None
+
+# Simple in-memory JWKS cache: keyed by JWKS URI → (keys_list, fetched_at).
+_jwks_cache: Dict[str, tuple] = {}
+_JWKS_CACHE_TTL = 3600  # 1 hour
+
+
+def get_mcp_jwt_signer() -> Optional["MCPJWTSigner"]:
+    """Return the active MCPJWTSigner singleton, or None if not initialized."""
+    return _mcp_jwt_signer_instance
+
+
+def _load_private_key_from_env(env_var: str) -> RSAPrivateKey:
+    """Load an RSA private key from an env var (PEM string or file:// path)."""
+    key_material = os.environ.get(env_var, "")
+    if not key_material:
+        raise ValueError(
+            f"MCPJWTSigner: environment variable '{env_var}' is set but empty."
+        )
+    if key_material.startswith("file://"):
+        path = key_material[len("file://") :]
+        with open(path, "rb") as f:
+            key_bytes = f.read()
+    else:
+        key_bytes = key_material.encode("utf-8")
+    return serialization.load_pem_private_key(key_bytes, password=None)  # type: ignore[return-value]
+
+
+def _generate_rsa_key_pair() -> RSAPrivateKey:
+    """Generate a new RSA-2048 private key."""
+    return rsa.generate_private_key(
+        public_exponent=65537,
+        key_size=2048,
+    )
+
+
+def _int_to_base64url(n: int) -> str:
+    """Encode an integer as a base64url string (no padding)."""
+    byte_length = (n.bit_length() + 7) // 8
+    return (
+        base64.urlsafe_b64encode(n.to_bytes(byte_length, byteorder="big"))
+        .rstrip(b"=")
+        .decode("ascii")
+    )
+
+
+def _compute_kid(public_key: Any) -> str:
+    """Derive a key ID from the public key's DER encoding (SHA-256, first 16 hex chars)."""
+    der_bytes = public_key.public_bytes(
+        encoding=serialization.Encoding.DER,
+        format=serialization.PublicFormat.SubjectPublicKeyInfo,
+    )
+    return hashlib.sha256(der_bytes).hexdigest()[:16]
+
+
+async def _fetch_jwks(jwks_uri: str) -> List[Dict[str, Any]]:
+    """
+    Fetch and cache a JWKS from the given URI.
+
+    Results are cached for _JWKS_CACHE_TTL seconds to avoid hammering the IdP.
+    """
+    now = time.time()
+    cached = _jwks_cache.get(jwks_uri)
+    if cached is not None:
+        keys, fetched_at = cached
+        if now - fetched_at < _JWKS_CACHE_TTL:
+            return keys  # type: ignore[return-value]
+
+    from litellm.llms.custom_httpx.http_handler import (
+        get_async_httpx_client,
+        httpxSpecialProvider,
+    )
+
+    client = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check)
+    resp = await client.get(jwks_uri, headers={"Accept": "application/json"})
+    resp.raise_for_status()
+    keys = resp.json().get("keys", [])
+    _jwks_cache[jwks_uri] = (keys, now)
+    return keys  # type: ignore[return-value]
+
+
+async def _fetch_oidc_discovery(discovery_uri: str) -> Dict[str, Any]:
+    """Fetch an OIDC discovery document and return its parsed JSON."""
+    from litellm.llms.custom_httpx.http_handler import (
+        get_async_httpx_client,
+        httpxSpecialProvider,
+    )
+
+    client = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check)
+    resp = await client.get(discovery_uri, headers={"Accept": "application/json"})
+    resp.raise_for_status()
+    return resp.json()  # type: ignore[return-value]
+
+
+class MCPJWTSigner(CustomGuardrail):
+    """
+    Built-in LiteLLM guardrail that signs outbound MCP requests with a
+    LiteLLM-issued RS256 JWT, enabling zero trust authentication.
+
+    MCP servers verify tokens using liteLLM's OIDC discovery endpoint and
+    JWKS endpoint rather than trusting each upstream IdP directly.
+
+    The signed JWT carries:
+      - iss: LiteLLM issuer identifier
+      - aud: MCP audience (configurable)
+      - sub: End-user identity (resolved via end_user_claim_sources, RFC 8693)
+      - act: Actor/agent identity (team_id or org_id, RFC 8693 delegation)
+      - scope: Tool-level access scopes (configurable via allowed_scopes)
+      - iat, exp, nbf: Standard timing claims
+
+    Feature set:
+      FR-5:  Verify + re-sign (access_token_discovery_uri, token_introspection_endpoint)
+      FR-9:  Debug headers (debug_headers)
+      FR-10: Configurable scopes (allowed_scopes)
+      FR-12: Configurable end-user identity mapping (end_user_claim_sources)
+      FR-13: Claim operations (add_claims, set_claims, remove_claims)
+      FR-14: Two-token model (channel_token_audience, channel_token_ttl)
+      FR-15: Incoming claim validation (required_claims, optional_claims)
+    """
+
+    ALGORITHM = "RS256"
+    DEFAULT_TTL = 300
+    DEFAULT_AUDIENCE = "mcp"
+    SIGNING_KEY_ENV = "MCP_JWT_SIGNING_KEY"
+
+    def __init__(
+        self,
+        # Core signing config
+        issuer: Optional[str] = None,
+        audience: Optional[str] = None,
+        ttl_seconds: Optional[int] = None,
+        # FR-5: Verify + re-sign
+        access_token_discovery_uri: Optional[str] = None,
+        token_introspection_endpoint: Optional[str] = None,
+        verify_issuer: Optional[str] = None,
+        verify_audience: Optional[str] = None,
+        # FR-12: End-user identity mapping
+        end_user_claim_sources: Optional[List[str]] = None,
+        # FR-13: Claim operations
+        add_claims: Optional[Dict[str, Any]] = None,
+        set_claims: Optional[Dict[str, Any]] = None,
+        remove_claims: Optional[List[str]] = None,
+        # FR-14: Two-token model
+        channel_token_audience: Optional[str] = None,
+        channel_token_ttl: Optional[int] = None,
+        # FR-15: Incoming claim validation
+        required_claims: Optional[List[str]] = None,
+        optional_claims: Optional[List[str]] = None,
+        # FR-9: Debug headers
+        debug_headers: bool = False,
+        # FR-10: Configurable scopes
+        allowed_scopes: Optional[List[str]] = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+
+        # --- Signing key setup ---
+        key_material = os.environ.get(self.SIGNING_KEY_ENV)
+        if key_material:
+            self._private_key = _load_private_key_from_env(self.SIGNING_KEY_ENV)
+            self._persistent_key: bool = True
+            verbose_proxy_logger.info(
+                "MCPJWTSigner: loaded RSA key from env var %s", self.SIGNING_KEY_ENV
+            )
+        else:
+            self._private_key = _generate_rsa_key_pair()
+            self._persistent_key = False
+            verbose_proxy_logger.info(
+                "MCPJWTSigner: auto-generated RSA-2048 keypair (set %s to use your own key)",
+                self.SIGNING_KEY_ENV,
+            )
+
+        self._public_key = self._private_key.public_key()
+        self._kid = _compute_kid(self._public_key)
+
+        # --- Core config ---
+        self.issuer: str = (
+            issuer
+            or os.environ.get("MCP_JWT_ISSUER")
+            or os.environ.get("LITELLM_EXTERNAL_URL")
+            or "litellm"
+        )
+        self.audience: str = (
+            audience or os.environ.get("MCP_JWT_AUDIENCE") or self.DEFAULT_AUDIENCE
+        )
+        resolved_ttl = int(
+            ttl_seconds
+            if ttl_seconds is not None
+            else os.environ.get("MCP_JWT_TTL_SECONDS", str(self.DEFAULT_TTL))
+        )
+        if resolved_ttl <= 0:
+            raise ValueError(
+                f"MCPJWTSigner: ttl_seconds must be > 0, got {resolved_ttl}"
+            )
+        self.ttl_seconds: int = resolved_ttl
+
+        # --- FR-5: Verify + re-sign ---
+        self.access_token_discovery_uri: Optional[str] = access_token_discovery_uri
+        self.token_introspection_endpoint: Optional[str] = token_introspection_endpoint
+        self.verify_issuer: Optional[str] = verify_issuer
+        self.verify_audience: Optional[str] = verify_audience
+        # Cached OIDC discovery document (fetched lazily, TTL = 24 h)
+        self._oidc_discovery_doc: Optional[Dict[str, Any]] = None
+        self._oidc_discovery_fetched_at: float = 0.0
+
+        # --- FR-12: End-user identity mapping ---
+        # Default chain: try incoming JWT sub, fall back to litellm user_id
+        self.end_user_claim_sources: List[str] = end_user_claim_sources or [
+            "token:sub",
+            "litellm:user_id",
+        ]
+
+        # --- FR-13: Claim operations ---
+        self.add_claims: Dict[str, Any] = add_claims or {}
+        self.set_claims: Dict[str, Any] = set_claims or {}
+        self.remove_claims: List[str] = remove_claims or []
+
+        # --- FR-14: Two-token model ---
+        self.channel_token_audience: Optional[str] = channel_token_audience
+        self.channel_token_ttl: int = (
+            channel_token_ttl if channel_token_ttl is not None else self.ttl_seconds
+        )
+
+        # --- FR-15: Incoming claim validation ---
+        self.required_claims: List[str] = required_claims or []
+        self.optional_claims: List[str] = optional_claims or []
+
+        # --- FR-9: Debug headers ---
+        self.debug_headers: bool = debug_headers
+
+        # --- FR-10: Configurable scopes ---
+        self.allowed_scopes: Optional[List[str]] = allowed_scopes
+
+        # Register singleton for JWKS/OIDC discovery endpoints.
+        global _mcp_jwt_signer_instance
+        if _mcp_jwt_signer_instance is not None:
+            verbose_proxy_logger.warning(
+                "MCPJWTSigner: replacing existing singleton — previously issued tokens "
+                "signed with the old key will fail JWKS verification. "
+                "Avoid configuring multiple mcp_jwt_signer guardrails."
+            )
+        _mcp_jwt_signer_instance = self
+
+        verbose_proxy_logger.info(
+            "MCPJWTSigner initialized: issuer=%s audience=%s ttl=%ds kid=%s "
+            "verify=%s channel_token=%s debug=%s",
+            self.issuer,
+            self.audience,
+            self.ttl_seconds,
+            self._kid,
+            bool(self.access_token_discovery_uri),
+            bool(self.channel_token_audience),
+            self.debug_headers,
+        )
+
+    # ------------------------------------------------------------------
+    # Public helpers (used by /.well-known/jwks.json endpoint)
+    # ------------------------------------------------------------------
+
+    @property
+    def jwks_max_age(self) -> int:
+        """
+        Recommended Cache-Control max-age for the JWKS response (seconds).
+
+        1 hour for persistent keys; 5 minutes for auto-generated keys so MCP
+        servers re-fetch quickly after a proxy restart.
+        """
+        return 3600 if self._persistent_key else 300
+
+    def get_jwks(self) -> Dict[str, Any]:
+        """
+        Return the JWKS for the RSA public key.
+        Used by GET /.well-known/jwks.json so MCP servers can verify tokens.
+        """
+        public_numbers = self._public_key.public_numbers()
+        return {
+            "keys": [
+                {
+                    "kty": "RSA",
+                    "alg": self.ALGORITHM,
+                    "use": "sig",
+                    "kid": self._kid,
+                    "n": _int_to_base64url(public_numbers.n),
+                    "e": _int_to_base64url(public_numbers.e),
+                }
+            ]
+        }
+
+    # ------------------------------------------------------------------
+    # FR-5: Verify + re-sign helpers
+    # ------------------------------------------------------------------
+
+    # 24-hour TTL for the OIDC discovery doc — long enough to avoid hammering
+    # the IdP, short enough to pick up jwks_uri changes after key rotation.
+    _OIDC_DISCOVERY_TTL = 86400
+
+    async def _get_oidc_discovery(self) -> Dict[str, Any]:
+        """Fetch and cache the OIDC discovery document with a 24-hour TTL.
+
+        Only caches when the doc contains a 'jwks_uri' so that a transient or
+        malformed response doesn't permanently disable JWT verification.
+        """
+        now = time.time()
+        cache_expired = (
+            now - self._oidc_discovery_fetched_at
+        ) >= self._OIDC_DISCOVERY_TTL
+        if (
+            self._oidc_discovery_doc is None or cache_expired
+        ) and self.access_token_discovery_uri:
+            doc = await _fetch_oidc_discovery(self.access_token_discovery_uri)
+            if "jwks_uri" in doc:
+                self._oidc_discovery_doc = doc
+                self._oidc_discovery_fetched_at = now
+            else:
+                return doc
+        return self._oidc_discovery_doc or {}
+
+    async def _verify_incoming_jwt(self, raw_token: str) -> Dict[str, Any]:
+        """
+        Verify an incoming Bearer JWT against the configured IdP's JWKS.
+
+        Returns the verified payload claims dict.
+        Raises jwt.PyJWTError (or subclass) if verification fails.
+        """
+        discovery = await self._get_oidc_discovery()
+        jwks_uri = discovery.get("jwks_uri")
+        if not jwks_uri:
+            raise ValueError(
+                "MCPJWTSigner: access_token_discovery_uri discovery document "
+                f"at {self.access_token_discovery_uri!r} has no 'jwks_uri'."
+            )
+
+        jwks_keys = await _fetch_jwks(jwks_uri)
+
+        # Only read `kid` from the unverified header — never `alg`.
+        # Reading `alg` from an attacker-controlled header enables algorithm
+        # confusion attacks (e.g. alg:none, HS256 with the public key as secret).
+        # The algorithm is determined from the JWKS key entry instead.
+        unverified_header = jwt.get_unverified_header(raw_token)
+        kid = unverified_header.get("kid")
+
+        # Build a JWKS object and pick the matching key.
+        # PyJWT's PyJWKSet handles key-type parsing and kid matching correctly.
+        from jwt import PyJWKSet
+
+        try:
+            jwks_set = PyJWKSet.from_dict({"keys": jwks_keys})
+        except Exception as exc:
+            raise jwt.exceptions.PyJWKSetError(  # type: ignore[attr-defined]
+                f"Failed to parse JWKS from {jwks_uri!r}: {exc}"
+            ) from exc
+
+        signing_jwk = None
+        for jwk_obj in jwks_set.keys:
+            if not kid or jwk_obj.key_id == kid:
+                signing_jwk = jwk_obj
+                break
+
+        if signing_jwk is None:
+            raise jwt.exceptions.PyJWKSetError(  # type: ignore[attr-defined]
+                f"No JWKS key matching kid={kid!r} at {jwks_uri!r}"
+            )
+
+        # Use the algorithm declared by the JWKS key entry, not the token header.
+        # PyJWT populates algorithm_name from the key's `alg` field; when absent
+        # it infers from the key type (RSAPublicKey → RS256).
+        alg = getattr(signing_jwk, "algorithm_name", None) or "RS256"
+
+        decode_options: Dict[str, Any] = {"verify_exp": True}
+        decode_kwargs: Dict[str, Any] = {
+            "algorithms": [alg],
+            "options": decode_options,
+        }
+        if self.verify_audience:
+            decode_kwargs["audience"] = self.verify_audience
+        else:
+            decode_options["verify_aud"] = False
+
+        if self.verify_issuer:
+            decode_kwargs["issuer"] = self.verify_issuer
+
+        payload: Dict[str, Any] = jwt.decode(
+            raw_token, signing_jwk.key, **decode_kwargs
+        )
+        return payload
+
+    async def _introspect_opaque_token(self, token: str) -> Dict[str, Any]:
+        """
+        Perform RFC 7662 token introspection for opaque (non-JWT) tokens.
+
+        Returns the introspection response dict.  Raises on HTTP error or
+        inactive token.
+        """
+        if not self.token_introspection_endpoint:
+            raise ValueError(
+                "MCPJWTSigner: token_introspection_endpoint is required for "
+                "opaque token verification but is not configured."
+            )
+
+        from litellm.llms.custom_httpx.http_handler import (
+            get_async_httpx_client,
+            httpxSpecialProvider,
+        )
+
+        client = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check)
+        resp = await client.post(
+            self.token_introspection_endpoint,
+            data={"token": token},
+            headers={"Accept": "application/json"},
+        )
+        resp.raise_for_status()
+        result: Dict[str, Any] = resp.json()
+        if not result.get("active", False):
+            raise jwt.exceptions.ExpiredSignatureError(  # type: ignore[attr-defined]
+                "MCPJWTSigner: incoming token is inactive (introspection returned active=false)"
+            )
+        return result
+
+    # ------------------------------------------------------------------
+    # FR-15: Incoming claim validation
+    # ------------------------------------------------------------------
+
+    def _validate_required_claims(
+        self,
+        jwt_claims: Optional[Dict[str, Any]],
+    ) -> None:
+        """
+        Raise HTTP 403 if any required_claims are absent from the verified
+        incoming token claims.
+        """
+        if not self.required_claims:
+            return
+
+        from fastapi import HTTPException
+
+        missing = [c for c in self.required_claims if not (jwt_claims or {}).get(c)]
+        if missing:
+            raise HTTPException(
+                status_code=403,
+                detail={
+                    "error": (
+                        f"MCPJWTSigner: incoming token is missing required claims: "
+                        f"{missing}. Configure the IdP to include these claims."
+                    )
+                },
+            )
+
+    # ------------------------------------------------------------------
+    # FR-12: End-user identity mapping
+    # ------------------------------------------------------------------
+
+    def _resolve_end_user_identity(
+        self,
+        user_api_key_dict: UserAPIKeyAuth,
+        jwt_claims: Optional[Dict[str, Any]],
+    ) -> str:
+        """
+        Resolve the outbound JWT 'sub' using the ordered end_user_claim_sources list.
+
+        Supported source prefixes:
+          token:<claim>       — from verified incoming JWT / introspection claims
+          litellm:user_id     — from UserAPIKeyAuth.user_id
+          litellm:email       — from UserAPIKeyAuth.user_email
+          litellm:end_user_id — from UserAPIKeyAuth.end_user_id
+          litellm:team_id     — from UserAPIKeyAuth.team_id
+
+        Falls back to a stable hash of the API token for service-account callers.
+        """
+        for source in self.end_user_claim_sources:
+            value: Optional[str] = None
+
+            if source.startswith("token:"):
+                claim_name = source[len("token:") :]
+                raw = (jwt_claims or {}).get(claim_name)
+                value = str(raw) if raw else None
+
+            elif source == "litellm:user_id":
+                uid = getattr(user_api_key_dict, "user_id", None)
+                value = str(uid) if uid else None
+
+            elif source == "litellm:email":
+                email = getattr(user_api_key_dict, "user_email", None)
+                value = str(email) if email else None
+
+            elif source == "litellm:end_user_id":
+                eid = getattr(user_api_key_dict, "end_user_id", None)
+                value = str(eid) if eid else None
+
+            elif source == "litellm:team_id":
+                tid = getattr(user_api_key_dict, "team_id", None)
+                value = str(tid) if tid else None
+
+            else:
+                verbose_proxy_logger.warning(
+                    "MCPJWTSigner: unknown end_user_claim_source %r — skipping", source
+                )
+                continue
+
+            if value:
+                return value
+
+        # Final fallback for service accounts with no user identity
+        token = getattr(user_api_key_dict, "token", None) or getattr(
+            user_api_key_dict, "api_key", None
+        )
+        if token:
+            return "apikey:" + hashlib.sha256(str(token).encode()).hexdigest()[:16]
+        return "litellm-proxy"
+
+    # ------------------------------------------------------------------
+    # FR-10: Scope building
+    # ------------------------------------------------------------------
+
+    def _build_scope(self, raw_tool_name: str) -> str:
+        """
+        Build the JWT scope string.
+
+        When allowed_scopes is configured: join them verbatim.
+        Otherwise auto-generate minimal, least-privilege scopes:
+          - Tool call   → mcp:tools/call  mcp:tools/<name>:call
+          - No tool     → mcp:tools/call  mcp:tools/list
+
+        NOTE: tools/list is intentionally NOT granted on tool-call JWTs to
+        prevent callers from enumerating tools they didn't ask to use.
+        """
+        if self.allowed_scopes is not None:
+            return " ".join(self.allowed_scopes)
+
+        tool_name = (
+            re.sub(r"[^a-zA-Z0-9_\-]", "_", raw_tool_name) if raw_tool_name else ""
+        )
+        if tool_name:
+            scopes = ["mcp:tools/call", f"mcp:tools/{tool_name}:call"]
+        else:
+            scopes = ["mcp:tools/call", "mcp:tools/list"]
+        return " ".join(scopes)
+
+    # ------------------------------------------------------------------
+    # FR-13: Claim operations
+    # ------------------------------------------------------------------
+
+    def _apply_claim_operations(self, claims: Dict[str, Any]) -> Dict[str, Any]:
+        """Apply add_claims, set_claims, and remove_claims to the claim dict."""
+        # add_claims: insert only when key is absent
+        for k, v in self.add_claims.items():
+            if k not in claims:
+                claims[k] = v
+
+        # set_claims: always override (highest priority)
+        claims = {**claims, **self.set_claims}
+
+        # remove_claims: delete listed keys
+        for k in self.remove_claims:
+            claims.pop(k, None)
+
+        return claims
+
+    # ------------------------------------------------------------------
+    # FR-15: optional_claims passthrough
+    # ------------------------------------------------------------------
+
+    def _passthrough_optional_claims(
+        self,
+        claims: Dict[str, Any],
+        jwt_claims: Optional[Dict[str, Any]],
+    ) -> Dict[str, Any]:
+        """Forward optional_claims from verified incoming token into the outbound JWT."""
+        if not self.optional_claims or not jwt_claims:
+            return claims
+        for claim in self.optional_claims:
+            if claim in jwt_claims and claim not in claims:
+                claims[claim] = jwt_claims[claim]
+        return claims
+
+    # ------------------------------------------------------------------
+    # Core JWT builder
+    # ------------------------------------------------------------------
+
+    def _build_claims(
+        self,
+        user_api_key_dict: UserAPIKeyAuth,
+        data: dict,
+        jwt_claims: Optional[Dict[str, Any]] = None,
+    ) -> Dict[str, Any]:
+        """
+        Build JWT claims for the outbound MCP access token.
+
+        Args:
+            user_api_key_dict: LiteLLM auth context for the current request.
+            data: Pre-call hook data dict (contains mcp_tool_name etc.).
+            jwt_claims: Verified incoming IdP claims (FR-5), or LiteLLM-decoded
+                        jwt_claims if available.  None for pure API-key requests.
+        """
+        now = int(time.time())
+        claims: Dict[str, Any] = {
+            "iss": self.issuer,
+            "aud": self.audience,
+            "iat": now,
+            "exp": now + self.ttl_seconds,
+            "nbf": now,
+        }
+
+        # sub — resolved via ordered claim sources (FR-12)
+        claims["sub"] = self._resolve_end_user_identity(user_api_key_dict, jwt_claims)
+
+        # email passthrough when available from LiteLLM context
+        user_email = getattr(user_api_key_dict, "user_email", None)
+        if user_email:
+            claims["email"] = user_email
+
+        # act — RFC 8693 delegation claim (team/org context)
+        team_id = getattr(user_api_key_dict, "team_id", None)
+        org_id = getattr(user_api_key_dict, "org_id", None)
+        act_sub = team_id or org_id or "litellm-proxy"
+        claims["act"] = {"sub": act_sub}
+
+        # end_user_id when set separately from user_id
+        end_user_id = getattr(user_api_key_dict, "end_user_id", None)
+        if end_user_id:
+            claims["end_user_id"] = end_user_id
+
+        # scope (FR-10)
+        raw_tool_name: str = data.get("mcp_tool_name", "")
+        claims["scope"] = self._build_scope(raw_tool_name)
+
+        # optional_claims passthrough (FR-15)
+        claims = self._passthrough_optional_claims(claims, jwt_claims)
+
+        # Claim operations — applied last so admin overrides take effect (FR-13)
+        claims = self._apply_claim_operations(claims)
+
+        return claims
+
+    def _build_channel_token_claims(
+        self,
+        base_claims: Dict[str, Any],
+    ) -> Dict[str, Any]:
+        """
+        Build claims for the channel token (FR-14 two-token model).
+
+        Inherits sub/act/scope from the access token but uses a separate
+        audience and TTL so the transport layer and resource layer receive
+        purpose-bound credentials.
+        """
+        now = int(time.time())
+        return {
+            **base_claims,
+            "aud": self.channel_token_audience,
+            "iat": now,
+            "exp": now + self.channel_token_ttl,
+            "nbf": now,
+        }
+
+    # ------------------------------------------------------------------
+    # FR-9: Debug header
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def _build_debug_header(claims: Dict[str, Any], kid: str) -> str:
+        """
+        Build the x-litellm-mcp-debug header value.
+
+        Format: v=1; kid=<kid>; sub=<sub>; iss=<iss>; exp=<exp>; scope=<scope>
+        Scope is truncated to 80 chars for header safety.
+        """
+        sub = claims.get("sub", "")
+        iss = claims.get("iss", "")
+        exp = claims.get("exp", 0)
+        scope = claims.get("scope", "")
+        if len(scope) > 80:
+            scope = scope[:77] + "..."
+        return f"v=1; kid={kid}; sub={sub}; iss={iss}; exp={exp}; scope={scope}"
+
+    # ------------------------------------------------------------------
+    # Guardrail hook
+    # ------------------------------------------------------------------
+
+    @log_guardrail_information
+    async def async_pre_call_hook(
+        self,
+        user_api_key_dict: UserAPIKeyAuth,
+        cache: DualCache,
+        data: dict,
+        call_type: CallTypesLiteral,
+    ) -> Optional[Union[Exception, str, dict]]:
+        """
+        Verifies the incoming token (when configured), validates required claims,
+        then signs an outbound JWT and injects it as the Authorization header.
+
+        All non-MCP call types pass through unchanged.
+        """
+        if call_type != "call_mcp_tool":
+            return data
+
+        # ------------------------------------------------------------------
+        # FR-5: Verify incoming token before re-signing
+        # ------------------------------------------------------------------
+        jwt_claims: Optional[Dict[str, Any]] = None
+        raw_token: Optional[str] = data.get("incoming_bearer_token")
+
+        if self.access_token_discovery_uri and raw_token:
+            # Three-dot pattern → JWT;  otherwise opaque.
+            is_jwt = raw_token.count(".") == 2
+            try:
+                if is_jwt:
+                    jwt_claims = await self._verify_incoming_jwt(raw_token)
+                elif self.token_introspection_endpoint:
+                    jwt_claims = await self._introspect_opaque_token(raw_token)
+                else:
+                    verbose_proxy_logger.warning(
+                        "MCPJWTSigner: access_token_discovery_uri is set but the "
+                        "incoming token appears to be opaque and no "
+                        "token_introspection_endpoint is configured. "
+                        "Proceeding without incoming token verification."
+                    )
+            except Exception as exc:
+                verbose_proxy_logger.error(
+                    "MCPJWTSigner: incoming token verification failed: %s", exc
+                )
+                from fastapi import HTTPException
+
+                raise HTTPException(
+                    status_code=401,
+                    detail={
+                        "error": (
+                            f"MCPJWTSigner: incoming token verification failed: {exc}"
+                        )
+                    },
+                )
+        elif not raw_token and self.access_token_discovery_uri:
+            verbose_proxy_logger.debug(
+                "MCPJWTSigner: access_token_discovery_uri configured but no Bearer "
+                "token found in request (API-key auth request — skipping verification)."
+            )
+
+        # Fall back to LiteLLM-decoded JWT claims (available when proxy uses JWT auth).
+        if jwt_claims is None:
+            jwt_claims = getattr(user_api_key_dict, "jwt_claims", None)
+
+        # ------------------------------------------------------------------
+        # FR-15: Validate required claims
+        # ------------------------------------------------------------------
+        self._validate_required_claims(jwt_claims)
+
+        # ------------------------------------------------------------------
+        # Build outbound access token
+        # ------------------------------------------------------------------
+        claims = self._build_claims(user_api_key_dict, data, jwt_claims)
+
+        signed_token = jwt.encode(
+            claims,
+            self._private_key,
+            algorithm=self.ALGORITHM,
+            headers={"kid": self._kid},
+        )
+
+        # Merge into existing extra_headers — a prior guardrail in the chain may
+        # have already injected tracing headers or correlation IDs.
+        existing_headers: Dict[str, str] = data.get("extra_headers") or {}
+        new_headers: Dict[str, str] = {
+            **existing_headers,
+            "Authorization": f"Bearer {signed_token}",
+        }
+
+        # ------------------------------------------------------------------
+        # FR-14: Two-token model — channel token
+        # ------------------------------------------------------------------
+        if self.channel_token_audience:
+            channel_claims = self._build_channel_token_claims(claims)
+            channel_token = jwt.encode(
+                channel_claims,
+                self._private_key,
+                algorithm=self.ALGORITHM,
+                headers={"kid": self._kid},
+            )
+            new_headers["x-mcp-channel-token"] = f"Bearer {channel_token}"
+
+        # ------------------------------------------------------------------
+        # FR-9: Debug header
+        # ------------------------------------------------------------------
+        if self.debug_headers:
+            new_headers["x-litellm-mcp-debug"] = self._build_debug_header(
+                claims, self._kid
+            )
+
+        data["extra_headers"] = new_headers
+
+        verbose_proxy_logger.debug(
+            "MCPJWTSigner: signed JWT sub=%s act=%s tool=%s exp=%d "
+            "verified=%s channel=%s",
+            claims.get("sub"),
+            claims.get("act", {}).get("sub"),
+            data.get("mcp_tool_name"),
+            claims["exp"],
+            jwt_claims is not None,
+            bool(self.channel_token_audience),
+        )
+
+        return data

litellm/proxy/management_endpoints/internal_user_endpoints.py

@@ -2142,8 +2142,7 @@ async def _resolve_org_filter_for_user_search(
     member_org_ids: List[str] = []
     if caller_user is not None:
         member_org_ids = [
-            m.organization_id
-            for m in (caller_user.organization_memberships or [])
+            m.organization_id for m in (caller_user.organization_memberships or [])
         ]
 
     if member_org_ids:

litellm/proxy/management_endpoints/key_management_endpoints.py

@@ -1863,16 +1863,10 @@ async def _validate_update_key_data(
     user_api_key_cache: Any,
 ) -> None:
     """Validate permissions and constraints for key update."""
-    _is_proxy_admin = (
-        user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
-    )
+    _is_proxy_admin = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
 
     # Prevent non-admin from removing user_id (setting to empty string) (LIT-1884)
-    if (
-        data.user_id is not None
-        and data.user_id == ""
-        and not _is_proxy_admin
-    ):
+    if data.user_id is not None and data.user_id == "" and not _is_proxy_admin:
         raise HTTPException(
             status_code=403,
             detail="Non-admin users cannot remove the user_id from a key.",

litellm/proxy/management_endpoints/team_endpoints.py

@@ -101,6 +101,7 @@ from litellm.types.proxy.management_endpoints.team_endpoints import (
     BulkTeamMemberAddRequest,
     BulkTeamMemberAddResponse,
     GetTeamMemberPermissionsResponse,
+    TeamListItem,
     TeamListResponse,
     TeamMemberAddResult,
     UpdateTeamMemberPermissionsRequest,
@@ -857,7 +858,13 @@ async def new_team(  # noqa: PLR0915
 
         # Apply defaults from litellm.default_team_params for any fields
         # not explicitly provided in the request.
-        for field in ("max_budget", "budget_duration", "tpm_limit", "rpm_limit", "team_member_permissions"):
+        for field in (
+            "max_budget",
+            "budget_duration",
+            "tpm_limit",
+            "rpm_limit",
+            "team_member_permissions",
+        ):
             if getattr(data, field, None) is None:
                 default_value = _get_default_team_param(field)
                 if default_value is not None:
@@ -3206,6 +3213,40 @@ async def list_available_teams(
     return available_teams_correct_type
 
 
+async def _get_org_admin_org_ids(
+    user_id: str,
+    prisma_client: Any,
+    user_api_key_cache: Any,
+    proxy_logging_obj: Any,
+) -> Optional[List[str]]:
+    """
+    Return the list of organization IDs where the user is an org admin.
+    Returns None if the user is not an org admin of any organization or if
+    the user cannot be found.
+    """
+    try:
+        caller_user = await get_user_object(
+            user_id=user_id,
+            prisma_client=prisma_client,
+            user_api_key_cache=user_api_key_cache,
+            user_id_upsert=False,
+            proxy_logging_obj=proxy_logging_obj,
+        )
+    except ValueError:
+        # get_user_object raises ValueError when the user doesn't exist
+        return None
+
+    if caller_user is None:
+        return None
+
+    org_ids = [
+        m.organization_id
+        for m in (caller_user.organization_memberships or [])
+        if m.user_role == LitellmUserRoles.ORG_ADMIN.value
+    ]
+    return org_ids if org_ids else None
+
+
 async def _build_team_list_where_conditions(
     prisma_client: PrismaClient,
     team_id: Optional[str],
@@ -3213,8 +3254,16 @@ async def _build_team_list_where_conditions(
     organization_id: Optional[str],
     user_id: Optional[str],
     use_deleted_table: bool,
-) -> Dict[str, Any]:
-    """Build where conditions for team list query."""
+    org_admin_org_ids: Optional[List[str]] = None,
+    user_api_key_cache: Optional[Any] = None,
+    proxy_logging_obj: Optional[Any] = None,
+) -> Optional[Dict[str, Any]]:
+    """
+    Build where conditions for team list query.
+
+    Returns None when the query is guaranteed to yield no results (e.g. user
+    has no team memberships), allowing the caller to skip the DB round-trip.
+    """
     where_conditions: Dict[str, Any] = {}
 
     if team_id:
@@ -3228,58 +3277,79 @@ async def _build_team_list_where_conditions(
 
     if organization_id:
         where_conditions["organization_id"] = organization_id
+    elif org_admin_org_ids is not None and not user_id:
+        # Org admin without explicit org or user filter: scope to their orgs.
+        # NOTE: when user_id is provided, no org filter is applied — the
+        # query returns all teams the target user belongs to across all
+        # organisations.  This matches the legacy /team/list behaviour in
+        # _authorize_and_filter_teams which fetches direct-membership teams
+        # without an org constraint.
+        where_conditions["organization_id"] = {"in": org_admin_org_ids}
 
     if user_id:
         try:
-            user_object = await prisma_client.db.litellm_usertable.find_unique(
-                where={"user_id": user_id}
+            user_object_correct_type = await get_user_object(
+                user_id=user_id,
+                prisma_client=prisma_client,
+                user_api_key_cache=user_api_key_cache,
+                user_id_upsert=False,
+                proxy_logging_obj=proxy_logging_obj,
             )
-        except Exception:
+        except ValueError:
             raise HTTPException(
                 status_code=404,
                 detail={"error": f"User not found, passed user_id={user_id}"},
             )
-        if user_object is None:
+        if user_object_correct_type is None:
             raise HTTPException(
                 status_code=404,
                 detail={"error": f"User not found, passed user_id={user_id}"},
             )
-        user_object_correct_type = LiteLLM_UserTable(**user_object.model_dump())
+        user_team_ids = user_object_correct_type.teams or []
 
         if use_deleted_table:
             where_conditions["members"] = {"has": user_id}
         else:
-            if team_id is None:
-                where_conditions["team_id"] = {"in": user_object_correct_type.teams}
-            elif team_id in user_object_correct_type.teams:
-                where_conditions["team_id"] = team_id
+            # When user_id is provided, filter by that user's direct team
+            # memberships. For org admins the access control gate in
+            # list_team_v2 already verified the caller's authority — the
+            # filter logic is the same as for regular users.
+            if not user_team_ids:
+                return None  # no memberships — skip the DB query
+            elif team_id is not None:
+                # team_id exact-match already in where_conditions; verify membership
+                if team_id not in user_team_ids:
+                    raise HTTPException(
+                        status_code=404,
+                        detail={"error": f"User is not a member of team_id={team_id}"},
+                    )
             else:
-                raise HTTPException(
-                    status_code=404,
-                    detail={"error": f"User is not a member of team_id={team_id}"},
-                )
+                where_conditions["team_id"] = {"in": user_team_ids}
 
     return where_conditions
 
 
-def _convert_teams_to_response(
-    teams: List[Any], use_deleted_table: bool
-) -> List[Union[LiteLLM_TeamTable, LiteLLM_DeletedTeamTable]]:
-    """Convert Prisma models to Pydantic models."""
-    team_list: List[Union[LiteLLM_TeamTable, LiteLLM_DeletedTeamTable]] = []
-    if teams:
-        for team in teams:
-            # Convert Prisma model to dict (supports both Pydantic v1 and v2)
-            try:
-                team_dict = team.model_dump()
-            except Exception:
-                # Fallback for Pydantic v1 compatibility
-                team_dict = team.dict()
-            if use_deleted_table:
-                # Use deleted team type to preserve deleted_at, deleted_by, etc.
-                team_list.append(LiteLLM_DeletedTeamTable(**team_dict))
-            else:
-                team_list.append(LiteLLM_TeamTable(**team_dict))
+def _convert_teams_to_response_models(
+    teams: list,
+    use_deleted_table: bool,
+) -> List[Union[TeamListItem, LiteLLM_TeamTable, LiteLLM_DeletedTeamTable]]:
+    """Convert raw Prisma team rows to response models."""
+    team_list: List[Union[TeamListItem, LiteLLM_TeamTable, LiteLLM_DeletedTeamTable]] = []
+    for team in teams:
+        try:
+            team_dict = team.model_dump()
+        except Exception:
+            team_dict = team.dict()
+
+        if use_deleted_table:
+            team_list.append(LiteLLM_DeletedTeamTable(**team_dict))
+        else:
+            members_with_roles = team_dict.get("members_with_roles")
+            if not isinstance(members_with_roles, list):
+                members_with_roles = []
+                team_dict["members_with_roles"] = members_with_roles
+            members_count = len(members_with_roles)
+            team_list.append(TeamListItem(**team_dict, members_count=members_count))
     return team_list
 
 
@@ -3347,7 +3417,11 @@ async def list_team_v2(
         status: Optional[str]
             Filter by status. Currently supports "deleted" to query deleted teams.
     """
-    from litellm.proxy.proxy_server import prisma_client
+    from litellm.proxy.proxy_server import (
+        prisma_client,
+        proxy_logging_obj,
+        user_api_key_cache,
+    )
 
     if prisma_client is None:
         raise HTTPException(
@@ -3355,20 +3429,55 @@ async def list_team_v2(
             detail={"error": f"No db connected. prisma client={prisma_client}"},
         )
 
-    if not allowed_route_check_inside_route(
-        user_api_key_dict=user_api_key_dict, requested_user_id=user_id
-    ):
-        raise HTTPException(
-            status_code=401,
-            detail={
-                "error": "Only admin users can query all teams/other teams. Your user role={}".format(
-                    user_api_key_dict.user_role
-                )
-            },
-        )
+    # --- Access control ---
+    # Proxy admins and admin viewers can query any teams.
+    # Org admins can query teams within their organizations.
+    # Regular users can only query their own teams.
+    is_proxy_admin = _user_has_admin_view(user_api_key_dict)
+    org_admin_org_ids: Optional[List[str]] = None
 
-    if user_id is None and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
-        user_id = user_api_key_dict.user_id
+    if not is_proxy_admin:
+        # Always check org admin status so that even own-queries see
+        # the full set of organisation teams, not just direct memberships.
+        if user_api_key_dict.user_id:
+            org_admin_org_ids = await _get_org_admin_org_ids(
+                user_id=user_api_key_dict.user_id,
+                prisma_client=prisma_client,
+                user_api_key_cache=user_api_key_cache,
+                proxy_logging_obj=proxy_logging_obj,
+            )
+
+        if org_admin_org_ids is not None:
+            # Org admin: validate org_id filter if provided
+            if organization_id and organization_id not in org_admin_org_ids:
+                raise HTTPException(
+                    status_code=403,
+                    detail={
+                        "error": "You can only view teams within your organizations."
+                    },
+                )
+            verbose_proxy_logger.debug(
+                "list_team_v2: org admin access for user=%s, org_ids=%s, user_id_filter=%s",
+                user_api_key_dict.user_id,
+                org_admin_org_ids,
+                user_id,
+            )
+        else:
+            # Not an org admin — fall back to standard route check
+            if not allowed_route_check_inside_route(
+                user_api_key_dict=user_api_key_dict, requested_user_id=user_id
+            ):
+                raise HTTPException(
+                    status_code=401,
+                    detail={
+                        "error": "Only admin users can query all teams/other teams. Your user role={}".format(
+                            user_api_key_dict.user_role
+                        )
+                    },
+                )
+            # Regular user — auto-inject caller's user_id
+            if user_id is None:
+                user_id = user_api_key_dict.user_id
 
     if status is not None and status != "deleted":
         raise HTTPException(
@@ -3383,7 +3492,8 @@ async def list_team_v2(
     # Calculate skip and take for pagination
     skip = (page - 1) * page_size
 
-    # Build where conditions based on provided parameters
+    # Build where conditions based on provided parameters.
+    # Returns None when the query is guaranteed to yield no results.
     where_conditions = await _build_team_list_where_conditions(
         prisma_client=prisma_client,
         team_id=team_id,
@@ -3391,8 +3501,20 @@ async def list_team_v2(
         organization_id=organization_id,
         user_id=user_id,
         use_deleted_table=use_deleted_table,
+        org_admin_org_ids=org_admin_org_ids,
+        user_api_key_cache=user_api_key_cache,
+        proxy_logging_obj=proxy_logging_obj,
     )
 
+    if where_conditions is None:
+        return {
+            "teams": [],
+            "total": 0,
+            "page": page,
+            "page_size": page_size,
+            "total_pages": 0,
+        }
+
     # Build order_by conditions
     valid_sort_columns = ["team_id", "team_alias", "created_at"]
     order_by = None
@@ -3428,8 +3550,8 @@ async def list_team_v2(
     # Calculate total pages
     total_pages = -(-total_count // page_size)  # Ceiling division
 
-    # Convert Prisma models to Pydantic models, preserving deleted fields when applicable
-    team_list = _convert_teams_to_response(teams, use_deleted_table)
+    # Convert Prisma models to response models with members_count
+    team_list = _convert_teams_to_response_models(teams, use_deleted_table)
 
     return {
         "teams": team_list,

litellm/proxy/openai_files_endpoints/common_utils.py

@@ -857,7 +857,10 @@ async def update_batch_in_database(
             # If the batch_processed column doesn't exist (old schema),
             # retry without it so the status update still succeeds.
             err_str = str(col_err).lower()
-            if "batch_processed" in err_str and update_data.get("batch_processed") is not None:
+            if (
+                "batch_processed" in err_str
+                and update_data.get("batch_processed") is not None
+            ):
                 verbose_proxy_logger.warning(
                     f"batch_processed column not found, retrying update without it: {col_err}"
                 )

litellm/proxy/proxy_cli.py

@@ -468,6 +468,12 @@ class ProxyInitializationHelpers:
     type=str,
     help="Path to the logging configuration file",
 )
+@click.option(
+    "--setup",
+    is_flag=True,
+    default=False,
+    help="Run the interactive setup wizard to configure providers and generate a config file",
+)
 @click.option(
     "--version",
     "-v",
@@ -598,6 +604,7 @@ def run_server(  # noqa: PLR0915
     num_requests,
     use_queue,
     health,
+    setup,
     version,
     run_gunicorn,
     run_hypercorn,
@@ -611,6 +618,12 @@ def run_server(  # noqa: PLR0915
     max_requests_before_restart,
     enforce_prisma_migration_check: bool,
 ):
+    if setup:
+        from litellm.setup_wizard import run_setup_wizard
+
+        run_setup_wizard()
+        return
+
     args = locals()
     if local:
         from proxy_server import (
@@ -904,7 +917,7 @@ def run_server(  # noqa: PLR0915
         # Auto-create PROMETHEUS_MULTIPROC_DIR for multi-worker setups
         ProxyInitializationHelpers._maybe_setup_prometheus_multiproc_dir(
             num_workers=num_workers,
-            litellm_settings=litellm_settings if config else None,
+            litellm_settings=litellm_settings if config else None,  # type: ignore[possibly-unbound]
         )
 
         # --- SEPARATE HEALTH APP LOGIC ---

litellm/proxy/response_polling/background_streaming.py

@@ -115,7 +115,9 @@ async def background_streaming_task(  # noqa: PLR0915
         UPDATE_INTERVAL = 0.150  # 150ms batching interval
 
         # Track the terminal event from the stream (may not be "completed")
-        terminal_status: Optional[ResponsesAPIStatus] = None  # Will be set by response.completed/failed/incomplete/cancelled
+        terminal_status: Optional[
+            ResponsesAPIStatus
+        ] = None  # Will be set by response.completed/failed/incomplete/cancelled
         terminal_error = None
         _event_to_status = {
             "response.completed": "completed",
@@ -259,7 +261,10 @@ async def background_streaming_task(  # noqa: PLR0915
                             )
 
                             # Extract error for failed and incomplete responses
-                            if event_type == "response.failed" or event_type == "response.incomplete":
+                            if (
+                                event_type == "response.failed"
+                                or event_type == "response.incomplete"
+                            ):
                                 terminal_error = response_data.get("error")
 
                             # Core response fields

litellm/proxy/schema.prisma

@@ -146,6 +146,10 @@ model LiteLLM_TeamTable {
     litellm_model_table LiteLLM_ModelTable? @relation(fields: [model_id], references: [id])
     object_permission LiteLLM_ObjectPermissionTable?   @relation(fields: [object_permission_id], references: [object_permission_id])
     projects LiteLLM_ProjectTable[]
+
+    @@index([organization_id])
+    @@index([team_alias])
+    @@index([created_at])
 }
 
 // Projects sit between teams and keys for use-case management

litellm/proxy/spend_tracking/vantage_endpoints.py

@@ -40,9 +40,7 @@ def _get_registered_vantage_logger():
     return None
 
 
-async def _set_vantage_settings(
-    api_key: str, integration_token: str, base_url: str
-):
+async def _set_vantage_settings(api_key: str, integration_token: str, base_url: str):
     """Store Vantage settings in the database with encrypted API key."""
     from litellm.proxy.proxy_server import prisma_client
 
@@ -341,9 +339,7 @@ async def init_vantage_settings(
     except HTTPException:
         raise
     except Exception as e:
-        verbose_proxy_logger.error(
-            f"Error initializing Vantage settings: {str(e)}"
-        )
+        verbose_proxy_logger.error(f"Error initializing Vantage settings: {str(e)}")
         raise HTTPException(
             status_code=500,
             detail={"error": f"Failed to initialize Vantage settings: {str(e)}"},
@@ -395,7 +391,8 @@ async def vantage_dry_run_export(
             """Cast Decimal columns to Float64 so .to_dicts() produces
             JSON-serializable float values instead of decimal.Decimal."""
             decimal_cols = [
-                col for col, dtype in zip(frame.columns, frame.dtypes)
+                col
+                for col, dtype in zip(frame.columns, frame.dtypes)
                 if isinstance(dtype, pl.Decimal)
             ]
             if decimal_cols:
@@ -404,8 +401,16 @@ async def vantage_dry_run_export(
                 )
             return frame.to_dicts()
 
-        usage_sample = _to_json_safe_dicts(data.head(min(50, len(data)))) if not data.is_empty() else []
-        normalized_sample = _to_json_safe_dicts(normalized.head(min(50, len(normalized)))) if not normalized.is_empty() else []
+        usage_sample = (
+            _to_json_safe_dicts(data.head(min(50, len(data))))
+            if not data.is_empty()
+            else []
+        )
+        normalized_sample = (
+            _to_json_safe_dicts(normalized.head(min(50, len(normalized))))
+            if not normalized.is_empty()
+            else []
+        )
 
         # Use the same pre-transform column names as
         # FocusExportEngine.dry_run_export_usage_data for consistency.
@@ -437,14 +442,10 @@ async def vantage_dry_run_export(
     except HTTPException:
         raise
     except Exception as e:
-        verbose_proxy_logger.error(
-            f"Error performing Vantage dry run export: {str(e)}"
-        )
+        verbose_proxy_logger.error(f"Error performing Vantage dry run export: {str(e)}")
         raise HTTPException(
             status_code=500,
-            detail={
-                "error": f"Failed to perform Vantage dry run export: {str(e)}"
-            },
+            detail={"error": f"Failed to perform Vantage dry run export: {str(e)}"},
         )
 
 

litellm/proxy/utils.py

@@ -454,8 +454,6 @@ class ProxyLogging:
 
         for hook in PROXY_HOOKS:
             proxy_hook = get_proxy_hook(hook)
-            import inspect
-
             expected_args = inspect.getfullargspec(proxy_hook).args
             passed_in_args: Dict[str, Any] = {}
             if "internal_usage_cache" in expected_args:
@@ -559,6 +557,10 @@ class ProxyLogging:
             "user_api_key_request_route": kwargs.get("user_api_key_request_route"),
             "mcp_tool_name": request_obj.tool_name,  # Keep original for reference
             "mcp_arguments": request_obj.arguments,  # Keep original for reference
+            # Raw Bearer token from the original HTTP request — allows guardrails
+            # (e.g. MCPJWTSigner) to independently verify the caller's identity
+            # before re-signing an outbound token (FR-5 verify+re-sign).
+            "incoming_bearer_token": kwargs.get("incoming_bearer_token"),
         }
 
         return synthetic_data
@@ -824,17 +826,30 @@ class ProxyLogging:
     ) -> dict:
         """
         Helper function to convert pre_call_hook response back to kwargs for MCP usage.
+
+        Supports:
+        - modified_arguments: Override tool call arguments
+        - extra_headers: Inject custom headers into the outbound MCP request
         """
         if not response_data:
             return original_kwargs
 
-        # Apply any argument modifications from the hook response
         modified_kwargs = original_kwargs.copy()
 
-        # If the response contains modified arguments, apply them
         if response_data.get("modified_arguments"):
             modified_kwargs["arguments"] = response_data["modified_arguments"]
 
+        if response_data.get("extra_headers"):
+            # Merge rather than replace — a prior guardrail in the chain may have
+            # already injected headers (e.g. tracing IDs).  Later guardrails win on
+            # key collisions so that the most-specific guardrail (e.g. JWT signer)
+            # takes precedence over earlier ones.
+            existing = modified_kwargs.get("extra_headers") or {}
+            modified_kwargs["extra_headers"] = {
+                **existing,
+                **response_data["extra_headers"],
+            }
+
         return modified_kwargs
 
     async def process_pre_call_hook_response(self, response, data, call_type):

litellm/proxy/video_endpoints/utils.py

@@ -7,7 +7,9 @@ from litellm.types.videos.utils import encode_character_id_with_provider
 
 def extract_model_from_target_model_names(target_model_names: Any) -> Optional[str]:
     if isinstance(target_model_names, str):
-        target_model_names = [m.strip() for m in target_model_names.split(",") if m.strip()]
+        target_model_names = [
+            m.strip() for m in target_model_names.split(",") if m.strip()
+        ]
     elif not isinstance(target_model_names, list):
         return None
     return target_model_names[0] if target_model_names else None

litellm/responses/main.py

@@ -692,11 +692,11 @@ def responses(
             return run_async_function(aresponses_api_with_mcp, **mcp_call_kwargs)
 
         # get provider config
-        responses_api_provider_config: Optional[BaseResponsesAPIConfig] = (
-            ProviderConfigManager.get_provider_responses_api_config(
-                model=model,
-                provider=custom_llm_provider,
-            )
+        responses_api_provider_config: Optional[
+            BaseResponsesAPIConfig
+        ] = ProviderConfigManager.get_provider_responses_api_config(
+            model=model,
+            provider=custom_llm_provider,
         )
 
         local_vars.update(kwargs)
@@ -908,11 +908,11 @@ def delete_responses(
             raise ValueError("custom_llm_provider is required but passed as None")
 
         # get provider config
-        responses_api_provider_config: Optional[BaseResponsesAPIConfig] = (
-            ProviderConfigManager.get_provider_responses_api_config(
-                model=None,
-                provider=custom_llm_provider,
-            )
+        responses_api_provider_config: Optional[
+            BaseResponsesAPIConfig
+        ] = ProviderConfigManager.get_provider_responses_api_config(
+            model=None,
+            provider=custom_llm_provider,
         )
 
         if responses_api_provider_config is None:
@@ -1089,11 +1089,11 @@ def get_responses(
             raise ValueError("custom_llm_provider is required but passed as None")
 
         # get provider config
-        responses_api_provider_config: Optional[BaseResponsesAPIConfig] = (
-            ProviderConfigManager.get_provider_responses_api_config(
-                model=None,
-                provider=custom_llm_provider,
-            )
+        responses_api_provider_config: Optional[
+            BaseResponsesAPIConfig
+        ] = ProviderConfigManager.get_provider_responses_api_config(
+            model=None,
+            provider=custom_llm_provider,
         )
 
         if responses_api_provider_config is None:
@@ -1247,11 +1247,11 @@ def list_input_items(
         if custom_llm_provider is None:
             raise ValueError("custom_llm_provider is required but passed as None")
 
-        responses_api_provider_config: Optional[BaseResponsesAPIConfig] = (
-            ProviderConfigManager.get_provider_responses_api_config(
-                model=None,
-                provider=custom_llm_provider,
-            )
+        responses_api_provider_config: Optional[
+            BaseResponsesAPIConfig
+        ] = ProviderConfigManager.get_provider_responses_api_config(
+            model=None,
+            provider=custom_llm_provider,
         )
 
         if responses_api_provider_config is None:
@@ -1406,11 +1406,11 @@ def cancel_responses(
             raise ValueError("custom_llm_provider is required but passed as None")
 
         # get provider config
-        responses_api_provider_config: Optional[BaseResponsesAPIConfig] = (
-            ProviderConfigManager.get_provider_responses_api_config(
-                model=None,
-                provider=custom_llm_provider,
-            )
+        responses_api_provider_config: Optional[
+            BaseResponsesAPIConfig
+        ] = ProviderConfigManager.get_provider_responses_api_config(
+            model=None,
+            provider=custom_llm_provider,
         )
 
         if responses_api_provider_config is None:
@@ -1594,11 +1594,11 @@ def compact_responses(
             raise ValueError("custom_llm_provider is required but passed as None")
 
         # get provider config
-        responses_api_provider_config: Optional[BaseResponsesAPIConfig] = (
-            ProviderConfigManager.get_provider_responses_api_config(
-                model=model,
-                provider=custom_llm_provider,
-            )
+        responses_api_provider_config: Optional[
+            BaseResponsesAPIConfig
+        ] = ProviderConfigManager.get_provider_responses_api_config(
+            model=model,
+            provider=custom_llm_provider,
         )
 
         if responses_api_provider_config is None:

litellm/router.py

@@ -8611,6 +8611,7 @@ class Router:
             _model_info = deployment.get("model_info", {})
 
             # see if we have the info for this model
+            _deployment_model = None  # per-deployment model name (avoids overwriting the outer `model` group name)
             try:
                 base_model = _model_info.get("base_model", None)
                 if base_model is None:
@@ -8618,7 +8619,7 @@ class Router:
                 model_info = self.get_router_model_info(
                     deployment=deployment, received_model_name=model
                 )
-                model = base_model or _litellm_params.get("model", None)
+                _deployment_model = base_model or _litellm_params.get("model", None)
 
                 if (
                     isinstance(model_info, dict)
@@ -8632,7 +8633,9 @@ class Router:
                         _context_window_error = True
                         _potential_error_str += (
                             "Model={}, Max Input Tokens={}, Got={}".format(
-                                model, model_info["max_input_tokens"], input_tokens
+                                _deployment_model,
+                                model_info["max_input_tokens"],
+                                input_tokens,
                             )
                         )
                         continue
@@ -8688,13 +8691,21 @@ class Router:
 
             ## INVALID PARAMS ## -> catch 'gpt-3.5-turbo-16k' not supporting 'response_format' param
             if request_kwargs is not None and litellm.drop_params is False:
-                # get supported params
-                model, custom_llm_provider, _, _ = litellm.get_llm_provider(
-                    model=model, litellm_params=LiteLLM_Params(**_litellm_params)
+                # get supported params — use per-deployment model to avoid overwriting the outer model group name
+                _dep_model_for_params = _deployment_model or model
+                (
+                    _dep_model_for_params,
+                    custom_llm_provider,
+                    _,
+                    _,
+                ) = litellm.get_llm_provider(
+                    model=_dep_model_for_params,
+                    litellm_params=LiteLLM_Params(**_litellm_params),
                 )
 
                 supported_openai_params = litellm.get_supported_openai_params(
-                    model=model, custom_llm_provider=custom_llm_provider
+                    model=_dep_model_for_params,
+                    custom_llm_provider=custom_llm_provider,
                 )
 
                 if supported_openai_params is None:

litellm/setup_wizard.py

@@ -0,0 +1,668 @@
+# ruff: noqa: T201
+# flake8: noqa: T201
+"""
+LiteLLM Interactive Setup Wizard
+
+Guides users through selecting LLM providers, entering API keys,
+and generating a proxy config file — mirroring the Claude Code onboarding UX.
+"""
+
+import importlib.metadata
+import os
+import re
+import secrets
+import sys
+import sysconfig
+from pathlib import Path
+from typing import Dict, List, Optional, Set
+
+# termios / tty are Unix-only; fall back gracefully on Windows
+try:
+    import termios
+    import tty
+
+    _HAS_RAW_TERMINAL: bool = True
+except ImportError:
+    termios = None  # type: ignore[assignment]
+    tty = None  # type: ignore[assignment]
+    _HAS_RAW_TERMINAL = False
+
+from litellm.utils import check_valid_key
+
+# ---------------------------------------------------------------------------
+# Provider definitions
+# ---------------------------------------------------------------------------
+# Each entry describes one provider card shown in the wizard.
+# `env_key`      — primary env var name (None = no key needed, e.g. Ollama)
+# `test_model`   — model passed to check_valid_key for credential validation
+#                  (None = skip validation, e.g. Azure needs a deployment name)
+# `models`       — default models written into the generated config
+# ---------------------------------------------------------------------------
+
+PROVIDERS: List[Dict] = [
+    {
+        "id": "openai",
+        "name": "OpenAI",
+        "description": "GPT-4o, GPT-4o-mini, o3-mini",
+        "env_key": "OPENAI_API_KEY",
+        "key_hint": "sk-...",
+        "test_model": "gpt-4o-mini",
+        "models": ["gpt-4o", "gpt-4o-mini"],
+    },
+    {
+        "id": "anthropic",
+        "name": "Anthropic",
+        "description": "Claude Opus 4.6, Sonnet 4.6, Haiku 4.5",
+        "env_key": "ANTHROPIC_API_KEY",
+        "key_hint": "sk-ant-...",
+        "test_model": "claude-haiku-4-5-20251001",
+        "models": ["claude-opus-4-6", "claude-sonnet-4-6", "claude-haiku-4-5-20251001"],
+    },
+    {
+        "id": "gemini",
+        "name": "Google Gemini",
+        "description": "Gemini 2.0 Flash, Gemini 2.5 Pro",
+        "env_key": "GEMINI_API_KEY",
+        "key_hint": "AIza...",
+        "test_model": "gemini/gemini-2.0-flash",
+        "models": ["gemini/gemini-2.0-flash", "gemini/gemini-2.5-pro"],
+    },
+    {
+        "id": "azure",
+        "name": "Azure OpenAI",
+        "description": "GPT-4o via Azure",
+        "env_key": "AZURE_API_KEY",
+        "key_hint": "your-azure-key",
+        "test_model": None,  # needs deployment name — skip validation
+        "models": [],
+        "needs_api_base": True,
+        "api_base_hint": "https://<resource>.openai.azure.com/",
+        "api_version": "2024-07-01-preview",
+    },
+    {
+        "id": "bedrock",
+        "name": "AWS Bedrock",
+        "description": "Claude 3.5, Llama 3 via AWS",
+        "env_key": "AWS_ACCESS_KEY_ID",
+        "key_hint": "AKIA...",
+        "test_model": None,  # multi-key auth — skip validation
+        "models": ["bedrock/anthropic.claude-3-5-sonnet-20241022-v2:0"],
+        "extra_keys": ["AWS_SECRET_ACCESS_KEY", "AWS_REGION_NAME"],
+        "extra_hints": ["your-secret-key", "us-east-1"],
+    },
+    {
+        "id": "ollama",
+        "name": "Ollama",
+        "description": "Local models (llama3.2, mistral, etc.)",
+        "env_key": None,
+        "key_hint": None,
+        "test_model": None,  # local — no remote validation
+        "models": ["ollama/llama3.2", "ollama/mistral"],
+        "api_base": "http://localhost:11434",
+    },
+]
+
+
+# ---------------------------------------------------------------------------
+# ANSI colour helpers
+# ---------------------------------------------------------------------------
+
+_ANSI_RE = re.compile(r"\033\[[^m]*m")
+
+_ORANGE = "\033[38;2;215;119;87m"
+_DIM = "\033[2m"
+_BOLD = "\033[1m"
+_GREEN = "\033[38;2;78;186;101m"
+_BLUE = "\033[38;2;177;185;249m"
+_GREY = "\033[38;2;153;153;153m"
+_RESET = "\033[0m"
+_CHECK = "✔"
+_CROSS = "✘"
+
+_CURSOR_HIDE = "\033[?25l"
+_CURSOR_SHOW = "\033[?25h"
+_MOVE_UP = "\033[{}A"
+
+
+def _supports_color() -> bool:
+    return sys.stdout.isatty() and os.environ.get("NO_COLOR") is None
+
+
+def _c(code: str, text: str) -> str:
+    return f"{code}{text}{_RESET}" if _supports_color() else text
+
+
+def orange(t: str) -> str:
+    return _c(_ORANGE, t)
+
+
+def bold(t: str) -> str:
+    return _c(_BOLD, t)
+
+
+def green(t: str) -> str:
+    return _c(_GREEN, t)
+
+
+def blue(t: str) -> str:
+    return _c(_BLUE, t)
+
+
+def grey(t: str) -> str:
+    return _c(_GREY, t)
+
+
+def dim(t: str) -> str:
+    return _c(_DIM, t)
+
+
+def _divider() -> str:
+    """Return a styled divider line (evaluated at call-time, not import-time)."""
+    return dim("  " + "╌" * 74)
+
+
+def _styled_input(prompt: str) -> str:
+    """
+    Like input() but wraps ANSI sequences in readline ignore markers
+    (\\001...\\002) so readline correctly tracks the cursor column.
+    In non-TTY contexts, strips ANSI entirely so no escape codes appear.
+    """
+    if sys.stdout.isatty():
+        rl_prompt = _ANSI_RE.sub(lambda m: f"\001{m.group()}\002", prompt)
+    else:
+        rl_prompt = _ANSI_RE.sub("", prompt)
+    return input(rl_prompt).strip()
+
+
+def _yaml_escape(value: str) -> str:
+    """Escape a string for safe embedding in a double-quoted YAML scalar."""
+    return (
+        value.replace("\\", "\\\\")
+        .replace('"', '\\"')
+        .replace("\n", "\\n")
+        .replace("\r", "\\r")
+        .replace("\t", "\\t")
+    )
+
+
+# ---------------------------------------------------------------------------
+# Layout constants
+# ---------------------------------------------------------------------------
+
+LITELLM_ASCII = r"""
+  ██╗     ██╗████████╗███████╗██╗     ██╗     ███╗   ███╗
+  ██║     ██║╚══██╔══╝██╔════╝██║     ██║     ████╗ ████║
+  ██║     ██║   ██║   █████╗  ██║     ██║     ██╔████╔██║
+  ██║     ██║   ██║   ██╔══╝  ██║     ██║     ██║╚██╔╝██║
+  ███████╗██║   ██║   ███████╗███████╗███████╗██║ ╚═╝ ██║
+  ╚══════╝╚═╝   ╚═╝   ╚══════╝╚══════╝╚══════╝╚═╝     ╚═╝
+"""
+
+
+# ---------------------------------------------------------------------------
+# Setup wizard
+# ---------------------------------------------------------------------------
+
+
+class SetupWizard:
+    """
+    Interactive onboarding wizard: provider selection → API keys → config file.
+
+    All methods are static — the class is purely a namespace with clear
+    single-responsibility sections. Entry point: SetupWizard.run().
+    """
+
+    # ── entry point ─────────────────────────────────────────────────────────
+
+    @staticmethod
+    def run() -> None:
+        try:
+            SetupWizard._wizard()
+        except (KeyboardInterrupt, EOFError):
+            print(f"\n\n  {grey('Setup cancelled.')}\n")
+
+    # ── wizard steps ────────────────────────────────────────────────────────
+
+    @staticmethod
+    def _wizard() -> None:
+        SetupWizard._print_welcome()
+        print(f"  {bold('Lets get started.')}")
+        print()
+
+        providers = SetupWizard._select_providers()
+        env_vars = SetupWizard._collect_keys(providers)
+        port, master_key = SetupWizard._proxy_settings()
+
+        config_path = Path(os.getcwd()) / "litellm_config.yaml"
+        try:
+            config_path.write_text(
+                SetupWizard._build_config(providers, env_vars, master_key)
+            )
+        except OSError as exc:
+            print(f"\n  {bold(_CROSS + ' Could not write config:')} {exc}")
+            print("  Try running from a directory you have write access to.\n")
+            return
+
+        SetupWizard._print_success(config_path, port, master_key)
+        SetupWizard._offer_start(config_path, port, master_key)
+
+    # ── welcome ─────────────────────────────────────────────────────────────
+
+    @staticmethod
+    def _print_welcome() -> None:
+        try:
+            version = importlib.metadata.version("litellm")
+        except Exception:
+            version = "unknown"
+        print()
+        print(orange(LITELLM_ASCII.rstrip("\n")))
+        print(f"  {orange('Welcome')} to {bold('LiteLLM')} {grey('v' + version)}")
+        print()
+        print(_divider())
+        print()
+
+    # ── provider selector ───────────────────────────────────────────────────
+
+    @staticmethod
+    def _select_providers() -> List[Dict]:
+        """Arrow-key multi-select. Falls back to number input if /dev/tty unavailable."""
+        if not _HAS_RAW_TERMINAL:
+            return SetupWizard._select_fallback()
+        try:
+            return SetupWizard._select_interactive()
+        except OSError:
+            return SetupWizard._select_fallback()
+
+    @staticmethod
+    def _read_key() -> str:
+        """Read one keypress from /dev/tty in raw mode."""
+        assert (
+            termios is not None and tty is not None
+        )  # only called when _HAS_RAW_TERMINAL
+        with open("/dev/tty", "rb") as tty_fh:
+            fd = tty_fh.fileno()
+            old = termios.tcgetattr(fd)
+            try:
+                tty.setraw(fd)
+                ch = tty_fh.read(1)
+                if ch == b"\x1b":
+                    ch2 = tty_fh.read(1)
+                    if ch2 == b"[":
+                        ch3 = tty_fh.read(1)
+                        return "\x1b[" + ch3.decode("utf-8", errors="replace")
+                    return "\x1b" + ch2.decode("utf-8", errors="replace")
+                return ch.decode("utf-8", errors="replace")
+            finally:
+                termios.tcsetattr(fd, termios.TCSADRAIN, old)
+
+    @staticmethod
+    def _render_selector(cursor: int, selected: Set[int], first_render: bool) -> int:
+        """Draw or redraw the provider list. Returns the number of lines printed."""
+        lines = [
+            f"\n  {bold('Add your first model')}\n",
+            grey("  ↑↓ to navigate · Space to select · Enter to confirm") + "\n",
+            "\n",
+        ]
+        for i, p in enumerate(PROVIDERS):
+            arrow = blue("❯") if i == cursor else " "
+            bullet = green("◉") if i in selected else grey("○")
+            name_str = bold(p["name"]) if i == cursor else p["name"]
+            lines.append(f"  {arrow} {bullet} {name_str}  {grey(p['description'])}\n")
+        lines.append("\n")
+
+        content = "".join(lines)
+        if not first_render and _supports_color():
+            sys.stdout.write(_MOVE_UP.format(content.count("\n")))
+        sys.stdout.write(content)
+        sys.stdout.flush()
+        return content.count("\n")
+
+    @staticmethod
+    def _select_interactive() -> List[Dict]:
+        cursor = 0
+        selected: set[int] = set()
+
+        if _supports_color():
+            sys.stdout.write(_CURSOR_HIDE)
+            sys.stdout.flush()
+        try:
+            SetupWizard._render_selector(cursor, selected, first_render=True)
+            while True:
+                key = SetupWizard._read_key()
+                dirty = False
+                if key == "\x1b[A":
+                    cursor = (cursor - 1) % len(PROVIDERS)
+                    dirty = True
+                elif key == "\x1b[B":
+                    cursor = (cursor + 1) % len(PROVIDERS)
+                    dirty = True
+                elif key == " ":
+                    selected.symmetric_difference_update({cursor})
+                    dirty = True
+                elif key in ("\r", "\n"):
+                    if not selected:
+                        selected.add(cursor)
+                    break
+                elif key in ("\x03", "\x04"):
+                    raise KeyboardInterrupt
+                if dirty:
+                    SetupWizard._render_selector(cursor, selected, first_render=False)
+        finally:
+            if _supports_color():
+                sys.stdout.write(_CURSOR_SHOW)
+                sys.stdout.flush()
+
+        return [PROVIDERS[i] for i in sorted(selected)]
+
+    @staticmethod
+    def _select_fallback() -> List[Dict]:
+        """Number-based fallback when raw terminal input is unavailable."""
+        print()
+        print(f"  {bold('Add your first model')}")
+        print(
+            grey(
+                "  Enter numbers separated by commas (e.g. 1,2). Press Enter to confirm."
+            )
+        )
+        print()
+        for i, p in enumerate(PROVIDERS, 1):
+            print(f"  {grey(str(i) + '.')} {bold(p['name'])}  {grey(p['description'])}")
+        print()
+
+        while True:
+            raw = _styled_input(f"  {blue('❯')} Provider(s): ")
+            if not raw:
+                print(grey("  Please select at least one provider."))
+                continue
+            try:
+                nums = [
+                    int(x.strip())
+                    for x in raw.replace(" ", ",").split(",")
+                    if x.strip()
+                ]
+                valid = sorted({n for n in nums if 1 <= n <= len(PROVIDERS)})
+                if not valid:
+                    print(grey(f"  Enter numbers between 1 and {len(PROVIDERS)}."))
+                    continue
+                return [PROVIDERS[i - 1] for i in valid]
+            except ValueError:
+                print(grey("  Enter numbers separated by commas, e.g. 1,3"))
+
+    # ── key collection ───────────────────────────────────────────────────────
+
+    @staticmethod
+    def _collect_keys(providers: List[Dict]) -> Dict[str, str]:
+        env_vars: Dict[str, str] = {}
+        print()
+        print(_divider())
+        print()
+        print(f"  {bold('Enter your API keys')}")
+        print(grey("  Keys are stored only in the generated config file."))
+        print(
+            grey(
+                "  Tip: add litellm_config.yaml to .gitignore to avoid committing secrets."
+            )
+        )
+        print()
+
+        for p in providers:
+            if p["env_key"] is None:
+                print(
+                    f"  {green(p['name'])}: {grey('no key needed (uses local Ollama)')}"
+                )
+                continue
+
+            key = SetupWizard._prompt_key(p)
+            if not key:
+                continue
+
+            for extra_key, extra_hint in zip(
+                p.get("extra_keys", []), p.get("extra_hints", [])
+            ):
+                val = _styled_input(f"  {blue('❯')} {extra_key} {grey(extra_hint)}: ")
+                if val:
+                    env_vars[extra_key] = val
+
+            if p.get("needs_api_base"):
+                api_base = _styled_input(
+                    f"  {blue('❯')} Azure endpoint URL {grey(p.get('api_base_hint', ''))}: "
+                )
+                if api_base:
+                    env_vars[f"_LITELLM_AZURE_API_BASE_{p['id'].upper()}"] = api_base
+                deployment = _styled_input(
+                    f"  {blue('❯')} Azure deployment name {grey('(e.g. my-gpt4o)')}: "
+                )
+                if deployment:
+                    env_vars[
+                        f"_LITELLM_AZURE_DEPLOYMENT_{p['id'].upper()}"
+                    ] = deployment
+
+            # Store the key returned by validation — may be a re-entered replacement
+            env_vars[p["env_key"]] = SetupWizard._validate_and_report(p, key)
+
+        return env_vars
+
+    @staticmethod
+    def _prompt_key(provider: Dict) -> str:
+        """Prompt for a provider's API key, with skip option. Returns the key or ''."""
+        hint = grey(provider.get("key_hint", ""))
+        while True:
+            key = _styled_input(
+                f"  {blue('❯')} {bold(provider['name'])} API key {hint}: "
+            )
+            if key:
+                return key
+            print(grey("  Key is required. Leave blank to skip this provider."))
+            if _styled_input(grey("  Skip? (y/N): ")).lower() == "y":
+                return ""
+
+    @staticmethod
+    def _validate_and_report(provider: Dict, api_key: str) -> str:
+        """
+        Validate credentials using litellm.utils.check_valid_key and print result.
+        Offers a re-entry loop on failure. Returns the final (possibly re-entered) key.
+        """
+        test_model: Optional[str] = provider.get("test_model")
+        if not test_model:
+            return api_key  # Azure / Bedrock / Ollama — skip validation
+
+        while True:
+            print(
+                f"  {grey('Testing connection to ' + provider['name'] + '...')}",
+                flush=True,
+            )
+            valid = check_valid_key(model=test_model, api_key=api_key)
+            if valid:
+                print(
+                    f"  {green(_CHECK)} {bold(provider['name'])} connected successfully"
+                )
+                return api_key
+
+            print(f"  {_CROSS} {bold(provider['name'])} {grey('— invalid API key')}")
+            if (
+                _styled_input(f"  {blue('❯')} Re-enter key? {grey('(y/N)')}: ").lower()
+                != "y"
+            ):
+                return api_key
+
+            hint = grey(provider.get("key_hint", ""))
+            new_key = _styled_input(
+                f"  {blue('❯')} {bold(provider['name'])} API key {hint}: "
+            )
+            if not new_key:
+                return api_key
+            api_key = new_key
+
+    # ── proxy settings ───────────────────────────────────────────────────────
+
+    @staticmethod
+    def _proxy_settings() -> "tuple[int, str]":
+        print()
+        print(_divider())
+        print()
+        print(f"  {bold('Proxy settings')}")
+        print()
+        port = 4000
+        while True:
+            port_raw = _styled_input(f"  {blue('❯')} Port {grey('[4000]')}: ")
+            if not port_raw:
+                break
+            if port_raw.isdigit() and 1 <= int(port_raw) <= 65535:
+                port = int(port_raw)
+                break
+            print(grey("  Enter a valid port number (1–65535)."))
+        key_raw = _styled_input(f"  {blue('❯')} Master key {grey('[auto-generate]')}: ")
+        master_key = key_raw if key_raw else f"sk-{secrets.token_urlsafe(32)}"
+        return port, master_key
+
+    # ── config generation ────────────────────────────────────────────────────
+
+    @staticmethod
+    def _build_config(
+        providers: List[Dict],
+        env_vars: Dict[str, str],
+        master_key: str,
+    ) -> str:
+        env_copy = dict(env_vars)  # work on a copy — do not mutate caller's dict
+        lines = ["model_list:"]
+        for p in providers:
+            # Only emit models for providers that actually have credentials
+            has_creds = p["env_key"] is None or p["env_key"] in env_copy
+            if not has_creds:
+                continue
+
+            if p["id"] == "azure":
+                deployment = env_copy.pop(
+                    f"_LITELLM_AZURE_DEPLOYMENT_{p['id'].upper()}", ""
+                )
+                if not deployment:
+                    continue  # skip Azure entirely if no deployment name was provided
+                models = [f"azure/{deployment}"]
+            else:
+                models = p["models"]
+
+            for model in models:
+                raw_display = model.split("/")[-1] if "/" in model else model
+                # Qualify azure display names to avoid collision with OpenAI model names
+                display = f"azure-{raw_display}" if p["id"] == "azure" else raw_display
+                lines += [
+                    f"  - model_name: {display}",
+                    "    litellm_params:",
+                    f"      model: {model}",
+                ]
+                if p["env_key"] and p["env_key"] in env_copy:
+                    lines.append(f"      api_key: os.environ/{p['env_key']}")
+                if p.get("api_base"):
+                    lines.append(
+                        f'      api_base: "{_yaml_escape(str(p["api_base"]))}"'
+                    )
+                elif p.get("needs_api_base"):
+                    azure_base_key = f"_LITELLM_AZURE_API_BASE_{p['id'].upper()}"
+                    if azure_base_key in env_copy:
+                        lines.append(
+                            f'      api_base: "{_yaml_escape(env_copy.pop(azure_base_key))}"'
+                        )
+                if p.get("api_version"):
+                    lines.append(f"      api_version: {p['api_version']}")
+
+        lines += [
+            "",
+            "general_settings:",
+            f'  master_key: "{_yaml_escape(master_key)}"',
+            "",
+        ]
+
+        real_vars = {k: v for k, v in env_copy.items() if not k.startswith("_LITELLM_")}
+        if real_vars:
+            lines.append("environment_variables:")
+            for k, v in real_vars.items():
+                lines.append(f'  {k}: "{_yaml_escape(v)}"')
+            lines.append("")
+
+        return "\n".join(lines)
+
+    # ── success + launch ─────────────────────────────────────────────────────
+
+    @staticmethod
+    def _print_success(config_path: Path, port: int, master_key: str) -> None:
+        print()
+        print(_divider())
+        print()
+        print(f"  {green(_CHECK + ' Config saved')} → {bold(str(config_path))}")
+        print()
+        print(f"  {bold('To start your proxy:')}")
+        print()
+        print(f"    {grey('$')} litellm --config {config_path} --port {port}")
+        print()
+        print(f"  {bold('Then set your client:')}")
+        print()
+        print(f"    export OPENAI_BASE_URL=http://localhost:{port}")
+        print(f"    export OPENAI_API_KEY={master_key}")
+        print()
+        print(_divider())
+        print()
+
+    @staticmethod
+    def _offer_start(config_path: Path, port: int, master_key: str) -> None:
+        start = _styled_input(
+            f"  {blue('❯')} Start the proxy now? {grey('(Y/n)')}: "
+        ).lower()
+        if start not in ("", "y", "yes"):
+            print()
+            print(
+                f"  Run {bold(f'litellm --config {config_path}')} whenever you're ready."
+            )
+            print()
+            print(
+                grey(f"  Quick test once running:  curl http://localhost:{port}/health")
+            )
+            print()
+            return
+
+        print()
+        print(_divider())
+        print()
+        print(f"  {bold('Proxy is starting on')} http://localhost:{port}")
+        print()
+        print(grey("  Your proxy is OpenAI-compatible. Point any OpenAI SDK at it:"))
+        print()
+        print(f"    export OPENAI_BASE_URL=http://localhost:{port}")
+        print(f"    export OPENAI_API_KEY={master_key}")
+        print()
+        print(grey("  Quick test (in another terminal):"))
+        print()
+        print(f"    curl http://localhost:{port}/health")
+        print()
+        print(grey("  Dashboard:"))
+        print()
+        print(f"    http://localhost:{port}/ui  {grey('(login with your master key)')}")
+        print()
+        print(_divider())
+        print()
+        print(f"  {green(_CHECK)} Starting…  {grey('(Ctrl+C to stop)')}")
+        print()
+
+        scripts_dir = sysconfig.get_path("scripts")
+        litellm_bin = os.path.join(scripts_dir or "", "litellm")
+        try:
+            os.execlp(
+                litellm_bin,
+                litellm_bin,
+                "--config",
+                str(config_path),
+                "--port",
+                str(port),
+            )  # noqa: S606
+        except OSError as exc:
+            print(f"\n  {bold(_CROSS + ' Could not start proxy:')} {exc}")
+            print(f"  Run manually:  litellm --config {config_path} --port {port}\n")
+
+
+# ---------------------------------------------------------------------------
+# Public entrypoint
+# ---------------------------------------------------------------------------
+
+
+def run_setup_wizard() -> None:
+    """Run the interactive setup wizard. Called by `litellm --setup`."""
+    SetupWizard.run()

litellm/types/guardrails.py

@@ -79,6 +79,7 @@ class SupportedGuardrailIntegrations(Enum):
     SEMANTIC_GUARD = "semantic_guard"
     MCP_END_USER_PERMISSION = "mcp_end_user_permission"
     BLOCK_CODE_EXECUTION = "block_code_execution"
+    MCP_JWT_SIGNER = "mcp_jwt_signer"
 
 
 class Role(Enum):

litellm/types/proxy/management_endpoints/team_endpoints.py

@@ -43,10 +43,16 @@ class UpdateTeamMemberPermissionsRequest(BaseModel):
     team_member_permissions: List[str]
 
 
+class TeamListItem(LiteLLM_TeamTable):
+    """A team item in the paginated list response, enriched with computed fields."""
+
+    members_count: int = 0
+
+
 class TeamListResponse(BaseModel):
     """Response to get the list of teams"""
 
-    teams: List[Union[LiteLLM_TeamTable, LiteLLM_DeletedTeamTable]]
+    teams: List[Union[TeamListItem, LiteLLM_TeamTable, LiteLLM_DeletedTeamTable]]
     total: int
     page: int
     page_size: int

litellm/types/proxy/vantage_endpoints.py

@@ -39,7 +39,8 @@ class VantageExportRequest(BaseModel):
     """Request model for Vantage export operations (actual export, no default limit)"""
 
     limit: Optional[int] = Field(
-        None, description="Optional limit on number of records to export (default: no limit)"
+        None,
+        description="Optional limit on number of records to export (default: no limit)",
     )
     start_time_utc: Optional[datetime] = Field(
         None, description="Start time for data export in UTC"

litellm/types/videos/utils.py

@@ -195,7 +195,9 @@ def decode_character_id_with_provider(encoded_character_id: str) -> DecodedChara
             character_id=decoded_character_id,
         )
     except Exception as e:
-        verbose_logger.debug(f"Error decoding character_id '{encoded_character_id}': {e}")
+        verbose_logger.debug(
+            f"Error decoding character_id '{encoded_character_id}': {e}"
+        )
         return DecodedCharacterId(
             custom_llm_provider=None,
             model_id=None,

litellm/videos/main.py

@@ -1186,13 +1186,17 @@ def video_create_character(
 
         litellm_params = GenericLiteLLMParams(**kwargs)
 
-        provider_config: Optional[BaseVideoConfig] = ProviderConfigManager.get_provider_video_config(
+        provider_config: Optional[
+            BaseVideoConfig
+        ] = ProviderConfigManager.get_provider_video_config(
             model=None,
             provider=litellm.LlmProviders(custom_llm_provider),
         )
 
         if provider_config is None:
-            raise ValueError(f"video create character is not supported for {custom_llm_provider}")
+            raise ValueError(
+                f"video create character is not supported for {custom_llm_provider}"
+            )
 
         local_vars.update(kwargs)
         request_params: Dict = {"name": name}
@@ -1311,13 +1315,17 @@ def video_get_character(
 
         litellm_params = GenericLiteLLMParams(**kwargs)
 
-        provider_config: Optional[BaseVideoConfig] = ProviderConfigManager.get_provider_video_config(
+        provider_config: Optional[
+            BaseVideoConfig
+        ] = ProviderConfigManager.get_provider_video_config(
             model=None,
             provider=litellm.LlmProviders(custom_llm_provider),
         )
 
         if provider_config is None:
-            raise ValueError(f"video get character is not supported for {custom_llm_provider}")
+            raise ValueError(
+                f"video get character is not supported for {custom_llm_provider}"
+            )
 
         local_vars.update(kwargs)
         request_params: Dict = {"character_id": character_id}
@@ -1439,7 +1447,9 @@ def video_edit(
 
         litellm_params = GenericLiteLLMParams(**kwargs)
 
-        provider_config: Optional[BaseVideoConfig] = ProviderConfigManager.get_provider_video_config(
+        provider_config: Optional[
+            BaseVideoConfig
+        ] = ProviderConfigManager.get_provider_video_config(
             model=None,
             provider=litellm.LlmProviders(custom_llm_provider),
         )
@@ -1572,16 +1582,24 @@ def video_extension(
 
         litellm_params = GenericLiteLLMParams(**kwargs)
 
-        provider_config: Optional[BaseVideoConfig] = ProviderConfigManager.get_provider_video_config(
+        provider_config: Optional[
+            BaseVideoConfig
+        ] = ProviderConfigManager.get_provider_video_config(
             model=None,
             provider=litellm.LlmProviders(custom_llm_provider),
         )
 
         if provider_config is None:
-            raise ValueError(f"video extension is not supported for {custom_llm_provider}")
+            raise ValueError(
+                f"video extension is not supported for {custom_llm_provider}"
+            )
 
         local_vars.update(kwargs)
-        request_params: Dict = {"video_id": video_id, "prompt": prompt, "seconds": seconds}
+        request_params: Dict = {
+            "video_id": video_id,
+            "prompt": prompt,
+            "seconds": seconds,
+        }
 
         litellm_logging_obj.update_environment_variables(
             model="",

model_prices_and_context_window.json

@@ -32354,6 +32354,53 @@
         "supports_vision": true,
         "supports_web_search": true
     },
+    "xai/grok-4.20-multi-agent-beta-0309": {
+        "cache_read_input_token_cost": 2e-07,
+        "input_cost_per_token": 2e-06,
+        "litellm_provider": "xai",
+        "max_input_tokens": 2000000,
+        "max_output_tokens": 2000000,
+        "max_tokens": 2000000,
+        "mode": "chat",
+        "output_cost_per_token": 6e-06,
+        "source": "https://docs.x.ai/docs/models",
+        "supports_function_calling": true,
+        "supports_reasoning": true,
+        "supports_tool_choice": true,
+        "supports_vision": true,
+        "supports_web_search": true
+    },
+    "xai/grok-4.20-beta-0309-reasoning": {
+        "cache_read_input_token_cost": 2e-07,
+        "input_cost_per_token": 2e-06,
+        "litellm_provider": "xai",
+        "max_input_tokens": 2000000,
+        "max_output_tokens": 2000000,
+        "max_tokens": 2000000,
+        "mode": "chat",
+        "output_cost_per_token": 6e-06,
+        "source": "https://docs.x.ai/docs/models",
+        "supports_function_calling": true,
+        "supports_reasoning": true,
+        "supports_tool_choice": true,
+        "supports_vision": true,
+        "supports_web_search": true
+    },
+    "xai/grok-4.20-beta-0309-non-reasoning": {
+        "cache_read_input_token_cost": 2e-07,
+        "input_cost_per_token": 2e-06,
+        "litellm_provider": "xai",
+        "max_input_tokens": 2000000,
+        "max_output_tokens": 2000000,
+        "max_tokens": 2000000,
+        "mode": "chat",
+        "output_cost_per_token": 6e-06,
+        "source": "https://docs.x.ai/docs/models",
+        "supports_function_calling": true,
+        "supports_tool_choice": true,
+        "supports_vision": true,
+        "supports_web_search": true
+    },
     "xai/grok-beta": {
         "input_cost_per_token": 5e-06,
         "litellm_provider": "xai",

poetry.lock

@@ -1,4 +1,4 @@
-# This file is automatically @generated by Poetry 2.2.1 and should not be changed by hand.
+# This file is automatically @generated by Poetry 2.3.2 and should not be changed by hand.
 
 [[package]]
 name = "a2a-sdk"
@@ -7,11 +7,11 @@ description = "A2A Python SDK"
 optional = false
 python-versions = ">=3.10"
 groups = ["main", "proxy-dev"]
-markers = "python_version >= \"3.10\""
 files = [
     {file = "a2a_sdk-0.3.22-py3-none-any.whl", hash = "sha256:b98701135bb90b0ff85d35f31533b6b7a299bf810658c1c65f3814a6c15ea385"},
     {file = "a2a_sdk-0.3.22.tar.gz", hash = "sha256:77a5694bfc4f26679c11b70c7f1062522206d430b34bc1215cfbb1eba67b7e7d"},
 ]
+markers = {main = "python_version >= \"3.10\" and extra == \"extra-proxy\"", proxy-dev = "python_version >= \"3.10\""}
 
 [package.dependencies]
 google-api-core = ">=1.26.0"
@@ -385,6 +385,7 @@ files = [
     {file = "azure_core-1.36.0-py3-none-any.whl", hash = "sha256:fee9923a3a753e94a259563429f3644aaf05c486d45b1215d098115102d91d3b"},
     {file = "azure_core-1.36.0.tar.gz", hash = "sha256:22e5605e6d0bf1d229726af56d9e92bc37b6e726b141a18be0b4d424131741b7"},
 ]
+markers = {main = "extra == \"proxy\" or extra == \"extra-proxy\""}
 
 [package.dependencies]
 requests = ">=2.21.0"
@@ -405,6 +406,7 @@ files = [
     {file = "azure_identity-1.25.1-py3-none-any.whl", hash = "sha256:e9edd720af03dff020223cd269fa3a61e8f345ea75443858273bcb44844ab651"},
     {file = "azure_identity-1.25.1.tar.gz", hash = "sha256:87ca8328883de6036443e1c37b40e8dc8fb74898240f61071e09d2e369361456"},
 ]
+markers = {main = "extra == \"proxy\" or extra == \"extra-proxy\""}
 
 [package.dependencies]
 azure-core = ">=1.31.0"
@@ -598,7 +600,7 @@ files = [
     {file = "cachetools-6.2.2-py3-none-any.whl", hash = "sha256:6c09c98183bf58560c97b2abfcedcbaf6a896a490f534b031b661d3723b45ace"},
     {file = "cachetools-6.2.2.tar.gz", hash = "sha256:8e6d266b25e539df852251cfd6f990b4bc3a141db73b939058d809ebd2590fc6"},
 ]
-markers = {main = "extra == \"google\" or extra == \"extra-proxy\" or python_version >= \"3.10\"", proxy-dev = "python_version >= \"3.10\""}
+markers = {main = "python_version >= \"3.10\" and (extra == \"extra-proxy\" or extra == \"google\" or extra == \"mlflow\") or extra == \"google\" or extra == \"extra-proxy\"", proxy-dev = "python_version >= \"3.10\""}
 
 [[package]]
 name = "certifi"
@@ -705,7 +707,7 @@ files = [
     {file = "cffi-2.0.0-cp39-cp39-win_amd64.whl", hash = "sha256:b882b3df248017dba09d6b16defe9b5c407fe32fc7c65a9c69798e6175601be9"},
     {file = "cffi-2.0.0.tar.gz", hash = "sha256:44d1b5909021139fe36001ae048dbdde8214afa20200eda0f64c068cac5d5529"},
 ]
-markers = {main = "platform_python_implementation != \"PyPy\" or extra == \"proxy\"", dev = "platform_python_implementation != \"PyPy\"", proxy-dev = "platform_python_implementation != \"PyPy\""}
+markers = {main = "(platform_python_implementation != \"PyPy\" or extra == \"proxy\") and (python_version >= \"3.10\" or extra == \"proxy\" or extra == \"extra-proxy\") and (extra == \"proxy\" or extra == \"extra-proxy\" or extra == \"mlflow\")", dev = "platform_python_implementation != \"PyPy\"", proxy-dev = "platform_python_implementation != \"PyPy\""}
 
 [package.dependencies]
 pycparser = {version = "*", markers = "implementation_name != \"PyPy\""}
@@ -1055,6 +1057,7 @@ files = [
     {file = "cryptography-43.0.3-pp39-pypy39_pp73-win_amd64.whl", hash = "sha256:2ce6fae5bdad59577b44e4dfed356944fbf1d925269114c28be377692643b4ff"},
     {file = "cryptography-43.0.3.tar.gz", hash = "sha256:315b9001266a492a6ff443b61238f956b214dbec9910a081ba5b6646a055a805"},
 ]
+markers = {main = "python_version >= \"3.10\" and (extra == \"proxy\" or extra == \"extra-proxy\" or extra == \"mlflow\") or extra == \"proxy\" or extra == \"extra-proxy\""}
 
 [package.dependencies]
 cffi = {version = ">=1.12", markers = "platform_python_implementation != \"PyPy\""}
@@ -1837,11 +1840,11 @@ description = "Google API client core library"
 optional = false
 python-versions = ">=3.7"
 groups = ["main", "proxy-dev"]
-markers = "python_version >= \"3.14\""
 files = [
     {file = "google_api_core-2.25.2-py3-none-any.whl", hash = "sha256:e9a8f62d363dc8424a8497f4c2a47d6bcda6c16514c935629c257ab5d10210e7"},
     {file = "google_api_core-2.25.2.tar.gz", hash = "sha256:1c63aa6af0d0d5e37966f157a77f9396d820fba59f9e43e9415bc3dc5baff300"},
 ]
+markers = {main = "python_version >= \"3.14\" and (extra == \"extra-proxy\" or extra == \"google\")", proxy-dev = "python_version >= \"3.14\""}
 
 [package.dependencies]
 google-auth = ">=2.14.1,<3.0.0"
@@ -1869,7 +1872,7 @@ files = [
     {file = "google_api_core-2.28.1-py3-none-any.whl", hash = "sha256:4021b0f8ceb77a6fb4de6fde4502cecab45062e66ff4f2895169e0b35bc9466c"},
     {file = "google_api_core-2.28.1.tar.gz", hash = "sha256:2b405df02d68e68ce0fbc138559e6036559e685159d148ae5861013dc201baf8"},
 ]
-markers = {main = "(python_version >= \"3.10\" or extra == \"google\" or extra == \"extra-proxy\") and python_version < \"3.14\"", proxy-dev = "python_version >= \"3.10\" and python_version < \"3.14\""}
+markers = {main = "python_version < \"3.14\" and (extra == \"extra-proxy\" or extra == \"google\")", proxy-dev = "python_version >= \"3.10\" and python_version < \"3.14\""}
 
 [package.dependencies]
 google-auth = ">=2.14.1,<3.0.0"
@@ -1906,7 +1909,7 @@ files = [
     {file = "google_auth-2.43.0-py2.py3-none-any.whl", hash = "sha256:af628ba6fa493f75c7e9dbe9373d148ca9f4399b5ea29976519e0a3848eddd16"},
     {file = "google_auth-2.43.0.tar.gz", hash = "sha256:88228eee5fc21b62a1b5fe773ca15e67778cb07dc8363adcb4a8827b52d81483"},
 ]
-markers = {main = "extra == \"google\" or extra == \"extra-proxy\" or python_version >= \"3.10\"", proxy-dev = "python_version >= \"3.10\""}
+markers = {main = "python_version >= \"3.10\" and (extra == \"extra-proxy\" or extra == \"google\" or extra == \"mlflow\") or extra == \"google\" or extra == \"extra-proxy\"", proxy-dev = "python_version >= \"3.10\""}
 
 [package.dependencies]
 cachetools = ">=2.0.0,<7.0"
@@ -2078,11 +2081,11 @@ files = [
 ]
 
 [package.dependencies]
-google-api-core = {version = ">=1.34.1,<2.0.dev0 || >=2.11.dev0,<3.0.0dev", extras = ["grpc"]}
-google-auth = ">=2.14.1,<2.24.0 || >2.24.0,<2.25.0 || >2.25.0,<3.0.0dev"
-grpc-google-iam-v1 = ">=0.12.4,<1.0.0dev"
-proto-plus = ">=1.22.3,<2.0.0dev"
-protobuf = ">=3.20.2,<4.21.0 || >4.21.0,<4.21.1 || >4.21.1,<4.21.2 || >4.21.2,<4.21.3 || >4.21.3,<4.21.4 || >4.21.4,<4.21.5 || >4.21.5,<6.0.0dev"
+google-api-core = {version = ">=1.34.1,<2.0.dev0 || >=2.11.dev0,<3.0.0.dev0", extras = ["grpc"]}
+google-auth = ">=2.14.1,<2.24.0 || >2.24.0,<2.25.0 || >2.25.0,<3.0.0.dev0"
+grpc-google-iam-v1 = ">=0.12.4,<1.0.0.dev0"
+proto-plus = ">=1.22.3,<2.0.0.dev0"
+protobuf = ">=3.20.2,<4.21.0 || >4.21.0,<4.21.1 || >4.21.1,<4.21.2 || >4.21.2,<4.21.3 || >4.21.3,<4.21.4 || >4.21.4,<4.21.5 || >4.21.5,<6.0.0.dev0"
 
 [[package]]
 name = "google-cloud-resource-manager"
@@ -2264,7 +2267,7 @@ files = [
     {file = "googleapis_common_protos-1.72.0-py3-none-any.whl", hash = "sha256:4299c5a82d5ae1a9702ada957347726b167f9f8d1fc352477702a1e851ff4038"},
     {file = "googleapis_common_protos-1.72.0.tar.gz", hash = "sha256:e55a601c1b32b52d7a3e65f43563e2aa61bcd737998ee672ac9b951cd49319f5"},
 ]
-markers = {main = "extra == \"google\" or extra == \"extra-proxy\" or python_version >= \"3.10\""}
+markers = {main = "python_version >= \"3.10\" and (extra == \"extra-proxy\" or extra == \"google\") or extra == \"google\" or extra == \"extra-proxy\""}
 
 [package.dependencies]
 grpcio = {version = ">=1.44.0,<2.0.0", optional = true, markers = "extra == \"grpc\""}
@@ -2673,11 +2676,11 @@ description = "Consume Server-Sent Event (SSE) messages with HTTPX."
 optional = false
 python-versions = ">=3.9"
 groups = ["main", "proxy-dev"]
-markers = "python_version >= \"3.10\""
 files = [
     {file = "httpx_sse-0.4.3-py3-none-any.whl", hash = "sha256:0ac1c9fe3c0afad2e0ebb25a934a59f4c7823b60792691f779fad2c5568830fc"},
     {file = "httpx_sse-0.4.3.tar.gz", hash = "sha256:9b1ed0127459a66014aec3c56bebd93da3c1bc8bb6618c8082039a44889a755d"},
 ]
+markers = {main = "python_version >= \"3.10\" and (extra == \"proxy\" or extra == \"extra-proxy\")", proxy-dev = "python_version >= \"3.10\""}
 
 [[package]]
 name = "huey"
@@ -3042,7 +3045,7 @@ files = [
 
 [package.dependencies]
 attrs = ">=22.2.0"
-jsonschema-specifications = ">=2023.03.6"
+jsonschema-specifications = ">=2023.3.6"
 referencing = ">=0.28.4"
 rpds-py = ">=0.7.1"
 
@@ -3219,15 +3222,15 @@ files = [
 
 [[package]]
 name = "litellm-proxy-extras"
-version = "0.4.56"
+version = "0.4.58"
 description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
 optional = true
 python-versions = "!=2.7.*,!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,!=3.6.*,!=3.7.*,>=3.8"
 groups = ["main"]
 markers = "extra == \"proxy\""
 files = [
-    {file = "litellm_proxy_extras-0.4.56-py3-none-any.whl", hash = "sha256:52dbe3b5358c790e77e12f1ec5ef8e7508b383c2aaf41299750b6fb400908ee7"},
-    {file = "litellm_proxy_extras-0.4.56.tar.gz", hash = "sha256:63ad59baa0defccc5c929cfd933ee7e32a6614b0fc5fa0fc45a12d7608e33f08"},
+    {file = "litellm_proxy_extras-0.4.58-py3-none-any.whl", hash = "sha256:8863e70de833c0e35119a1cbbf583619bdebe52222efd5654586519175ba403b"},
+    {file = "litellm_proxy_extras-0.4.58.tar.gz", hash = "sha256:84a67483329eced8be4fc61c4e43f117287aa4e3deeb8ddf8fe8cdc9a8508836"},
 ]
 
 [[package]]
@@ -3713,6 +3716,7 @@ files = [
     {file = "msal-1.34.0-py3-none-any.whl", hash = "sha256:f669b1644e4950115da7a176441b0e13ec2975c29528d8b9e81316023676d6e1"},
     {file = "msal-1.34.0.tar.gz", hash = "sha256:76ba83b716ea5a6d75b0279c0ac353a0e05b820ca1f6682c0eb7f45190c43c2f"},
 ]
+markers = {main = "extra == \"proxy\" or extra == \"extra-proxy\""}
 
 [package.dependencies]
 cryptography = ">=2.5,<49"
@@ -3733,6 +3737,7 @@ files = [
     {file = "msal_extensions-1.3.1-py3-none-any.whl", hash = "sha256:96d3de4d034504e969ac5e85bae8106c8373b5c6568e4c8fa7af2eca9dbe6bca"},
     {file = "msal_extensions-1.3.1.tar.gz", hash = "sha256:c5b0fd10f65ef62b5f1d62f4251d51cbcaf003fcedae8c91b040a488614be1a4"},
 ]
+markers = {main = "extra == \"proxy\" or extra == \"extra-proxy\""}
 
 [package.dependencies]
 msal = ">=1.29,<2"
@@ -3983,6 +3988,7 @@ files = [
     {file = "nodeenv-1.9.1-py2.py3-none-any.whl", hash = "sha256:ba11c9782d29c27c70ffbdda2d7415098754709be8a7056d79a737cd901155c9"},
     {file = "nodeenv-1.9.1.tar.gz", hash = "sha256:6ec12890a2dab7946721edbfbcd91f3319c6ccc9aec47be7c7e6b7011ee6645f"},
 ]
+markers = {main = "extra == \"extra-proxy\""}
 
 [[package]]
 name = "numpy"
@@ -4105,7 +4111,7 @@ files = [
     {file = "opentelemetry_api-1.39.1-py3-none-any.whl", hash = "sha256:2edd8463432a7f8443edce90972169b195e7d6a05500cd29e6d13898187c9950"},
     {file = "opentelemetry_api-1.39.1.tar.gz", hash = "sha256:fbde8c80e1b937a2c61f20347e91c0c18a1940cecf012d62e65a7caf08967c9c"},
 ]
-markers = {main = "python_version >= \"3.10\""}
+markers = {main = "python_version >= \"3.10\" and extra == \"mlflow\""}
 
 [package.dependencies]
 importlib-metadata = ">=6.0,<8.8.0"
@@ -4220,7 +4226,7 @@ files = [
     {file = "opentelemetry_sdk-1.39.1-py3-none-any.whl", hash = "sha256:4d5482c478513ecb0a5d938dcc61394e647066e0cc2676bee9f3af3f3f45f01c"},
     {file = "opentelemetry_sdk-1.39.1.tar.gz", hash = "sha256:cf4d4563caf7bff906c9f7967e2be22d0d6b349b908be0d90fb21c8e9c995cc6"},
 ]
-markers = {main = "python_version >= \"3.10\""}
+markers = {main = "python_version >= \"3.10\" and extra == \"mlflow\""}
 
 [package.dependencies]
 opentelemetry-api = "1.39.1"
@@ -4238,7 +4244,7 @@ files = [
     {file = "opentelemetry_semantic_conventions-0.60b1-py3-none-any.whl", hash = "sha256:9fa8c8b0c110da289809292b0591220d3a7b53c1526a23021e977d68597893fb"},
     {file = "opentelemetry_semantic_conventions-0.60b1.tar.gz", hash = "sha256:87c228b5a0669b748c76d76df6c364c369c28f1c465e50f661e39737e84bc953"},
 ]
-markers = {main = "python_version >= \"3.10\""}
+markers = {main = "python_version >= \"3.10\" and extra == \"mlflow\""}
 
 [package.dependencies]
 opentelemetry-api = "1.39.1"
@@ -4455,6 +4461,21 @@ sql-other = ["SQLAlchemy (>=2.0.0)", "adbc-driver-postgresql (>=0.8.0)", "adbc-d
 test = ["hypothesis (>=6.46.1)", "pytest (>=7.3.2)", "pytest-xdist (>=2.2.0)"]
 xml = ["lxml (>=4.9.2)"]
 
+[[package]]
+name = "parameterized"
+version = "0.9.0"
+description = "Parameterized testing with any Python test framework"
+optional = false
+python-versions = ">=3.7"
+groups = ["dev"]
+files = [
+    {file = "parameterized-0.9.0-py2.py3-none-any.whl", hash = "sha256:4e0758e3d41bea3bbd05ec14fc2c24736723f243b28d702081aef438c9372b1b"},
+    {file = "parameterized-0.9.0.tar.gz", hash = "sha256:7fc905272cefa4f364c1a3429cbbe9c0f98b793988efb5bf90aac80f08db09b1"},
+]
+
+[package.extras]
+dev = ["jinja2"]
+
 [[package]]
 name = "pathspec"
 version = "0.12.1"
@@ -4722,6 +4743,7 @@ files = [
     {file = "prisma-0.11.0-py3-none-any.whl", hash = "sha256:22bb869e59a2968b99f3483bb417717273ffbc569fd1e9ceed95e5614cbaf53a"},
     {file = "prisma-0.11.0.tar.gz", hash = "sha256:3f2f2fd2361e1ec5ff655f2a04c7860c2f2a5bc4c91f78ca9c5c6349735bf693"},
 ]
+markers = {main = "extra == \"extra-proxy\""}
 
 [package.dependencies]
 click = ">=7.1.2"
@@ -4895,7 +4917,7 @@ files = [
     {file = "proto_plus-1.26.1-py3-none-any.whl", hash = "sha256:13285478c2dcf2abb829db158e1047e2f1e8d63a077d94263c2b88b043c75a66"},
     {file = "proto_plus-1.26.1.tar.gz", hash = "sha256:21a515a4c4c0088a773899e23c7bbade3d18f9c66c73edd4c7ee3816bc96a012"},
 ]
-markers = {main = "extra == \"google\" or extra == \"extra-proxy\" or python_version >= \"3.10\"", proxy-dev = "python_version >= \"3.10\""}
+markers = {main = "extra == \"google\" or extra == \"extra-proxy\"", proxy-dev = "python_version >= \"3.10\""}
 
 [package.dependencies]
 protobuf = ">=3.19.0,<7.0.0"
@@ -4923,7 +4945,7 @@ files = [
     {file = "protobuf-5.29.5-py3-none-any.whl", hash = "sha256:6cf42630262c59b2d8de33954443d94b746c952b01434fc58a417fdbd2e84bd5"},
     {file = "protobuf-5.29.5.tar.gz", hash = "sha256:bc1463bafd4b0929216c35f437a8e28731a2b7fe3d98bb77a600efced5a15c84"},
 ]
-markers = {main = "extra == \"google\" or extra == \"extra-proxy\" or python_version >= \"3.10\""}
+markers = {main = "python_version >= \"3.10\" and (extra == \"extra-proxy\" or extra == \"google\" or extra == \"mlflow\") or extra == \"google\" or extra == \"extra-proxy\""}
 
 [[package]]
 name = "psutil"
@@ -5083,7 +5105,7 @@ files = [
     {file = "pyasn1-0.6.1-py3-none-any.whl", hash = "sha256:0d632f46f2ba09143da3a8afe9e33fb6f92fa2320ab7e886e2d0f7672af84629"},
     {file = "pyasn1-0.6.1.tar.gz", hash = "sha256:6f580d2bdd84365380830acf45550f2511469f673cb4a5ae3857a3170128b034"},
 ]
-markers = {main = "extra == \"google\" or extra == \"extra-proxy\" or python_version >= \"3.10\"", proxy-dev = "python_version >= \"3.10\""}
+markers = {main = "python_version >= \"3.10\" and (extra == \"extra-proxy\" or extra == \"google\" or extra == \"mlflow\") or extra == \"google\" or extra == \"extra-proxy\"", proxy-dev = "python_version >= \"3.10\""}
 
 [[package]]
 name = "pyasn1-modules"
@@ -5096,7 +5118,7 @@ files = [
     {file = "pyasn1_modules-0.4.2-py3-none-any.whl", hash = "sha256:29253a9207ce32b64c3ac6600edc75368f98473906e8fd1043bd6b5b1de2c14a"},
     {file = "pyasn1_modules-0.4.2.tar.gz", hash = "sha256:677091de870a80aae844b1ca6134f54652fa2c8c5a52aa396440ac3106e941e6"},
 ]
-markers = {main = "extra == \"google\" or extra == \"extra-proxy\" or python_version >= \"3.10\"", proxy-dev = "python_version >= \"3.10\""}
+markers = {main = "python_version >= \"3.10\" and (extra == \"extra-proxy\" or extra == \"google\" or extra == \"mlflow\") or extra == \"google\" or extra == \"extra-proxy\"", proxy-dev = "python_version >= \"3.10\""}
 
 [package.dependencies]
 pyasn1 = ">=0.6.1,<0.7.0"
@@ -5124,7 +5146,7 @@ files = [
     {file = "pycparser-2.23-py3-none-any.whl", hash = "sha256:e5c6e8d3fbad53479cab09ac03729e0a9faf2bee3db8208a550daf5af81a5934"},
     {file = "pycparser-2.23.tar.gz", hash = "sha256:78816d4f24add8f10a06d6f05b4d424ad9e96cfebf68a4ddc99c65c0720d00c2"},
 ]
-markers = {main = "implementation_name != \"PyPy\" and (platform_python_implementation != \"PyPy\" or extra == \"proxy\")", dev = "platform_python_implementation != \"PyPy\" and implementation_name != \"PyPy\"", proxy-dev = "platform_python_implementation != \"PyPy\" and implementation_name != \"PyPy\""}
+markers = {main = "implementation_name != \"PyPy\" and (platform_python_implementation != \"PyPy\" or extra == \"proxy\") and (python_version >= \"3.10\" or extra == \"proxy\" or extra == \"extra-proxy\") and (extra == \"proxy\" or extra == \"extra-proxy\" or extra == \"mlflow\")", dev = "platform_python_implementation != \"PyPy\" and implementation_name != \"PyPy\"", proxy-dev = "platform_python_implementation != \"PyPy\" and implementation_name != \"PyPy\""}
 
 [[package]]
 name = "pydantic"
@@ -5347,6 +5369,7 @@ files = [
     {file = "pyjwt-2.12.1-py3-none-any.whl", hash = "sha256:28ca37c070cad8ba8cd9790cd940535d40274d22f80ab87f3ac6a713e6e8454c"},
     {file = "pyjwt-2.12.1.tar.gz", hash = "sha256:c74a7a2adf861c04d002db713dd85f84beb242228e671280bf709d765b03672b"},
 ]
+markers = {main = "(python_version <= \"3.13\" or extra == \"proxy\" or extra == \"extra-proxy\") and (extra == \"extra-proxy\" or extra == \"proxy\")"}
 
 [package.dependencies]
 cryptography = {version = ">=3.4.0", optional = true, markers = "extra == \"crypto\""}
@@ -6290,7 +6313,7 @@ files = [
     {file = "rsa-4.9.1-py3-none-any.whl", hash = "sha256:68635866661c6836b8d39430f97a996acbd61bfa49406748ea243539fe239762"},
     {file = "rsa-4.9.1.tar.gz", hash = "sha256:e7bdbfdb5497da4c07dfd35530e1a902659db6ff241e39d9953cad06ebd0ae75"},
 ]
-markers = {main = "extra == \"google\" or extra == \"extra-proxy\" or python_version >= \"3.10\"", proxy-dev = "python_version >= \"3.10\""}
+markers = {main = "python_version >= \"3.10\" and (extra == \"extra-proxy\" or extra == \"google\" or extra == \"mlflow\") or extra == \"google\" or extra == \"extra-proxy\"", proxy-dev = "python_version >= \"3.10\""}
 
 [package.dependencies]
 pyasn1 = ">=0.1.3"
@@ -6336,10 +6359,10 @@ files = [
 ]
 
 [package.dependencies]
-botocore = ">=1.37.4,<2.0a.0"
+botocore = ">=1.37.4,<2.0a0"
 
 [package.extras]
-crt = ["botocore[crt] (>=1.37.4,<2.0a.0)"]
+crt = ["botocore[crt] (>=1.37.4,<2.0a0)"]
 
 [[package]]
 name = "scikit-learn"
@@ -6492,9 +6515,9 @@ tornado = ">=6.4.2,<7"
 urllib3 = ">=1.26,<3"
 
 [package.extras]
-all = ["boto3 (>=1.34.98,<2)", "botocore (>=1.34.110,<2)", "cohere (>=5.9.4,<6.00)", "dagger-io (>=0.1.1) ; python_version >= \"3.11\"", "fastembed (>=0.3.0,<0.4) ; python_version < \"3.13\"", "google-cloud-aiplatform (>=1.45.0,<2)", "ipykernel (>=6.25.0,<7)", "llama-cpp-python (>=0.2.28,<0.2.86) ; python_version < \"3.13\"", "mistralai (>=0.0.12,<0.1.0)", "mypy (>=1.7.1,<2)", "ollama (>=0.1.7)", "pillow (>=10.2.0,<11.0.0) ; python_version < \"3.13\"", "pinecone[asyncio] (>=7.0.0,<8.0.0)", "psycopg[binary] (>=3.1.0,<4)", "pytest (>=8.2,<9.0)", "pytest-asyncio (>=0.24.0,<0.25)", "pytest-cov (>=4.1.0,<5)", "pytest-mock (>=3.12.0,<4)", "pytest-timeout", "pytest-xdist (>=3.5.0,<4)", "python-dotenv (>=1.0.0,<2)", "qdrant-client (>=1.11.1,<2)", "requests-mock (>=1.12.1,<2)", "ruff (>=0.11.2,<0.12)", "sentence-transformers (>=5.0.0) ; python_version < \"3.13\"", "tokenizers (>=0.19) ; python_version < \"3.13\"", "torch (>=2.6.0) ; python_version < \"3.13\"", "torchvision (>=0.17.0) ; python_version < \"3.13\"", "transformers (>=4.36.2) ; python_version < \"3.13\"", "types-pyyaml (>=6.0.12.12,<7)", "types-requests (>=2.31.0,<3)"]
+all = ["boto3 (>=1.34.98,<2)", "botocore (>=1.34.110,<2)", "cohere (>=5.9.4,<6.0)", "dagger-io (>=0.1.1) ; python_version >= \"3.11\"", "fastembed (>=0.3.0,<0.4) ; python_version < \"3.13\"", "google-cloud-aiplatform (>=1.45.0,<2)", "ipykernel (>=6.25.0,<7)", "llama-cpp-python (>=0.2.28,<0.2.86) ; python_version < \"3.13\"", "mistralai (>=0.0.12,<0.1.0)", "mypy (>=1.7.1,<2)", "ollama (>=0.1.7)", "pillow (>=10.2.0,<11.0.0) ; python_version < \"3.13\"", "pinecone[asyncio] (>=7.0.0,<8.0.0)", "psycopg[binary] (>=3.1.0,<4)", "pytest (>=8.2,<9.0)", "pytest-asyncio (>=0.24.0,<0.25)", "pytest-cov (>=4.1.0,<5)", "pytest-mock (>=3.12.0,<4)", "pytest-timeout", "pytest-xdist (>=3.5.0,<4)", "python-dotenv (>=1.0.0,<2)", "qdrant-client (>=1.11.1,<2)", "requests-mock (>=1.12.1,<2)", "ruff (>=0.11.2,<0.12)", "sentence-transformers (>=5.0.0) ; python_version < \"3.13\"", "tokenizers (>=0.19) ; python_version < \"3.13\"", "torch (>=2.6.0) ; python_version < \"3.13\"", "torchvision (>=0.17.0) ; python_version < \"3.13\"", "transformers (>=4.36.2) ; python_version < \"3.13\"", "types-pyyaml (>=6.0.12.12,<7)", "types-requests (>=2.31.0,<3)"]
 bedrock = ["boto3 (>=1.34.98,<2)", "botocore (>=1.34.110,<2)"]
-cohere = ["cohere (>=5.9.4,<6.00)"]
+cohere = ["cohere (>=5.9.4,<6.0)"]
 dev = ["dagger-io (>=0.1.1) ; python_version >= \"3.11\"", "ipykernel (>=6.25.0,<7)", "mypy (>=1.7.1,<2)", "pytest (>=8.2,<9.0)", "pytest-asyncio (>=0.24.0,<0.25)", "pytest-cov (>=4.1.0,<5)", "pytest-mock (>=3.12.0,<4)", "pytest-timeout", "pytest-xdist (>=3.5.0,<4)", "python-dotenv (>=1.0.0,<2)", "requests-mock (>=1.12.1,<2)", "ruff (>=0.11.2,<0.12)", "types-pyyaml (>=6.0.12.12,<7)", "types-requests (>=2.31.0,<3)"]
 docs = ["pydoc-markdown (>=4.8.2) ; python_version < \"3.12\""]
 fastembed = ["fastembed (>=0.3.0,<0.4) ; python_version < \"3.13\""]
@@ -7222,6 +7245,7 @@ files = [
     {file = "tomlkit-0.13.3-py3-none-any.whl", hash = "sha256:c89c649d79ee40629a9fda55f8ace8c6a1b42deb912b2a8fd8d942ddadb606b0"},
     {file = "tomlkit-0.13.3.tar.gz", hash = "sha256:430cf247ee57df2b94ee3fbe588e71d362a941ebb545dec29b53961d61add2a1"},
 ]
+markers = {main = "extra == \"extra-proxy\""}
 
 [[package]]
 name = "tornado"
@@ -7994,4 +8018,4 @@ utils = ["numpydoc"]
 [metadata]
 lock-version = "2.1"
 python-versions = ">=3.9,<4.0"
-content-hash = "1f3bbf967451633fb6290ba88980bdf4fbf83420024b14e862d1da717d903684"
+content-hash = "eda34dfd8b35474beffee18893d6782c7b3d0d3d2c610f66237eb97176f43527"

pyproject.toml

@@ -61,7 +61,7 @@ boto3 = { version = "^1.40.76", optional = true }
 redisvl = {version = "^0.4.1", optional = true, markers = "python_version >= '3.9' and python_version < '3.14'"}
 mcp = {version = ">=1.25.0,<2.0.0", optional = true, python = ">=3.10"}
 a2a-sdk = {version = "^0.3.22", optional = true, python = ">=3.10"}
-litellm-proxy-extras = {version = "^0.4.57", optional = true}
+litellm-proxy-extras = {version = "^0.4.58", optional = true}
 rich = {version = "^13.7.1", optional = true}
 litellm-enterprise = {version = "^0.1.33", optional = true}
 diskcache = {version = "^5.6.1", optional = true}
@@ -167,6 +167,7 @@ langfuse = "^2.45.0"
 fastapi-offline = "^1.7.3"
 fakeredis = "^2.27.1"
 pytest-rerunfailures = "^14.0"
+parameterized = "^0.9.0"
 
 [tool.poetry.group.proxy-dev.dependencies]
 prisma = "0.11.0"

requirements.txt

@@ -57,7 +57,7 @@ grpcio>=1.75.0; python_version >= "3.14"
 sentry_sdk==2.21.0 # for sentry error handling
 detect-secrets==1.5.0 # Enterprise - secret detection / masking in LLM requests
 tzdata==2025.1 # IANA time zone database
-litellm-proxy-extras==0.4.57 # for proxy extras - e.g. prisma migrations
+litellm-proxy-extras==0.4.58 # for proxy extras - e.g. prisma migrations
 llm-sandbox==0.3.31 # for skill execution in sandbox
 ### LITELLM PACKAGE DEPENDENCIES
 python-dotenv==1.0.1 # for env

schema.prisma

@@ -146,6 +146,10 @@ model LiteLLM_TeamTable {
     litellm_model_table LiteLLM_ModelTable? @relation(fields: [model_id], references: [id])
     object_permission LiteLLM_ObjectPermissionTable?   @relation(fields: [object_permission_id], references: [object_permission_id])
     projects LiteLLM_ProjectTable[]
+
+    @@index([organization_id])
+    @@index([team_alias])
+    @@index([created_at])
 }
 
 // Projects sit between teams and keys for use-case management

scripts/install.sh

@@ -0,0 +1,143 @@
+#!/usr/bin/env bash
+# LiteLLM Installer
+# Usage: curl -fsSL https://raw.githubusercontent.com/BerriAI/litellm/main/scripts/install.sh | sh
+#
+# NOTE: set -e without pipefail for POSIX sh compatibility (dash on Ubuntu/Debian
+# ignores the shebang when invoked as `sh` and does not support `pipefail`).
+set -eu
+
+MIN_PYTHON_MAJOR=3
+MIN_PYTHON_MINOR=9
+
+# NOTE: before merging, this must stay as "litellm[proxy]" to install from PyPI.
+LITELLM_PACKAGE="litellm[proxy]"
+
+# ── colours ────────────────────────────────────────────────────────────────
+if [ -t 1 ]; then
+  BOLD='\033[1m'
+  GREEN='\033[38;2;78;186;101m'
+  GREY='\033[38;2;153;153;153m'
+  RESET='\033[0m'
+else
+  BOLD='' GREEN='' GREY='' RESET=''
+fi
+
+info()    { printf "${GREY}  %s${RESET}\n" "$*"; }
+success() { printf "${GREEN}  ✔ %s${RESET}\n" "$*"; }
+header()  { printf "${BOLD}  %s${RESET}\n" "$*"; }
+die()     { printf "\n  Error: %s\n\n" "$*" >&2; exit 1; }
+
+# ── banner ─────────────────────────────────────────────────────────────────
+echo ""
+cat << 'EOF'
+  ██╗     ██╗████████╗███████╗██╗     ██╗     ███╗   ███╗
+  ██║     ██║╚══██╔══╝██╔════╝██║     ██║     ████╗ ████║
+  ██║     ██║   ██║   █████╗  ██║     ██║     ██╔████╔██║
+  ██║     ██║   ██║   ██╔══╝  ██║     ██║     ██║╚██╔╝██║
+  ███████╗██║   ██║   ███████╗███████╗███████╗██║ ╚═╝ ██║
+  ╚══════╝╚═╝   ╚═╝   ╚══════╝╚══════╝╚══════╝╚═╝     ╚═╝
+EOF
+printf "  ${BOLD}LiteLLM Installer${RESET}  ${GREY}— unified gateway for 100+ LLM providers${RESET}\n\n"
+
+# ── OS detection ───────────────────────────────────────────────────────────
+OS="$(uname -s)"
+ARCH="$(uname -m)"
+
+case "$OS" in
+  Darwin)  PLATFORM="macOS ($ARCH)" ;;
+  Linux)   PLATFORM="Linux ($ARCH)" ;;
+  *)       die "Unsupported OS: $OS. LiteLLM supports macOS and Linux." ;;
+esac
+
+info "Platform: $PLATFORM"
+
+# ── Python detection ───────────────────────────────────────────────────────
+PYTHON_BIN=""
+for candidate in python3 python; do
+  if command -v "$candidate" >/dev/null 2>&1; then
+    major="$("$candidate" -c 'import sys; print(sys.version_info.major)' 2>/dev/null || true)"
+    minor="$("$candidate" -c 'import sys; print(sys.version_info.minor)' 2>/dev/null || true)"
+    if [ "${major:-0}" -ge "$MIN_PYTHON_MAJOR" ] && [ "${minor:-0}" -ge "$MIN_PYTHON_MINOR" ]; then
+      PYTHON_BIN="$(command -v "$candidate")"
+      info "Python: $("$candidate" --version 2>&1)"
+      break
+    fi
+  fi
+done
+
+if [ -z "$PYTHON_BIN" ]; then
+  die "Python ${MIN_PYTHON_MAJOR}.${MIN_PYTHON_MINOR}+ is required but not found.
+  Install it from https://python.org/downloads or via your package manager:
+    macOS:  brew install python@3
+    Ubuntu: sudo apt install python3 python3-pip"
+fi
+
+# ── pip detection ──────────────────────────────────────────────────────────
+if ! "$PYTHON_BIN" -m pip --version >/dev/null 2>&1; then
+  die "pip is not available. Install it with:
+    $PYTHON_BIN -m ensurepip --upgrade"
+fi
+
+# ── install ────────────────────────────────────────────────────────────────
+echo ""
+header "Installing litellm[proxy]…"
+echo ""
+
+"$PYTHON_BIN" -m pip install --upgrade "${LITELLM_PACKAGE}" \
+  || die "pip install failed. Try manually: $PYTHON_BIN -m pip install '${LITELLM_PACKAGE}'"
+
+# ── find the litellm binary installed by pip for this Python ───────────────
+# sysconfig.get_path('scripts') is where pip puts console scripts — reliable
+# even when the Python lives in a libexec/ symlink tree (e.g. Homebrew).
+SCRIPTS_DIR="$("$PYTHON_BIN" -c 'import sysconfig; print(sysconfig.get_path("scripts"))')"
+LITELLM_BIN="${SCRIPTS_DIR}/litellm"
+
+if [ ! -x "$LITELLM_BIN" ]; then
+  # Fall back to user-base bin (pip install --user)
+  USER_BIN="$("$PYTHON_BIN" -c 'import site; print(site.getuserbase())')/bin"
+  LITELLM_BIN="${USER_BIN}/litellm"
+fi
+
+if [ ! -x "$LITELLM_BIN" ]; then
+  die "litellm binary not found after install. Try: $PYTHON_BIN -m pip install --user '${LITELLM_PACKAGE}'"
+fi
+
+# ── success banner ─────────────────────────────────────────────────────────
+echo ""
+success "LiteLLM installed"
+
+installed_ver="$("$LITELLM_BIN" --version 2>&1 | grep -oE '[0-9]+\.[0-9]+\.[0-9]+' | head -1 || true)"
+[ -n "$installed_ver" ] && info "Version: $installed_ver"
+
+# ── PATH hint ──────────────────────────────────────────────────────────────
+if ! command -v litellm >/dev/null 2>&1; then
+  info "Note: add litellm to your PATH:  export PATH=\"\$PATH:${SCRIPTS_DIR}\""
+fi
+
+# ── launch setup wizard ────────────────────────────────────────────────────
+echo ""
+printf "  ${BOLD}Run the interactive setup wizard?${RESET} ${GREY}(Y/n)${RESET}: "
+# /dev/tty may be unavailable in Docker/CI — default to yes if it can't be read
+answer=""
+if [ -r /dev/tty ]; then
+  read -r answer </dev/tty || answer=""
+fi
+
+if [ -z "$answer" ] || [ "$answer" = "y" ] || [ "$answer" = "Y" ]; then
+  echo ""
+  # Use /dev/tty for interactive input when available (stdin is a pipe from curl)
+  if [ -r /dev/tty ]; then
+    exec "$LITELLM_BIN" --setup </dev/tty
+  else
+    exec "$LITELLM_BIN" --setup
+  fi
+else
+  echo ""
+  header "Quick start:"
+  echo ""
+  info "  litellm --setup          # interactive wizard"
+  info "  litellm --model gpt-4o   # single-model quickstart"
+  echo ""
+  info "Docs: https://docs.litellm.ai"
+  echo ""
+fi

tests/local_testing/test_router.py

@@ -907,8 +907,9 @@ def test_router_region_pre_call_check(allowed_model_region):
         {
             "model_name": "gpt-3.5-turbo-large",  # openai model name
             "litellm_params": {  # params for litellm completion/embedding call
-                "model": "gpt-3.5-turbo-1106",
+                "model": "gpt-4.1-mini",
                 "api_key": os.getenv("OPENAI_API_KEY"),
+                "mock_response": "This is a mock response.",
             },
             "model_info": {"id": "2"},
         },

tests/logging_callback_tests/test_langfuse_e2e_test.py

@@ -60,10 +60,33 @@ def assert_langfuse_request_matches_expected(
             )
         ]
 
+    # When aggregating from multiple flush cycles, deduplicate by keeping
+    # only one trace-create and one generation-create per trace_id.
+    seen_types: dict = {}
+    deduped_batch: list = []
+    for item in actual_request_body["batch"]:
+        item_type = item["type"]
+        if item_type not in seen_types:
+            seen_types[item_type] = True
+            deduped_batch.append(item)
+    actual_request_body["batch"] = deduped_batch
+
+    # Ensure canonical order: trace-create first, generation-create second
+    actual_request_body["batch"].sort(
+        key=lambda x: 0 if x["type"] == "trace-create" else 1
+    )
+
     print(
         "actual_request_body after filtering", json.dumps(actual_request_body, indent=4)
     )
 
+    assert len(actual_request_body["batch"]) >= 2, (
+        f"Expected at least 2 batch items (trace-create + generation-create) "
+        f"after filtering by trace_id={trace_id}, "
+        f"but got {len(actual_request_body['batch'])}. "
+        f"Items: {json.dumps(actual_request_body['batch'], indent=2)}"
+    )
+
     # Replace dynamic values in actual request body
     for item in actual_request_body["batch"]:
 
@@ -150,19 +173,36 @@ class TestLangfuseLogging:
         """Helper method to verify Langfuse API calls"""
         await asyncio.sleep(3)
 
-        # Verify the call
+        # Verify at least one call was made
         assert mock_post.call_count >= 1
-        url = mock_post.call_args[0][0]
-        request_body = mock_post.call_args[1].get("content")
 
-        # Parse the JSON string into a dict for assertions
-        actual_request_body = json.loads(request_body)
+        # Aggregate batch items from ALL calls — the Langfuse SDK may split
+        # trace-create and generation-create across separate HTTP flushes.
+        langfuse_url = "https://us.cloud.langfuse.com/api/public/ingestion"
+        all_batch_items: list = []
+        metadata: Optional[dict] = None
+        for call in mock_post.call_args_list:
+            url = call[0][0]
+            if url != langfuse_url:
+                continue
+            request_body = call[1].get("content")
+            if request_body:
+                body = json.loads(request_body)
+                all_batch_items.extend(body.get("batch", []))
+                if metadata is None:
+                    metadata = body.get("metadata")
 
-        print("\nMocked Request Details:")
-        print(f"URL: {url}")
+        assert len(all_batch_items) > 0, "No Langfuse ingestion calls found"
+        assert metadata is not None, "No metadata found in Langfuse calls"
+
+        actual_request_body = {
+            "batch": all_batch_items,
+            "metadata": metadata,
+        }
+
+        print("\nMocked Request Details (aggregated from all calls):")
         print(f"Request Body: {json.dumps(actual_request_body, indent=4)}")
 
-        assert url == "https://us.cloud.langfuse.com/api/public/ingestion"
         assert_langfuse_request_matches_expected(
             actual_request_body,
             expected_file_name,
@@ -170,6 +210,7 @@ class TestLangfuseLogging:
         )
 
     @pytest.mark.asyncio
+    @pytest.mark.flaky(retries=3, delay=1)
     async def test_langfuse_logging_completion(self, mock_setup):
         """Test Langfuse logging for chat completion"""
         setup = mock_setup
@@ -185,6 +226,7 @@ class TestLangfuseLogging:
             )
 
     @pytest.mark.asyncio
+    @pytest.mark.flaky(retries=3, delay=1)
     async def test_langfuse_logging_completion_with_tags(self, mock_setup):
         """Test Langfuse logging for chat completion with tags"""
         setup = mock_setup
@@ -203,6 +245,7 @@ class TestLangfuseLogging:
             )
 
     @pytest.mark.asyncio
+    @pytest.mark.flaky(retries=3, delay=1)
     async def test_langfuse_logging_completion_with_tags_stream(self, mock_setup):
         """Test Langfuse logging for chat completion with tags"""
         setup = mock_setup
@@ -223,6 +266,7 @@ class TestLangfuseLogging:
             )
 
     @pytest.mark.asyncio
+    @pytest.mark.flaky(retries=3, delay=1)
     async def test_langfuse_logging_completion_with_langfuse_metadata(self, mock_setup):
         """Test Langfuse logging for chat completion with metadata for langfuse"""
         setup = mock_setup
@@ -252,6 +296,7 @@ class TestLangfuseLogging:
             )
 
     @pytest.mark.asyncio
+    @pytest.mark.flaky(retries=3, delay=1)
     async def test_langfuse_logging_with_non_serializable_metadata(self, mock_setup):
         """Test Langfuse logging with metadata that requires preparation (Pydantic models, sets, etc)"""
         from pydantic import BaseModel
@@ -358,6 +403,7 @@ class TestLangfuseLogging:
             )
 
     @pytest.mark.asyncio
+    @pytest.mark.flaky(retries=3, delay=1)
     async def test_langfuse_logging_completion_with_malformed_llm_response(
         self, mock_setup
     ):
@@ -387,6 +433,7 @@ class TestLangfuseLogging:
             )
     
     @pytest.mark.asyncio
+    @pytest.mark.flaky(retries=3, delay=1)
     async def test_langfuse_logging_completion_with_bedrock_llm_response(
         self, mock_setup
     ):
@@ -418,6 +465,7 @@ class TestLangfuseLogging:
                 setup["mock_post"], "completion_with_bedrock_call.json", setup["trace_id"]
             )
     @pytest.mark.asyncio
+    @pytest.mark.flaky(retries=3, delay=1)
     async def test_langfuse_logging_completion_with_vertex_llm_response(
         self, mock_setup
     ):
@@ -449,6 +497,7 @@ class TestLangfuseLogging:
             )
 
     @pytest.mark.asyncio
+    @pytest.mark.flaky(retries=3, delay=1)
     async def test_langfuse_logging_vllm_embedding(self, mock_setup):
         """
         Test that the request sent to the vllm embedding endpoint is correct.
@@ -500,6 +549,7 @@ class TestLangfuseLogging:
         )
 
     @pytest.mark.asyncio
+    @pytest.mark.flaky(retries=3, delay=1)
     async def test_langfuse_logging_with_router(self, mock_setup):
         """Test Langfuse logging with router"""
         litellm._turn_on_debug()

tests/spend_tracking_tests/test_spend_accuracy_tests.py

@@ -308,16 +308,19 @@ async def test_long_term_spend_accuracy_with_bursts():
             response = await chat_completion(session, key)
             print(f"Burst 2 - Request {i + 1}/{BURST_2_REQUESTS} completed")
 
-        # Poll until key spend reflects burst 2
-        burst_1_spend = intermediate_key_info["info"]["spend"]
+        # Poll until key spend reaches expected total (burst 1 + burst 2)
         start = time.time()
         while time.time() - start < 120:
             key_info_check = await get_spend_info(session, "key", key)
             current_spend = key_info_check["info"]["spend"]
-            if current_spend > burst_1_spend:
-                print(f"Key spend increased to {current_spend} after {time.time() - start:.1f}s")
+            if abs(current_spend - expected_spend) < TOLERANCE:
+                print(
+                    f"Total spend reached expected {expected_spend} after {time.time() - start:.1f}s"
+                )
                 break
-            print(f"Key spend still {current_spend}, waiting for burst 2 flush...")
+            print(
+                f"Key spend {current_spend}, expected {expected_spend}, waiting..."
+            )
             await asyncio.sleep(10)
 
         # Allow extra time for all entity spend aggregations

tests/test_litellm/integrations/test_opentelemetry.py

@@ -3,7 +3,8 @@ import os
 import sys
 import time
 import unittest
-from datetime import datetime, timedelta
+from datetime import datetime, timedelta, timezone
+from parameterized import parameterized
 from unittest.mock import MagicMock, patch
 
 # Adds the grandparent directory to sys.path to allow importing project modules
@@ -425,6 +426,7 @@ class TestOpenTelemetry(unittest.TestCase):
         ) as mock_get_headers, patch.object(
             otel, "_get_tracer_with_dynamic_headers"
         ) as mock_get_tracer:
+
             # Test case 1: With dynamic headers
             mock_get_headers.return_value = {
                 "arize-space-id": "test-space",
@@ -2165,25 +2167,30 @@ class TestOpenTelemetrySemanticConventions138(unittest.TestCase):
 
         mock_span.set_attribute.assert_any_call("gen_ai.operation.name", "chat")
 
-    def test_handle_failure_langfuse_otel_nulls_parent_span(self):
+    @parameterized.expand([("_handle_success",), ("_handle_failure",)])
+    def test_handle_success_failure_nulls_parent_span_if_ignore_context_propagation(
+        self, handle_method: str
+    ):
         """
-        For langfuse_otel, _handle_failure should ignore parent spans from other providers
-        and create a root-level error span (symmetric with _handle_success).
+        If ignore_context_propagation is True, _handle_success should ignore any parent span
+        and create a root-level span. This could be useful for langfuse_otel where
+        _handle_success may ignore parent spans from other providers and create a root-level
+        span (symmetric with _handle_failure).
         """
         span_exporter = InMemorySpanExporter()
         tracer_provider = TracerProvider()
         tracer_provider.add_span_processor(SimpleSpanProcessor(span_exporter))
 
         otel = OpenTelemetry(
-            callback_name="langfuse_otel",
+            config=OpenTelemetryConfig(ignore_context_propagation=True),
             tracer_provider=tracer_provider,
         )
         otel.tracer = tracer_provider.get_tracer("litellm")
 
         other_tracer = tracer_provider.get_tracer("other_provider")
-        other_span = other_tracer.start_span("other_provider_span")
+        other_span = other_tracer.start_span("parent_span")
 
-        start = datetime.utcnow()
+        start = datetime.now(timezone.utc)
         end = start + timedelta(seconds=1)
 
         kwargs = {
@@ -2202,23 +2209,33 @@ class TestOpenTelemetrySemanticConventions138(unittest.TestCase):
             "exception": Exception("test error"),
         }
 
-        otel._handle_failure(kwargs, None, start, end)
+        with patch.dict(os.environ, {"USE_OTEL_LITELLM_REQUEST_SPAN": "true"}):
+            if handle_method == "_handle_success":
+                otel._handle_success(kwargs, None, start, end)
+            elif handle_method == "_handle_failure":
+                otel._handle_failure(kwargs, None, start, end)
+            else:
+                self.fail(f"Invalid handle_method: {handle_method}")
 
         other_span.end()
 
         spans = span_exporter.get_finished_spans()
-        failure_spans = [s for s in spans if s.name != "other_provider_span"]
+        child_spans = [s for s in spans if s.name != "parent_span"]
+        child_span_ids = {s.context.span_id for s in child_spans if s.context}
 
-        self.assertTrue(failure_spans, "Expected at least one failure span")
-        for span in failure_spans:
-            self.assertIsNone(
-                span.parent,
-                f"langfuse_otel failure span should be a root span, but has parent: {span.parent}",
-            )
+        self.assertTrue(child_spans, "Expected at least one child span")
+        for span in child_spans:
+            assert (
+                span.parent is None or span.parent.span_id in child_span_ids
+            ), f"if ignore_context_propagation is True, span should not have parent from other providers, but got parent: {span.parent}"
 
-    def test_handle_failure_non_langfuse_preserves_parent_span(self):
+    @parameterized.expand([("_handle_success",), ("_handle_failure",)])
+    def test_handle_success_failure_default_preserves_parent_span(
+        self, handle_method: str
+    ):
         """
-        For non-langfuse_otel callbacks, _handle_failure should still use parent spans normally.
+        For default otel callbacks, _handle_success should use parent spans normally.
+        (symmetric with _handle_failure)
         """
         span_exporter = InMemorySpanExporter()
         tracer_provider = TracerProvider()
@@ -2229,7 +2246,7 @@ class TestOpenTelemetrySemanticConventions138(unittest.TestCase):
 
         parent_span = otel.tracer.start_span("parent_span")
 
-        start = datetime.utcnow()
+        start = datetime.now(timezone.utc)
         end = start + timedelta(seconds=1)
 
         kwargs = {
@@ -2249,19 +2266,81 @@ class TestOpenTelemetrySemanticConventions138(unittest.TestCase):
         }
 
         with patch.dict(os.environ, {"USE_OTEL_LITELLM_REQUEST_SPAN": "true"}):
-            otel._handle_failure(kwargs, None, start, end)
+            if handle_method == "_handle_success":
+                otel._handle_success(kwargs, None, start, end)
+            elif handle_method == "_handle_failure":
+                otel._handle_failure(kwargs, None, start, end)
+            else:
+                self.fail(f"Invalid handle_method: {handle_method}")
 
         parent_span.end()
 
         spans = span_exporter.get_finished_spans()
         child_spans = [s for s in spans if s.name != "parent_span"]
 
-        self.assertTrue(child_spans, "Expected at least one child failure span")
+        self.assertTrue(child_spans, "Expected at least one child span")
         for span in child_spans:
-            self.assertIsNotNone(
-                span.parent,
-                "Non-langfuse_otel failure span should have a parent",
-            )
+            assert (
+                span.parent is not None
+            ), f"By default parent span should be preserved, but got None parent for span: {span.name}"
+
+    @parameterized.expand([("_handle_success",), ("_handle_failure",)])
+    def test_handle_success_failure_with_context_propagation_preserves_parent_span(
+        self, handle_method: str
+    ):
+        """
+        For otel callbacks with context propagation enabled, _handle_success should
+        use parent spans normally. (symmetric with _handle_failure)
+        """
+        span_exporter = InMemorySpanExporter()
+        tracer_provider = TracerProvider()
+        tracer_provider.add_span_processor(SimpleSpanProcessor(span_exporter))
+
+        otel = OpenTelemetry(
+            config=OpenTelemetryConfig(ignore_context_propagation=False),
+            tracer_provider=tracer_provider,
+        )
+        otel.tracer = tracer_provider.get_tracer("litellm")
+
+        parent_span = otel.tracer.start_span("parent_span")
+
+        start = datetime.now(timezone.utc)
+        end = start + timedelta(seconds=1)
+
+        kwargs = {
+            "model": "gpt-4",
+            "messages": [{"role": "user", "content": "Hello"}],
+            "optional_params": {},
+            "litellm_params": {
+                "custom_llm_provider": "openai",
+                "metadata": {"litellm_parent_otel_span": parent_span},
+            },
+            "standard_logging_object": {
+                "id": "test-id",
+                "call_type": "completion",
+                "metadata": {},
+            },
+            "exception": Exception("test error"),
+        }
+
+        with patch.dict(os.environ, {"USE_OTEL_LITELLM_REQUEST_SPAN": "true"}):
+            if handle_method == "_handle_success":
+                otel._handle_success(kwargs, None, start, end)
+            elif handle_method == "_handle_failure":
+                otel._handle_failure(kwargs, None, start, end)
+            else:
+                self.fail(f"Invalid handle_method: {handle_method}")
+
+        parent_span.end()
+
+        spans = span_exporter.get_finished_spans()
+        child_spans = [s for s in spans if s.name != "parent_span"]
+
+        self.assertTrue(child_spans, "Expected at least one child span")
+        for span in child_spans:
+            assert (
+                span.parent is not None
+            ), f"If ignore_context_propagation is False, parent span should be preserved, but got None parent for span: {span.name}"
 
     def test_handle_failure_hasattr_guard_on_parent_name(self):
         """
@@ -2431,9 +2510,7 @@ class TestNoParentSpanDuplication(unittest.TestCase):
         otel._handle_success(kwargs, response_obj, start, end)
 
         spans = span_exporter.get_finished_spans()
-        proxy_spans = [
-            s for s in spans if s.name == LITELLM_PROXY_REQUEST_SPAN_NAME
-        ]
+        proxy_spans = [s for s in spans if s.name == LITELLM_PROXY_REQUEST_SPAN_NAME]
         self.assertEqual(len(proxy_spans), 1, "Should have exactly one proxy span")
 
         proxy_attrs = proxy_spans[0].attributes or {}

tests/test_litellm/interactions/test_openapi_compliance.py

@@ -77,6 +77,11 @@ class TestRequestCompliance:
         schema = spec_dict["components"]["schemas"]["CreateModelInteractionParams"]
         input_schema = schema["properties"]["input"]
         
+        # The input property may be inline oneOf or a $ref to InteractionsInput
+        if "$ref" in input_schema:
+            ref_name = input_schema["$ref"].split("/")[-1]
+            input_schema = spec_dict["components"]["schemas"][ref_name]
+
         # Should be oneOf with multiple types
         assert "oneOf" in input_schema
         
@@ -100,10 +105,21 @@ class TestRequestCompliance:
         assert "discriminator" in content_schema
         assert content_schema["discriminator"]["propertyName"] == "type"
         
-        # Check TextContent is an option
-        mapping = content_schema["discriminator"]["mapping"]
-        assert "text" in mapping
-        print(f"Content type discriminator mapping: {list(mapping.keys())}")
+        # Check TextContent is an option (via mapping if present, or via oneOf refs)
+        mapping = content_schema["discriminator"].get("mapping")
+        if mapping:
+            assert "text" in mapping
+            print(f"Content type discriminator mapping: {list(mapping.keys())}")
+        else:
+            # Discriminator without explicit mapping — verify via oneOf
+            one_of = content_schema.get("oneOf", [])
+            ref_names = [
+                opt["$ref"].split("/")[-1] for opt in one_of if "$ref" in opt
+            ]
+            assert "TextContent" in ref_names, (
+                f"TextContent not found in oneOf refs: {ref_names}"
+            )
+            print(f"Content type discriminator (no mapping), oneOf refs: {ref_names}")
 
     def test_text_content_schema(self, spec_dict):
         """Verify TextContent schema."""

tests/test_litellm/llms/xai/test_xai_cost_calculator.py

@@ -386,8 +386,56 @@ class TestXAICostCalculator:
             completion_tokens=50,
             total_tokens=150,
         )
-        
+
         web_search_cost = cost_per_web_search_request(usage=usage, model_info={})
-        
+
         # Expected cost: No web search data = $0.0
         assert web_search_cost == 0.0
+
+    def test_grok_4_20_beta_reasoning_cost_calculation(self):
+        """Test cost calculation for grok-4.20-beta-0309-reasoning model."""
+        usage = Usage(prompt_tokens=100, completion_tokens=200, total_tokens=300)
+
+        prompt_cost, completion_cost = cost_per_token(
+            model="grok-4.20-beta-0309-reasoning", usage=usage
+        )
+
+        # Input: 100 tokens * $2e-6 = $0.0002
+        # Output: 200 tokens * $6e-6 = $0.0012
+        expected_prompt_cost = 100 * 2e-6
+        expected_completion_cost = 200 * 6e-6
+
+        assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-10)
+        assert math.isclose(completion_cost, expected_completion_cost, rel_tol=1e-10)
+
+    def test_grok_4_20_beta_non_reasoning_cost_calculation(self):
+        """Test cost calculation for grok-4.20-beta-0309-non-reasoning model."""
+        usage = Usage(prompt_tokens=50, completion_tokens=100, total_tokens=150)
+
+        prompt_cost, completion_cost = cost_per_token(
+            model="grok-4.20-beta-0309-non-reasoning", usage=usage
+        )
+
+        # Input: 50 tokens * $2e-6 = $0.0001
+        # Output: 100 tokens * $6e-6 = $0.0006
+        expected_prompt_cost = 50 * 2e-6
+        expected_completion_cost = 100 * 6e-6
+
+        assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-10)
+        assert math.isclose(completion_cost, expected_completion_cost, rel_tol=1e-10)
+
+    def test_grok_4_20_multi_agent_cost_calculation(self):
+        """Test cost calculation for grok-4.20-multi-agent-beta-0309 model."""
+        usage = Usage(prompt_tokens=200, completion_tokens=300, total_tokens=500)
+
+        prompt_cost, completion_cost = cost_per_token(
+            model="grok-4.20-multi-agent-beta-0309", usage=usage
+        )
+
+        # Input: 200 tokens * $2e-6 = $0.0004
+        # Output: 300 tokens * $6e-6 = $0.0018
+        expected_prompt_cost = 200 * 2e-6
+        expected_completion_cost = 300 * 6e-6
+
+        assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-10)
+        assert math.isclose(completion_cost, expected_completion_cost, rel_tol=1e-10)

tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py

@@ -0,0 +1,707 @@
+"""
+Tests for pre_mcp_call guardrail hook header mutation support.
+
+Validates that:
+1. _convert_mcp_hook_response_to_kwargs extracts extra_headers from hook response
+2. pre_call_tool_check returns hook-provided extra_headers AND modified arguments
+3. call_tool flows hook headers and modified arguments downstream
+4. Hook-provided headers take highest priority (merge after static_headers)
+5. OpenAPI-backed servers log a warning and continue (skip injection) when hook headers are present
+6. JWT claims are propagated in both standard and virtual-key fast paths
+7. Backward compatibility: hooks without extra_headers continue to work
+"""
+
+import asyncio
+from typing import Any, Dict, Optional
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+from fastapi import HTTPException
+
+from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager
+from litellm.proxy._types import UserAPIKeyAuth
+from litellm.proxy.utils import ProxyLogging
+from litellm.types.mcp import MCPAuth, MCPTransport
+from litellm.types.mcp_server.mcp_server_manager import MCPServer
+
+
+class TestConvertMcpHookResponseToKwargs:
+    """Tests for ProxyLogging._convert_mcp_hook_response_to_kwargs"""
+
+    def setup_method(self):
+        self.proxy_logging = ProxyLogging(user_api_key_cache=MagicMock())
+
+    def test_returns_original_kwargs_when_response_is_none(self):
+        original = {"arguments": {"key": "val"}, "name": "tool"}
+        result = self.proxy_logging._convert_mcp_hook_response_to_kwargs(
+            None, original
+        )
+        assert result == original
+
+    def test_returns_original_kwargs_when_response_is_empty_dict(self):
+        original = {"arguments": {"key": "val"}}
+        result = self.proxy_logging._convert_mcp_hook_response_to_kwargs({}, original)
+        assert result == original
+
+    def test_extracts_modified_arguments(self):
+        original = {"arguments": {"old": "value"}}
+        response = {"modified_arguments": {"new": "value"}}
+        result = self.proxy_logging._convert_mcp_hook_response_to_kwargs(
+            response, original
+        )
+        assert result["arguments"] == {"new": "value"}
+
+    def test_extracts_extra_headers(self):
+        original = {"arguments": {"key": "val"}}
+        response = {"extra_headers": {"Authorization": "Bearer signed-jwt"}}
+        result = self.proxy_logging._convert_mcp_hook_response_to_kwargs(
+            response, original
+        )
+        assert result["extra_headers"] == {"Authorization": "Bearer signed-jwt"}
+
+    def test_extracts_both_arguments_and_headers(self):
+        original = {"arguments": {"old": "value"}}
+        response = {
+            "modified_arguments": {"new": "value"},
+            "extra_headers": {"X-Custom": "header-val"},
+        }
+        result = self.proxy_logging._convert_mcp_hook_response_to_kwargs(
+            response, original
+        )
+        assert result["arguments"] == {"new": "value"}
+        assert result["extra_headers"] == {"X-Custom": "header-val"}
+
+    def test_no_extra_headers_key_preserves_original(self):
+        """Backward compat: hooks that only return modified_arguments still work."""
+        original = {"arguments": {"key": "val"}}
+        response = {"modified_arguments": {"key": "new_val"}}
+        result = self.proxy_logging._convert_mcp_hook_response_to_kwargs(
+            response, original
+        )
+        assert "extra_headers" not in result
+        assert result["arguments"] == {"key": "new_val"}
+
+    def test_empty_extra_headers_not_set(self):
+        """Empty dict for extra_headers is falsy and should not be set."""
+        original = {"arguments": {"key": "val"}}
+        response = {"extra_headers": {}}
+        result = self.proxy_logging._convert_mcp_hook_response_to_kwargs(
+            response, original
+        )
+        assert "extra_headers" not in result
+
+
+class TestPreCallToolCheckReturnsHeaders:
+    """Tests that pre_call_tool_check returns hook-provided headers."""
+
+    def _make_server(self, name="test_server"):
+        return MCPServer(
+            server_id="test-id",
+            name=name,
+            server_name=name,
+            url="https://example.com",
+            transport=MCPTransport.http,
+            auth_type=MCPAuth.none,
+        )
+
+    @pytest.mark.asyncio
+    async def test_returns_empty_dict_when_hook_has_no_headers(self):
+        manager = MCPServerManager()
+        server = self._make_server()
+
+        proxy_logging = MagicMock(spec=ProxyLogging)
+        proxy_logging._create_mcp_request_object_from_kwargs = MagicMock(
+            return_value=MagicMock()
+        )
+        proxy_logging._convert_mcp_to_llm_format = MagicMock(
+            return_value={"model": "fake"}
+        )
+        proxy_logging.pre_call_hook = AsyncMock(
+            return_value={"modified_arguments": {"key": "val"}}
+        )
+        proxy_logging._convert_mcp_hook_response_to_kwargs = MagicMock(
+            return_value={"arguments": {"key": "val"}}
+        )
+
+        with patch.object(manager, "check_allowed_or_banned_tools", return_value=True):
+            with patch.object(
+                manager,
+                "check_tool_permission_for_key_team",
+                new_callable=AsyncMock,
+            ):
+                with patch.object(manager, "validate_allowed_params"):
+                    result = await manager.pre_call_tool_check(
+                        name="test_tool",
+                        arguments={"key": "val"},
+                        server_name="test_server",
+                        user_api_key_auth=None,
+                        proxy_logging_obj=proxy_logging,
+                        server=server,
+                    )
+
+        assert result == {}
+
+    @pytest.mark.asyncio
+    async def test_returns_extra_headers_from_hook(self):
+        manager = MCPServerManager()
+        server = self._make_server()
+
+        hook_headers = {"Authorization": "Bearer signed-jwt", "X-Trace-Id": "abc123"}
+
+        proxy_logging = MagicMock(spec=ProxyLogging)
+        proxy_logging._create_mcp_request_object_from_kwargs = MagicMock(
+            return_value=MagicMock()
+        )
+        proxy_logging._convert_mcp_to_llm_format = MagicMock(
+            return_value={"model": "fake"}
+        )
+        proxy_logging.pre_call_hook = AsyncMock(
+            return_value={"extra_headers": hook_headers}
+        )
+        proxy_logging._convert_mcp_hook_response_to_kwargs = MagicMock(
+            return_value={"arguments": {"key": "val"}, "extra_headers": hook_headers}
+        )
+
+        with patch.object(manager, "check_allowed_or_banned_tools", return_value=True):
+            with patch.object(
+                manager,
+                "check_tool_permission_for_key_team",
+                new_callable=AsyncMock,
+            ):
+                with patch.object(manager, "validate_allowed_params"):
+                    result = await manager.pre_call_tool_check(
+                        name="test_tool",
+                        arguments={"key": "val"},
+                        server_name="test_server",
+                        user_api_key_auth=None,
+                        proxy_logging_obj=proxy_logging,
+                        server=server,
+                    )
+
+        assert result["extra_headers"] == hook_headers
+
+    @pytest.mark.asyncio
+    async def test_returns_empty_dict_when_hook_returns_none(self):
+        manager = MCPServerManager()
+        server = self._make_server()
+
+        proxy_logging = MagicMock(spec=ProxyLogging)
+        proxy_logging._create_mcp_request_object_from_kwargs = MagicMock(
+            return_value=MagicMock()
+        )
+        proxy_logging._convert_mcp_to_llm_format = MagicMock(
+            return_value={"model": "fake"}
+        )
+        proxy_logging.pre_call_hook = AsyncMock(return_value=None)
+
+        with patch.object(manager, "check_allowed_or_banned_tools", return_value=True):
+            with patch.object(
+                manager,
+                "check_tool_permission_for_key_team",
+                new_callable=AsyncMock,
+            ):
+                with patch.object(manager, "validate_allowed_params"):
+                    result = await manager.pre_call_tool_check(
+                        name="test_tool",
+                        arguments={"key": "val"},
+                        server_name="test_server",
+                        user_api_key_auth=None,
+                        proxy_logging_obj=proxy_logging,
+                        server=server,
+                    )
+
+        assert result == {}
+
+    @pytest.mark.asyncio
+    async def test_returns_modified_arguments_from_hook(self):
+        """Modified arguments from the hook must be returned so the caller can use them."""
+        manager = MCPServerManager()
+        server = self._make_server()
+
+        original_args = {"key": "original"}
+        modified_args = {"key": "modified", "extra": "added"}
+
+        proxy_logging = MagicMock(spec=ProxyLogging)
+        proxy_logging._create_mcp_request_object_from_kwargs = MagicMock(
+            return_value=MagicMock()
+        )
+        proxy_logging._convert_mcp_to_llm_format = MagicMock(
+            return_value={"model": "fake"}
+        )
+        proxy_logging.pre_call_hook = AsyncMock(
+            return_value={"modified_arguments": modified_args}
+        )
+        proxy_logging._convert_mcp_hook_response_to_kwargs = MagicMock(
+            return_value={"arguments": modified_args}
+        )
+
+        with patch.object(manager, "check_allowed_or_banned_tools", return_value=True):
+            with patch.object(
+                manager,
+                "check_tool_permission_for_key_team",
+                new_callable=AsyncMock,
+            ):
+                with patch.object(manager, "validate_allowed_params"):
+                    result = await manager.pre_call_tool_check(
+                        name="test_tool",
+                        arguments=original_args,
+                        server_name="test_server",
+                        user_api_key_auth=None,
+                        proxy_logging_obj=proxy_logging,
+                        server=server,
+                    )
+
+        assert result["arguments"] == modified_args
+
+    @pytest.mark.asyncio
+    async def test_returns_both_modified_arguments_and_headers(self):
+        """Hook can modify both arguments and inject headers simultaneously."""
+        manager = MCPServerManager()
+        server = self._make_server()
+
+        modified_args = {"key": "modified"}
+        hook_headers = {"Authorization": "Bearer jwt"}
+
+        proxy_logging = MagicMock(spec=ProxyLogging)
+        proxy_logging._create_mcp_request_object_from_kwargs = MagicMock(
+            return_value=MagicMock()
+        )
+        proxy_logging._convert_mcp_to_llm_format = MagicMock(
+            return_value={"model": "fake"}
+        )
+        proxy_logging.pre_call_hook = AsyncMock(return_value={"dummy": True})
+        proxy_logging._convert_mcp_hook_response_to_kwargs = MagicMock(
+            return_value={"arguments": modified_args, "extra_headers": hook_headers}
+        )
+
+        with patch.object(manager, "check_allowed_or_banned_tools", return_value=True):
+            with patch.object(
+                manager,
+                "check_tool_permission_for_key_team",
+                new_callable=AsyncMock,
+            ):
+                with patch.object(manager, "validate_allowed_params"):
+                    result = await manager.pre_call_tool_check(
+                        name="test_tool",
+                        arguments={"key": "original"},
+                        server_name="test_server",
+                        user_api_key_auth=None,
+                        proxy_logging_obj=proxy_logging,
+                        server=server,
+                    )
+
+        assert result["arguments"] == modified_args
+        assert result["extra_headers"] == hook_headers
+
+
+class TestCallToolFlowsHookHeaders:
+    """Tests that call_tool passes hook_extra_headers to _call_regular_mcp_tool."""
+
+    def _make_server(self, name="test_server"):
+        return MCPServer(
+            server_id="test-id",
+            name=name,
+            server_name=name,
+            url="https://example.com",
+            transport=MCPTransport.http,
+            auth_type=MCPAuth.none,
+        )
+
+    @pytest.mark.asyncio
+    async def test_hook_headers_passed_to_call_regular_mcp_tool(self):
+        """Verify that hook_extra_headers kwarg is forwarded."""
+        manager = MCPServerManager()
+        server = self._make_server()
+
+        hook_headers = {"Authorization": "Bearer signed-jwt"}
+
+        with patch.object(
+            manager,
+            "_get_mcp_server_from_tool_name",
+            return_value=server,
+        ):
+            with patch.object(
+                manager,
+                "pre_call_tool_check",
+                new_callable=AsyncMock,
+                return_value={"extra_headers": hook_headers},
+            ):
+                with patch.object(
+                    manager,
+                    "_create_during_hook_task",
+                    return_value=asyncio.create_task(asyncio.sleep(0)),
+                ):
+                    with patch.object(
+                        manager,
+                        "_call_regular_mcp_tool",
+                        new_callable=AsyncMock,
+                        return_value=MagicMock(),
+                    ) as mock_call:
+                        proxy_logging = MagicMock(spec=ProxyLogging)
+
+                        await manager.call_tool(
+                            server_name="test_server",
+                            name="test_tool",
+                            arguments={"key": "val"},
+                            proxy_logging_obj=proxy_logging,
+                        )
+
+                        mock_call.assert_called_once()
+                        call_kwargs = mock_call.call_args
+                        assert call_kwargs.kwargs.get("hook_extra_headers") == hook_headers
+
+    @pytest.mark.asyncio
+    async def test_no_hook_headers_when_no_proxy_logging(self):
+        """Without proxy_logging_obj, no pre_call_tool_check runs."""
+        manager = MCPServerManager()
+        server = self._make_server()
+
+        with patch.object(
+            manager,
+            "_get_mcp_server_from_tool_name",
+            return_value=server,
+        ):
+            with patch.object(
+                manager,
+                "_call_regular_mcp_tool",
+                new_callable=AsyncMock,
+                return_value=MagicMock(),
+            ) as mock_call:
+                await manager.call_tool(
+                    server_name="test_server",
+                    name="test_tool",
+                    arguments={"key": "val"},
+                    proxy_logging_obj=None,
+                )
+
+                mock_call.assert_called_once()
+                call_kwargs = mock_call.call_args
+                assert call_kwargs.kwargs.get("hook_extra_headers") is None
+
+    @pytest.mark.asyncio
+    async def test_modified_arguments_passed_to_downstream(self):
+        """Hook-modified arguments must be used for the actual tool call."""
+        manager = MCPServerManager()
+        server = self._make_server()
+
+        modified_args = {"key": "modified_by_hook"}
+
+        with patch.object(
+            manager,
+            "_get_mcp_server_from_tool_name",
+            return_value=server,
+        ):
+            with patch.object(
+                manager,
+                "pre_call_tool_check",
+                new_callable=AsyncMock,
+                return_value={"arguments": modified_args},
+            ):
+                with patch.object(
+                    manager,
+                    "_create_during_hook_task",
+                    return_value=asyncio.create_task(asyncio.sleep(0)),
+                ):
+                    with patch.object(
+                        manager,
+                        "_call_regular_mcp_tool",
+                        new_callable=AsyncMock,
+                        return_value=MagicMock(),
+                    ) as mock_call:
+                        proxy_logging = MagicMock(spec=ProxyLogging)
+
+                        await manager.call_tool(
+                            server_name="test_server",
+                            name="test_tool",
+                            arguments={"key": "original"},
+                            proxy_logging_obj=proxy_logging,
+                        )
+
+                        mock_call.assert_called_once()
+                        call_kwargs = mock_call.call_args
+                        assert call_kwargs.kwargs.get("arguments") == modified_args
+
+    @pytest.mark.asyncio
+    async def test_openapi_server_warns_and_continues_on_hook_headers(self):
+        """OpenAPI-backed servers log a warning and continue when hook injects headers."""
+        manager = MCPServerManager()
+        server = MCPServer(
+            server_id="test-id",
+            name="openapi_server",
+            server_name="openapi_server",
+            url="https://example.com",
+            transport=MCPTransport.http,
+            auth_type=MCPAuth.none,
+            spec_path="/path/to/spec.yaml",
+        )
+
+        with patch.object(
+            manager, "_get_mcp_server_from_tool_name", return_value=server
+        ):
+            with patch.object(
+                manager,
+                "pre_call_tool_check",
+                new_callable=AsyncMock,
+                return_value={"extra_headers": {"Authorization": "Bearer jwt"}},
+            ):
+                with patch.object(
+                    manager,
+                    "_create_during_hook_task",
+                    return_value=asyncio.create_task(asyncio.sleep(0)),
+                ):
+                    with patch.object(
+                        manager,
+                        "_call_openapi_tool_handler",
+                        new_callable=AsyncMock,
+                        return_value=MagicMock(),
+                    ):
+                        import litellm.proxy._experimental.mcp_server.mcp_server_manager as mgr_mod
+
+                        proxy_logging = MagicMock(spec=ProxyLogging)
+
+                        with patch.object(mgr_mod, "verbose_logger") as mock_logger:
+                            # Should NOT raise — just warn and proceed
+                            await manager.call_tool(
+                                server_name="openapi_server",
+                                name="test_tool",
+                                arguments={},
+                                proxy_logging_obj=proxy_logging,
+                            )
+                            mock_logger.warning.assert_called_once()
+                            assert "header injection is not supported" in mock_logger.warning.call_args[0][0]
+
+    @pytest.mark.asyncio
+    async def test_openapi_server_no_error_without_hook_headers(self):
+        """No exception when OpenAPI server has no hook-injected headers."""
+        manager = MCPServerManager()
+        server = MCPServer(
+            server_id="test-id",
+            name="openapi_server",
+            server_name="openapi_server",
+            url="https://example.com",
+            transport=MCPTransport.http,
+            auth_type=MCPAuth.none,
+            spec_path="/path/to/spec.yaml",
+        )
+
+        with patch.object(
+            manager, "_get_mcp_server_from_tool_name", return_value=server
+        ):
+            with patch.object(
+                manager,
+                "pre_call_tool_check",
+                new_callable=AsyncMock,
+                return_value={},
+            ):
+                with patch.object(
+                    manager,
+                    "_create_during_hook_task",
+                    return_value=asyncio.create_task(asyncio.sleep(0)),
+                ):
+                    with patch.object(
+                        manager,
+                        "_call_openapi_tool_handler",
+                        new_callable=AsyncMock,
+                        return_value=MagicMock(),
+                    ):
+                        proxy_logging = MagicMock(spec=ProxyLogging)
+
+                        await manager.call_tool(
+                            server_name="openapi_server",
+                            name="test_tool",
+                            arguments={},
+                            proxy_logging_obj=proxy_logging,
+                        )
+
+
+class TestHookHeaderMergePriority:
+    """Tests that hook-provided headers have highest priority in _call_regular_mcp_tool."""
+
+    def _make_server(
+        self,
+        static_headers: Optional[Dict[str, str]] = None,
+        extra_headers_config: Optional[list] = None,
+    ):
+        return MCPServer(
+            server_id="test-id",
+            name="Test Server",
+            server_name="test_server",
+            url="https://example.com",
+            transport=MCPTransport.http,
+            auth_type=MCPAuth.none,
+            static_headers=static_headers,
+            extra_headers=extra_headers_config,
+        )
+
+    @pytest.mark.asyncio
+    async def test_hook_headers_override_static_headers(self):
+        """Hook headers should take precedence over static_headers."""
+        manager = MCPServerManager()
+        server = self._make_server(
+            static_headers={"Authorization": "Bearer static-token", "X-Static": "yes"}
+        )
+
+        hook_headers = {"Authorization": "Bearer hook-signed-jwt"}
+
+        captured_extra_headers: Dict[str, Any] = {}
+
+        async def fake_create_mcp_client(
+            server, mcp_auth_header=None, extra_headers=None, stdio_env=None
+        ):
+            captured_extra_headers["value"] = extra_headers
+            mock_client = MagicMock()
+            mock_client.call_tool = AsyncMock(return_value=MagicMock())
+            return mock_client
+
+        with patch.object(
+            manager, "_create_mcp_client", side_effect=fake_create_mcp_client
+        ):
+            with patch.object(manager, "_build_stdio_env", return_value=None):
+                try:
+                    await manager._call_regular_mcp_tool(
+                        mcp_server=server,
+                        original_tool_name="test_tool",
+                        arguments={"key": "val"},
+                        tasks=[],
+                        mcp_auth_header=None,
+                        mcp_server_auth_headers=None,
+                        oauth2_headers=None,
+                        raw_headers=None,
+                        proxy_logging_obj=None,
+                        hook_extra_headers=hook_headers,
+                    )
+                except Exception:
+                    pass
+
+        headers = captured_extra_headers.get("value", {})
+        assert headers["Authorization"] == "Bearer hook-signed-jwt"
+        assert headers["X-Static"] == "yes"
+
+    @pytest.mark.asyncio
+    async def test_no_hook_headers_preserves_existing_behavior(self):
+        """When hook_extra_headers is None, existing header logic is unchanged."""
+        manager = MCPServerManager()
+        server = self._make_server(
+            static_headers={"X-Static": "static-value"}
+        )
+
+        captured_extra_headers: Dict[str, Any] = {}
+
+        async def fake_create_mcp_client(
+            server, mcp_auth_header=None, extra_headers=None, stdio_env=None
+        ):
+            captured_extra_headers["value"] = extra_headers
+            mock_client = MagicMock()
+            mock_client.call_tool = AsyncMock(return_value=MagicMock())
+            return mock_client
+
+        with patch.object(
+            manager, "_create_mcp_client", side_effect=fake_create_mcp_client
+        ):
+            with patch.object(manager, "_build_stdio_env", return_value=None):
+                try:
+                    await manager._call_regular_mcp_tool(
+                        mcp_server=server,
+                        original_tool_name="test_tool",
+                        arguments={"key": "val"},
+                        tasks=[],
+                        mcp_auth_header=None,
+                        mcp_server_auth_headers=None,
+                        oauth2_headers=None,
+                        raw_headers=None,
+                        proxy_logging_obj=None,
+                        hook_extra_headers=None,
+                    )
+                except Exception:
+                    pass
+
+        headers = captured_extra_headers.get("value", {})
+        assert headers == {"X-Static": "static-value"}
+
+    @pytest.mark.asyncio
+    async def test_hook_headers_merge_with_oauth2(self):
+        """Hook headers merge on top of OAuth2 headers."""
+        manager = MCPServerManager()
+        server = MCPServer(
+            server_id="test-id",
+            name="Test Server",
+            server_name="test_server",
+            url="https://example.com",
+            transport=MCPTransport.http,
+            auth_type=MCPAuth.oauth2,
+        )
+
+        captured_extra_headers: Dict[str, Any] = {}
+
+        async def fake_create_mcp_client(
+            server, mcp_auth_header=None, extra_headers=None, stdio_env=None
+        ):
+            captured_extra_headers["value"] = extra_headers
+            mock_client = MagicMock()
+            mock_client.call_tool = AsyncMock(return_value=MagicMock())
+            return mock_client
+
+        with patch.object(
+            manager, "_create_mcp_client", side_effect=fake_create_mcp_client
+        ):
+            with patch.object(manager, "_build_stdio_env", return_value=None):
+                try:
+                    await manager._call_regular_mcp_tool(
+                        mcp_server=server,
+                        original_tool_name="test_tool",
+                        arguments={"key": "val"},
+                        tasks=[],
+                        mcp_auth_header=None,
+                        mcp_server_auth_headers=None,
+                        oauth2_headers={
+                            "Authorization": "Bearer oauth2-token",
+                            "X-OAuth": "yes",
+                        },
+                        raw_headers=None,
+                        proxy_logging_obj=None,
+                        hook_extra_headers={
+                            "Authorization": "Bearer hook-jwt",
+                            "X-Trace-Id": "trace-123",
+                        },
+                    )
+                except Exception:
+                    pass
+
+        headers = captured_extra_headers.get("value", {})
+        assert headers["Authorization"] == "Bearer hook-jwt"
+        assert headers["X-OAuth"] == "yes"
+        assert headers["X-Trace-Id"] == "trace-123"
+
+
+class TestUserAPIKeyAuthJwtClaims:
+    """Tests that UserAPIKeyAuth correctly carries jwt_claims."""
+
+    def test_jwt_claims_field_defaults_to_none(self):
+        auth = UserAPIKeyAuth(api_key="test-key")
+        assert auth.jwt_claims is None
+
+    def test_jwt_claims_field_accepts_dict(self):
+        claims = {"sub": "user-123", "iss": "litellm", "exp": 9999999999}
+        auth = UserAPIKeyAuth(api_key="test-key", jwt_claims=claims)
+        assert auth.jwt_claims == claims
+        assert auth.jwt_claims["sub"] == "user-123"
+
+    def test_jwt_claims_backward_compatible_without_field(self):
+        """Existing code that doesn't pass jwt_claims should still work."""
+        auth = UserAPIKeyAuth(
+            api_key="test-key",
+            user_id="user-1",
+            team_id="team-1",
+        )
+        assert auth.jwt_claims is None
+        assert auth.user_id == "user-1"
+
+    def test_jwt_claims_set_after_construction(self):
+        """Virtual-key fast path sets jwt_claims after the object is created."""
+        auth = UserAPIKeyAuth(api_key="test-key")
+        assert auth.jwt_claims is None
+
+        claims = {"sub": "user-456", "iss": "okta", "groups": ["admin"]}
+        auth.jwt_claims = claims
+        assert auth.jwt_claims == claims
+        assert auth.jwt_claims["groups"] == ["admin"]

tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py

@@ -2062,7 +2062,14 @@ async def test_list_team_v2_security_check_non_admin_user():
         user_id="non_admin_user_123",
     )
 
-    with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client:
+    with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client, \
+         patch("litellm.proxy.proxy_server.user_api_key_cache"), \
+         patch("litellm.proxy.proxy_server.proxy_logging_obj"), \
+         patch(
+             "litellm.proxy.management_endpoints.team_endpoints.get_user_object",
+             new_callable=AsyncMock,
+             return_value=None,
+         ):
         mock_prisma_client.return_value = MagicMock()  # Mock non-None prisma client
 
         # Should raise HTTPException with 401 status
@@ -2103,7 +2110,14 @@ async def test_list_team_v2_security_check_non_admin_user_other_user():
         user_id="non_admin_user_123",
     )
 
-    with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client:
+    with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client, \
+         patch("litellm.proxy.proxy_server.user_api_key_cache"), \
+         patch("litellm.proxy.proxy_server.proxy_logging_obj"), \
+         patch(
+             "litellm.proxy.management_endpoints.team_endpoints.get_user_object",
+             new_callable=AsyncMock,
+             return_value=None,
+         ):
         mock_prisma_client.return_value = MagicMock()  # Mock non-None prisma client
 
         # Should raise HTTPException with 401 status
@@ -2142,19 +2156,21 @@ async def test_list_team_v2_security_check_non_admin_user_own_teams():
         user_id="non_admin_user_123",
     )
 
-    with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client:
+    with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client, \
+         patch("litellm.proxy.proxy_server.user_api_key_cache"), \
+         patch("litellm.proxy.proxy_server.proxy_logging_obj"):
         # Mock prisma client and database operations
         mock_db = Mock()
         mock_prisma_client.db = mock_db
-        
-        # Mock user lookup
-        mock_user_object = Mock()
-        mock_user_object.model_dump.return_value = {
-            "user_id": "non_admin_user_123",
-            "teams": ["team_1", "team_2"],
-        }
-        mock_db.litellm_usertable.find_unique = AsyncMock(return_value=mock_user_object)
-        
+
+        # Mock get_user_object to return a user with teams
+        from litellm.proxy._types import LiteLLM_UserTable
+
+        mock_user = LiteLLM_UserTable(
+            user_id="non_admin_user_123",
+            teams=["team_1", "team_2"],
+        )
+
         # Mock team lookup
         mock_teams = [
             Mock(model_dump=lambda: {"team_id": "team_1", "team_alias": "Team 1"}),
@@ -2163,21 +2179,26 @@ async def test_list_team_v2_security_check_non_admin_user_own_teams():
         mock_db.litellm_teamtable.find_many = AsyncMock(return_value=mock_teams)
         mock_db.litellm_teamtable.count = AsyncMock(return_value=2)
 
-        # Should NOT raise an exception
-        result = await list_team_v2(
-            http_request=mock_request,
-            user_id="non_admin_user_123",  # Non-admin querying their own teams
-            user_api_key_dict=mock_user_api_key_dict_non_admin,
-            team_id=None,
-            page=1,
-            page_size=10,
-            status=None,
-        )
+        with patch(
+            "litellm.proxy.management_endpoints.team_endpoints.get_user_object",
+            new_callable=AsyncMock,
+            return_value=mock_user,
+        ):
+            # Should NOT raise an exception
+            result = await list_team_v2(
+                http_request=mock_request,
+                user_id="non_admin_user_123",  # Non-admin querying their own teams
+                user_api_key_dict=mock_user_api_key_dict_non_admin,
+                team_id=None,
+                page=1,
+                page_size=10,
+                status=None,
+            )
 
-        # Should return results without error
-        assert "teams" in result
-        assert "total" in result
-        assert result["total"] == 2
+            # Should return results without error
+            assert "teams" in result
+            assert "total" in result
+            assert result["total"] == 2
 
 
 @pytest.mark.asyncio
@@ -2293,27 +2314,280 @@ async def test_list_team_v2_with_status_deleted():
         assert len(result["teams"]) == 2
 
 
+@pytest.mark.asyncio
+async def test_list_team_v2_org_admin_sees_org_teams():
+    """
+    Test that an org admin (internal_user role with org_admin membership)
+    can list teams scoped to their organisations without getting a 401.
+    """
+    from datetime import datetime
+    from unittest.mock import AsyncMock, Mock, patch
+
+    from fastapi import Request
+
+    from litellm.proxy._types import (
+        LiteLLM_OrganizationMembershipTable,
+        LiteLLM_UserTable,
+        LitellmUserRoles,
+        UserAPIKeyAuth,
+    )
+    from litellm.proxy.management_endpoints.team_endpoints import list_team_v2
+
+    mock_request = Mock(spec=Request)
+    mock_user_api_key_dict = UserAPIKeyAuth(
+        user_role=LitellmUserRoles.INTERNAL_USER,
+        user_id="org_admin_user",
+    )
+
+    mock_user = LiteLLM_UserTable(
+        user_id="org_admin_user",
+        teams=[],
+        organization_memberships=[
+            LiteLLM_OrganizationMembershipTable(
+                user_id="org_admin_user",
+                organization_id="org_A",
+                user_role="org_admin",
+                spend=0.0,
+                created_at=datetime.now(),
+                updated_at=datetime.now(),
+            ),
+        ],
+    )
+
+    with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, \
+         patch("litellm.proxy.proxy_server.user_api_key_cache"), \
+         patch("litellm.proxy.proxy_server.proxy_logging_obj"), \
+         patch(
+             "litellm.proxy.management_endpoints.team_endpoints.get_user_object",
+             new_callable=AsyncMock,
+             return_value=mock_user,
+         ):
+        mock_db = Mock()
+        mock_prisma.db = mock_db
+
+        mock_team = Mock()
+        mock_team.model_dump.return_value = {
+            "team_id": "team_in_org_A",
+            "team_alias": "Org A Team",
+            "organization_id": "org_A",
+            "members_with_roles": [{"user_id": "u1", "role": "user"}],
+        }
+        mock_db.litellm_teamtable.find_many = AsyncMock(return_value=[mock_team])
+        mock_db.litellm_teamtable.count = AsyncMock(return_value=1)
+
+        result = await list_team_v2(
+            http_request=mock_request,
+            user_id=None,
+            organization_id=None,
+            team_id=None,
+            team_alias=None,
+            user_api_key_dict=mock_user_api_key_dict,
+            page=1,
+            page_size=10,
+            sort_by=None,
+            sort_order="asc",
+            status=None,
+        )
+
+        assert result["total"] == 1
+        assert len(result["teams"]) == 1
+        assert result["teams"][0].members_count == 1
+
+        # Verify org-scoped where clause
+        where = mock_db.litellm_teamtable.find_many.call_args.kwargs["where"]
+        assert where["organization_id"] == {"in": ["org_A"]}
+
+
+@pytest.mark.asyncio
+async def test_list_team_v2_org_admin_cannot_view_other_orgs():
+    """
+    Test that an org admin is rejected with 403 when filtering by an
+    organisation they do not administer.
+    """
+    from datetime import datetime
+    from unittest.mock import AsyncMock, Mock, patch
+
+    from fastapi import HTTPException, Request
+
+    from litellm.proxy._types import (
+        LiteLLM_OrganizationMembershipTable,
+        LiteLLM_UserTable,
+        LitellmUserRoles,
+        UserAPIKeyAuth,
+    )
+    from litellm.proxy.management_endpoints.team_endpoints import list_team_v2
+
+    mock_request = Mock(spec=Request)
+    mock_user_api_key_dict = UserAPIKeyAuth(
+        user_role=LitellmUserRoles.INTERNAL_USER,
+        user_id="org_admin_user",
+    )
+
+    mock_user = LiteLLM_UserTable(
+        user_id="org_admin_user",
+        teams=[],
+        organization_memberships=[
+            LiteLLM_OrganizationMembershipTable(
+                user_id="org_admin_user",
+                organization_id="org_A",
+                user_role="org_admin",
+                spend=0.0,
+                created_at=datetime.now(),
+                updated_at=datetime.now(),
+            ),
+        ],
+    )
+
+    with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, \
+         patch("litellm.proxy.proxy_server.user_api_key_cache"), \
+         patch("litellm.proxy.proxy_server.proxy_logging_obj"), \
+         patch(
+             "litellm.proxy.management_endpoints.team_endpoints.get_user_object",
+             new_callable=AsyncMock,
+             return_value=mock_user,
+         ):
+        mock_prisma.db = Mock()
+
+        with pytest.raises(HTTPException) as exc_info:
+            await list_team_v2(
+                http_request=mock_request,
+                user_id=None,
+                organization_id="org_B",  # not their org
+                team_id=None,
+                team_alias=None,
+                user_api_key_dict=mock_user_api_key_dict,
+                page=1,
+                page_size=10,
+                sort_by=None,
+                sort_order="asc",
+                status=None,
+            )
+
+        assert exc_info.value.status_code == 403
+        assert "only view teams within your organizations" in str(
+            exc_info.value.detail
+        ).lower()
+
+
+@pytest.mark.asyncio
+async def test_list_team_v2_org_admin_with_user_id_returns_user_teams():
+    """
+    Test that an org admin passing user_id gets that user's direct team
+    memberships (not all org teams).
+    """
+    from datetime import datetime
+    from unittest.mock import AsyncMock, Mock, patch
+
+    from fastapi import Request
+
+    from litellm.proxy._types import (
+        LiteLLM_OrganizationMembershipTable,
+        LiteLLM_UserTable,
+        LitellmUserRoles,
+        UserAPIKeyAuth,
+    )
+    from litellm.proxy.management_endpoints.team_endpoints import list_team_v2
+
+    mock_request = Mock(spec=Request)
+    mock_user_api_key_dict = UserAPIKeyAuth(
+        user_role=LitellmUserRoles.INTERNAL_USER,
+        user_id="org_admin_user",
+    )
+
+    mock_org_admin = LiteLLM_UserTable(
+        user_id="org_admin_user",
+        teams=["team_1"],
+        organization_memberships=[
+            LiteLLM_OrganizationMembershipTable(
+                user_id="org_admin_user",
+                organization_id="org_A",
+                user_role="org_admin",
+                spend=0.0,
+                created_at=datetime.now(),
+                updated_at=datetime.now(),
+            ),
+        ],
+    )
+
+    # The target user whose teams we want to list
+    mock_target_user = LiteLLM_UserTable(
+        user_id="target_user",
+        teams=["team_X", "team_Y"],
+    )
+
+    call_count = 0
+
+    async def mock_get_user_object(**kwargs):
+        nonlocal call_count
+        call_count += 1
+        # First call: org admin lookup in list_team_v2
+        # Second call: target user lookup in _build_team_list_where_conditions
+        if call_count == 1:
+            return mock_org_admin
+        return mock_target_user
+
+    with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, \
+         patch("litellm.proxy.proxy_server.user_api_key_cache"), \
+         patch("litellm.proxy.proxy_server.proxy_logging_obj"), \
+         patch(
+             "litellm.proxy.management_endpoints.team_endpoints.get_user_object",
+             side_effect=mock_get_user_object,
+         ):
+        mock_db = Mock()
+        mock_prisma.db = mock_db
+
+        mock_team = Mock()
+        mock_team.model_dump.return_value = {
+            "team_id": "team_X",
+            "team_alias": "Target Team",
+            "members_with_roles": [{"user_id": "target_user", "role": "user"}],
+        }
+        mock_db.litellm_teamtable.find_many = AsyncMock(return_value=[mock_team])
+        mock_db.litellm_teamtable.count = AsyncMock(return_value=1)
+
+        result = await list_team_v2(
+            http_request=mock_request,
+            user_id="target_user",
+            organization_id=None,
+            team_id=None,
+            team_alias=None,
+            user_api_key_dict=mock_user_api_key_dict,
+            page=1,
+            page_size=10,
+            sort_by=None,
+            sort_order="asc",
+            status=None,
+        )
+
+        assert result["total"] == 1
+
+        # Verify the where clause filters by user's teams, not org scope
+        where = mock_db.litellm_teamtable.find_many.call_args.kwargs["where"]
+        assert where["team_id"] == {"in": ["team_X", "team_Y"]}
+        assert "organization_id" not in where
+
+
 @pytest.mark.asyncio
 async def test_list_team_v2_with_invalid_status():
     """
     Test that invalid status parameter raises HTTPException.
     """
     from unittest.mock import Mock, patch
-    
+
     from fastapi import HTTPException, Request
-    
+
     from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
     from litellm.proxy.management_endpoints.team_endpoints import list_team_v2
-    
+
     # Mock request
     mock_request = Mock(spec=Request)
-    
+
     # Mock admin user
     mock_user_api_key_dict_admin = UserAPIKeyAuth(
         user_role=LitellmUserRoles.PROXY_ADMIN,
         user_id="admin_user_123",
     )
-    
+
     mock_prisma_client = Mock()
     
     # Mock prisma_client to be non-None

tests/test_litellm/test_setup_wizard.py

@@ -0,0 +1,188 @@
+"""Unit tests for litellm.setup_wizard — pure functions only, no network calls."""
+
+from litellm.setup_wizard import SetupWizard, _yaml_escape
+
+# ---------------------------------------------------------------------------
+# _yaml_escape
+# ---------------------------------------------------------------------------
+
+
+def test_yaml_escape_plain():
+    assert _yaml_escape("sk-abc123") == "sk-abc123"
+
+
+def test_yaml_escape_double_quote():
+    assert _yaml_escape('sk-ab"cd') == 'sk-ab\\"cd'
+
+
+def test_yaml_escape_backslash():
+    assert _yaml_escape("sk-ab\\cd") == "sk-ab\\\\cd"
+
+
+def test_yaml_escape_combined():
+    assert _yaml_escape('ab\\"cd') == 'ab\\\\\\"cd'
+
+
+def test_yaml_escape_newline():
+    assert _yaml_escape("sk-abc\ndef") == "sk-abc\\ndef"
+
+
+def test_yaml_escape_carriage_return():
+    assert _yaml_escape("sk-abc\rdef") == "sk-abc\\rdef"
+
+
+def test_yaml_escape_tab():
+    assert _yaml_escape("sk-abc\tdef") == "sk-abc\\tdef"
+
+
+# ---------------------------------------------------------------------------
+# SetupWizard._build_config
+# ---------------------------------------------------------------------------
+
+_OPENAI = {
+    "id": "openai",
+    "name": "OpenAI",
+    "env_key": "OPENAI_API_KEY",
+    "models": ["gpt-4o", "gpt-4o-mini"],
+    "test_model": "gpt-4o-mini",
+}
+
+_ANTHROPIC = {
+    "id": "anthropic",
+    "name": "Anthropic",
+    "env_key": "ANTHROPIC_API_KEY",
+    "models": ["claude-opus-4-6"],
+    "test_model": "claude-haiku-4-5-20251001",
+}
+
+_AZURE = {
+    "id": "azure",
+    "name": "Azure OpenAI",
+    "env_key": "AZURE_API_KEY",
+    "models": [],
+    "test_model": None,
+    "needs_api_base": True,
+    "api_base_hint": "https://<resource>.openai.azure.com/",
+    "api_version": "2024-07-01-preview",
+}
+
+_OLLAMA = {
+    "id": "ollama",
+    "name": "Ollama",
+    "env_key": None,
+    "models": ["ollama/llama3.2"],
+    "test_model": None,
+    "api_base": "http://localhost:11434",
+}
+
+
+def test_build_config_basic_openai():
+    config = SetupWizard._build_config(
+        [_OPENAI],
+        {"OPENAI_API_KEY": "sk-test"},
+        "sk-master",
+    )
+    assert "model_list:" in config
+    assert "model_name: gpt-4o" in config
+    assert "model: gpt-4o" in config
+    assert "api_key: os.environ/OPENAI_API_KEY" in config
+    assert 'master_key: "sk-master"' in config
+
+
+def test_build_config_skipped_provider_omitted():
+    """Provider with no key in env_vars should not appear in model_list."""
+    config = SetupWizard._build_config(
+        [_OPENAI, _ANTHROPIC],
+        {"ANTHROPIC_API_KEY": "sk-ant-test"},  # OpenAI key missing
+        "sk-master",
+    )
+    assert "gpt-4o" not in config
+    assert "claude-opus-4-6" in config
+
+
+def test_build_config_env_vars_written_escaped():
+    """API keys with special chars should be YAML-escaped."""
+    config = SetupWizard._build_config(
+        [_OPENAI],
+        {"OPENAI_API_KEY": 'sk-ab"cd'},
+        "sk-master",
+    )
+    assert 'OPENAI_API_KEY: "sk-ab\\"cd"' in config
+
+
+def test_build_config_master_key_quoted():
+    """master_key must be quoted in YAML to handle special characters."""
+    config = SetupWizard._build_config(
+        [_OPENAI],
+        {"OPENAI_API_KEY": "sk-test"},
+        'sk-master"special',
+    )
+    assert 'master_key: "sk-master\\"special"' in config
+
+
+def test_build_config_does_not_mutate_env_vars():
+    """_build_config must not modify the caller's env_vars dict."""
+    env_vars = {
+        "AZURE_API_KEY": "az-key",
+        "_LITELLM_AZURE_API_BASE_AZURE": "https://my.azure.com",
+        "_LITELLM_AZURE_DEPLOYMENT_AZURE": "my-deployment",
+    }
+    original_keys = set(env_vars.keys())
+    SetupWizard._build_config([_AZURE], env_vars, "sk-master")
+    assert set(env_vars.keys()) == original_keys
+
+
+def test_build_config_azure_uses_deployment_name():
+    env_vars = {
+        "AZURE_API_KEY": "az-key",
+        "_LITELLM_AZURE_API_BASE_AZURE": "https://my.azure.com",
+        "_LITELLM_AZURE_DEPLOYMENT_AZURE": "my-gpt4o",
+    }
+    config = SetupWizard._build_config([_AZURE], env_vars, "sk-master")
+    assert "model: azure/my-gpt4o" in config
+    assert "model_name: azure-my-gpt4o" in config
+    # api_base must be quoted to survive YAML special chars
+    assert 'api_base: "https://my.azure.com"' in config
+
+
+def test_build_config_azure_no_deployment_skipped():
+    """Azure without a deployment name should emit nothing (not fallback to gpt-4o)."""
+    env_vars = {"AZURE_API_KEY": "az-key"}  # no deployment sentinel
+    config = SetupWizard._build_config([_AZURE], env_vars, "sk-master")
+    # No azure model entry should be emitted when deployment name is absent
+    assert "model: azure/" not in config
+
+
+def test_build_config_no_display_name_collision_openai_and_azure():
+    """OpenAI gpt-4o and azure gpt-4o should get distinct model_name values."""
+    env_vars = {
+        "OPENAI_API_KEY": "sk-openai",
+        "AZURE_API_KEY": "az-key",
+        "_LITELLM_AZURE_DEPLOYMENT_AZURE": "gpt-4o",
+    }
+    config = SetupWizard._build_config([_OPENAI, _AZURE], env_vars, "sk-master")
+    assert "model_name: gpt-4o" in config  # OpenAI
+    assert "model_name: azure-gpt-4o" in config  # Azure — qualified
+
+
+def test_build_config_ollama_no_api_key_line():
+    """Ollama has no env_key — config should not contain an api_key line for it."""
+    config = SetupWizard._build_config([_OLLAMA], {}, "sk-master")
+    assert "ollama/llama3.2" in config
+    assert "api_key:" not in config
+
+
+def test_build_config_master_key_in_general_settings():
+    """master_key is written to general_settings."""
+    config = SetupWizard._build_config([_OPENAI], {"OPENAI_API_KEY": "k"}, "sk-m")
+    assert 'master_key: "sk-m"' in config
+
+
+def test_build_config_internal_sentinel_keys_excluded():
+    """_LITELLM_ prefixed sentinel keys must not appear in environment_variables."""
+    env_vars = {
+        "OPENAI_API_KEY": "sk-real",
+        "_LITELLM_AZURE_API_BASE_AZURE": "https://x.azure.com",
+    }
+    config = SetupWizard._build_config([_OPENAI], env_vars, "sk-master")
+    assert "_LITELLM_" not in config

ui/litellm-dashboard/src/app/(dashboard)/components/Sidebar2.tsx

@@ -20,7 +20,6 @@ import {
   ToolOutlined,
   TagsOutlined,
   AuditOutlined,
-  MessageOutlined,
 } from "@ant-design/icons";
 // import {
 //   all_admin_roles,
@@ -466,41 +465,6 @@ const Sidebar2: React.FC<SidebarProps> = ({ accessToken, userRole, defaultSelect
         </ConfigProvider>
         {isAdminRole(userRole) && !collapsed && <UsageIndicator accessToken={accessToken} width={220} />}
 
-        {/* Pinned "Open Chat" button at bottom */}
-        <div style={{
-          padding: collapsed ? "10px 8px" : "10px 12px",
-          borderTop: "1px solid #f0f0f0",
-          flexShrink: 0,
-        }}>
-          <a
-            href={toHref("chat")}
-            target="_blank"
-            rel="noopener noreferrer"
-            style={{
-              display: "flex",
-              alignItems: "center",
-              justifyContent: collapsed ? "center" : "flex-start",
-              gap: 8,
-              padding: collapsed ? "8px 0" : "8px 10px",
-              borderRadius: 8,
-              background: "#1677ff",
-              color: "#fff",
-              textDecoration: "none",
-              fontSize: 13,
-              fontWeight: 600,
-              transition: "background 0.15s",
-            }}
-            onMouseEnter={(e) => {
-              (e.currentTarget as HTMLAnchorElement).style.background = "#0958d9";
-            }}
-            onMouseLeave={(e) => {
-              (e.currentTarget as HTMLAnchorElement).style.background = "#1677ff";
-            }}
-          >
-            <MessageOutlined style={{ fontSize: 16, flexShrink: 0 }} />
-            {!collapsed && <span>Open Chat</span>}
-          </a>
-        </div>
       </Sider>
     </Layout>
   );

ui/litellm-dashboard/src/app/(dashboard)/playground/page.tsx

@@ -8,8 +8,6 @@ import ComplianceUI from "@/components/playground/complianceUI/ComplianceUI";
 import { TabGroup, TabList, Tab, TabPanels, TabPanel } from "@tremor/react";
 import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
 import { fetchProxySettings } from "@/utils/proxyUtils";
-import { useUIConfig } from "@/app/(dashboard)/hooks/uiConfig/useUIConfig";
-import { MessageOutlined, CloseOutlined } from "@ant-design/icons";
 
 interface ProxySettings {
   PROXY_BASE_URL?: string;
@@ -19,12 +17,6 @@ interface ProxySettings {
 export default function PlaygroundPage() {
   const { accessToken, userRole, userId, disabledPersonalKeyCreation, token } = useAuthorized();
   const [proxySettings, setProxySettings] = useState<ProxySettings | undefined>(undefined);
-  const [chatBannerDismissed, setChatBannerDismissed] = useState(false);
-  const { data: uiConfig } = useUIConfig();
-  const uiRoot = uiConfig?.server_root_path && uiConfig.server_root_path !== "/"
-    ? uiConfig.server_root_path.replace(/\/+$/, "")
-    : "";
-  const chatHref = `${uiRoot}/ui/chat`;
 
   useEffect(() => {
     const initializeProxySettings = async () => {
@@ -44,64 +36,6 @@ export default function PlaygroundPage() {
 
   return (
     <div className="h-full w-full flex flex-col">
-      {!chatBannerDismissed && (
-        <div style={{
-          display: "flex",
-          alignItems: "center",
-          gap: 16,
-          padding: "10px 20px",
-          background: "#f0f9ff",
-          borderBottom: "1px solid #bae6fd",
-          flexShrink: 0,
-        }}>
-          <span style={{
-            fontSize: 10,
-            fontWeight: 700,
-            color: "#fff",
-            background: "#0ea5e9",
-            borderRadius: 4,
-            padding: "2px 7px",
-            letterSpacing: "0.08em",
-            textTransform: "uppercase",
-            flexShrink: 0,
-            lineHeight: "18px",
-          }}>
-            New
-          </span>
-          <span style={{ flex: 1, color: "#0c4a6e", fontSize: 13.5, lineHeight: 1.5 }}>
-            <strong>Chat UI</strong>
-            {" "}— a ChatGPT-like interface for your users to chat with AI models and MCP tools. Share it with your team.
-          </span>
-          <a
-            href={chatHref}
-            target="_blank"
-            rel="noopener noreferrer"
-            style={{
-              display: "inline-flex",
-              alignItems: "center",
-              gap: 5,
-              padding: "5px 14px",
-              borderRadius: 6,
-              background: "#0ea5e9",
-              color: "#fff",
-              fontSize: 12.5,
-              fontWeight: 600,
-              textDecoration: "none",
-              whiteSpace: "nowrap",
-              flexShrink: 0,
-            }}
-          >
-            Open Chat UI →
-          </a>
-          <button
-            onClick={() => setChatBannerDismissed(true)}
-            style={{ background: "none", border: "none", cursor: "pointer", color: "#64748b", padding: 4, flexShrink: 0, lineHeight: 1 }}
-            aria-label="Dismiss"
-          >
-            <CloseOutlined style={{ fontSize: 13 }} />
-          </button>
-        </div>
-      )}
     <TabGroup className="w-full" style={{ flex: 1, minHeight: 0, display: "flex", flexDirection: "column" }}>
       <TabList className="mb-0">
         <Tab>Chat</Tab>

ui/litellm-dashboard/src/components/DebugWarningBanner.test.tsx

@@ -2,37 +2,40 @@ import { renderWithProviders, screen } from "../../tests/test-utils";
 import { vi } from "vitest";
 import { DebugWarningBanner } from "./DebugWarningBanner";
 
-const mockUseHealthReadiness = vi.fn();
 vi.mock("@/app/(dashboard)/hooks/healthReadiness/useHealthReadiness", () => ({
-  useHealthReadiness: () => mockUseHealthReadiness(),
+  useHealthReadiness: vi.fn(),
 }));
 
+import { useHealthReadiness } from "@/app/(dashboard)/hooks/healthReadiness/useHealthReadiness";
+
 describe("DebugWarningBanner", () => {
-  afterEach(() => {
-    vi.resetAllMocks();
-  });
-
   it("should render", () => {
-    mockUseHealthReadiness.mockReturnValue({ data: { is_detailed_debug: true } });
+    vi.mocked(useHealthReadiness).mockReturnValue({ data: { is_detailed_debug: true } } as any);
     renderWithProviders(<DebugWarningBanner />);
-    expect(screen.getByText(/Performance Warning/i)).toBeInTheDocument();
+    expect(screen.getByRole("alert")).toBeInTheDocument();
   });
 
-  it("should render nothing when debug mode is disabled", () => {
-    mockUseHealthReadiness.mockReturnValue({ data: { is_detailed_debug: false } });
-    const { container } = renderWithProviders(<DebugWarningBanner />);
-    expect(container.firstChild).toBeNull();
-  });
-
-  it("should render nothing when health data is undefined", () => {
-    mockUseHealthReadiness.mockReturnValue({ data: undefined });
-    const { container } = renderWithProviders(<DebugWarningBanner />);
-    expect(container.firstChild).toBeNull();
+  it("should show warning when detailed debug mode is active", () => {
+    vi.mocked(useHealthReadiness).mockReturnValue({ data: { is_detailed_debug: true } } as any);
+    renderWithProviders(<DebugWarningBanner />);
+    expect(screen.getByText(/Performance Warning: Detailed Debug Mode Active/i)).toBeInTheDocument();
   });
 
   it("should mention LITELLM_LOG=DEBUG in the description", () => {
-    mockUseHealthReadiness.mockReturnValue({ data: { is_detailed_debug: true } });
+    vi.mocked(useHealthReadiness).mockReturnValue({ data: { is_detailed_debug: true } } as any);
     renderWithProviders(<DebugWarningBanner />);
-    expect(screen.getByText(/LITELLM_LOG=DEBUG/)).toBeInTheDocument();
+    expect(screen.getByText("LITELLM_LOG=DEBUG")).toBeInTheDocument();
+  });
+
+  it("should render nothing when is_detailed_debug is false", () => {
+    vi.mocked(useHealthReadiness).mockReturnValue({ data: { is_detailed_debug: false } } as any);
+    const { container } = renderWithProviders(<DebugWarningBanner />);
+    expect(container).toBeEmptyDOMElement();
+  });
+
+  it("should render nothing when health data is undefined", () => {
+    vi.mocked(useHealthReadiness).mockReturnValue({ data: undefined } as any);
+    const { container } = renderWithProviders(<DebugWarningBanner />);
+    expect(container).toBeEmptyDOMElement();
   });
 });

ui/litellm-dashboard/src/components/EntityUsageExport/ExportFormatSelector.test.tsx

@@ -1,34 +1,20 @@
-import { render, screen } from "@testing-library/react";
-import userEvent from "@testing-library/user-event";
+import { renderWithProviders, screen } from "../../../tests/test-utils";
 import { vi } from "vitest";
 import ExportFormatSelector from "./ExportFormatSelector";
 
 describe("ExportFormatSelector", () => {
   it("should render", () => {
-    render(<ExportFormatSelector value="csv" onChange={vi.fn()} />);
+    renderWithProviders(<ExportFormatSelector value="csv" onChange={vi.fn()} />);
     expect(screen.getByText("Format")).toBeInTheDocument();
   });
 
-  it("should display the current value as csv", () => {
-    render(<ExportFormatSelector value="csv" onChange={vi.fn()} />);
+  it("should display the current value", () => {
+    renderWithProviders(<ExportFormatSelector value="csv" onChange={vi.fn()} />);
     expect(screen.getByText("CSV (Excel, Google Sheets)")).toBeInTheDocument();
   });
 
-  it("should display the current value as json", () => {
-    render(<ExportFormatSelector value="json" onChange={vi.fn()} />);
+  it("should display JSON label when json is selected", () => {
+    renderWithProviders(<ExportFormatSelector value="json" onChange={vi.fn()} />);
     expect(screen.getByText("JSON (includes metadata)")).toBeInTheDocument();
   });
-
-  it("should call onChange when a different format is selected", async () => {
-    const onChange = vi.fn();
-    const user = userEvent.setup();
-    render(<ExportFormatSelector value="csv" onChange={onChange} />);
-
-    // Open the Ant Design Select dropdown
-    await user.click(screen.getByText("CSV (Excel, Google Sheets)"));
-    // Select JSON option from the dropdown
-    const jsonOption = await screen.findByText("JSON (includes metadata)");
-    await user.click(jsonOption);
-    expect(onChange).toHaveBeenCalledWith("json", expect.anything());
-  });
 });

ui/litellm-dashboard/src/components/EntityUsageExport/ExportSummary.test.tsx

@@ -1,37 +1,59 @@
-import { render, screen } from "@testing-library/react";
+import { renderWithProviders, screen } from "../../../tests/test-utils";
 import ExportSummary from "./ExportSummary";
 
 describe("ExportSummary", () => {
-  const dateRange = {
-    from: new Date("2025-01-01"),
-    to: new Date("2025-01-31"),
-  };
-
   it("should render", () => {
-    render(<ExportSummary dateRange={dateRange} selectedFilters={[]} />);
-    expect(screen.getByText(/2025/)).toBeInTheDocument();
+    const dateRange = {
+      from: new Date("2024-01-01"),
+      to: new Date("2024-01-31"),
+    };
+    const { container } = renderWithProviders(
+      <ExportSummary dateRange={dateRange} selectedFilters={[]} />
+    );
+    expect(container).not.toBeEmptyDOMElement();
   });
 
-  it("should display formatted date range", () => {
-    render(<ExportSummary dateRange={dateRange} selectedFilters={[]} />);
-    // Pin locale to en-US so test is deterministic regardless of CI runner locale
-    const expectedFrom = dateRange.from!.toLocaleDateString("en-US");
-    const expectedTo = dateRange.to!.toLocaleDateString("en-US");
-    expect(screen.getByText(`${expectedFrom} - ${expectedTo}`)).toBeInTheDocument();
+  it("should display the date range", () => {
+    const from = new Date(2024, 0, 1);
+    const to = new Date(2024, 0, 31);
+    const dateRange = { from, to };
+    renderWithProviders(
+      <ExportSummary dateRange={dateRange} selectedFilters={[]} />
+    );
+    expect(screen.getByText(new RegExp(from.toLocaleDateString()))).toBeInTheDocument();
+    expect(screen.getByText(new RegExp(to.toLocaleDateString()))).toBeInTheDocument();
+  });
+
+  it("should show filter count when filters are selected", () => {
+    const dateRange = {
+      from: new Date("2024-01-01"),
+      to: new Date("2024-01-31"),
+    };
+    renderWithProviders(
+      <ExportSummary dateRange={dateRange} selectedFilters={["team-a", "team-b", "team-c"]} />
+    );
+    expect(screen.getByText(/3 filters/)).toBeInTheDocument();
   });
 
   it("should show singular 'filter' for one filter", () => {
-    render(<ExportSummary dateRange={dateRange} selectedFilters={["team-a"]} />);
+    const dateRange = {
+      from: new Date("2024-01-01"),
+      to: new Date("2024-01-31"),
+    };
+    renderWithProviders(
+      <ExportSummary dateRange={dateRange} selectedFilters={["team-a"]} />
+    );
     expect(screen.getByText(/1 filter$/)).toBeInTheDocument();
   });
 
-  it("should show plural 'filters' for multiple filters", () => {
-    render(<ExportSummary dateRange={dateRange} selectedFilters={["team-a", "team-b"]} />);
-    expect(screen.getByText(/2 filters/)).toBeInTheDocument();
-  });
-
-  it("should not show filter text when no filters applied", () => {
-    render(<ExportSummary dateRange={dateRange} selectedFilters={[]} />);
+  it("should not show filter count when no filters selected", () => {
+    const dateRange = {
+      from: new Date("2024-01-01"),
+      to: new Date("2024-01-31"),
+    };
+    renderWithProviders(
+      <ExportSummary dateRange={dateRange} selectedFilters={[]} />
+    );
     expect(screen.queryByText(/filter/)).not.toBeInTheDocument();
   });
 });

ui/litellm-dashboard/src/components/EntityUsageExport/ExportTypeSelector.test.tsx

@@ -1,37 +1,46 @@
-import { render, screen } from "@testing-library/react";
+import { renderWithProviders, screen } from "../../../tests/test-utils";
 import userEvent from "@testing-library/user-event";
 import { vi } from "vitest";
 import ExportTypeSelector from "./ExportTypeSelector";
 
 describe("ExportTypeSelector", () => {
   it("should render", () => {
-    render(<ExportTypeSelector value="daily" onChange={vi.fn()} entityType="team" />);
+    renderWithProviders(
+      <ExportTypeSelector value="daily" onChange={vi.fn()} entityType="team" />
+    );
     expect(screen.getByText("Export type")).toBeInTheDocument();
   });
 
-  it("should render all three radio options", () => {
-    render(<ExportTypeSelector value="daily" onChange={vi.fn()} entityType="team" />);
-    expect(screen.getAllByRole("radio")).toHaveLength(3);
+  it("should display entity type in radio labels", () => {
+    renderWithProviders(
+      <ExportTypeSelector value="daily" onChange={vi.fn()} entityType="team" />
+    );
+    expect(screen.getByText(/Day-by-day breakdown by team$/)).toBeInTheDocument();
+    expect(screen.getByText(/Day-by-day breakdown by team and key/)).toBeInTheDocument();
+    expect(screen.getByText(/Day-by-day by team and model/)).toBeInTheDocument();
   });
 
-  it("should interpolate entity type in labels", () => {
-    render(<ExportTypeSelector value="daily" onChange={vi.fn()} entityType="organization" />);
+  it("should display the correct entity type for different entities", () => {
+    renderWithProviders(
+      <ExportTypeSelector value="daily" onChange={vi.fn()} entityType="organization" />
+    );
     expect(screen.getByText(/Day-by-day breakdown by organization$/)).toBeInTheDocument();
-    expect(screen.getByText(/organization and key/)).toBeInTheDocument();
-    expect(screen.getByText(/organization and model/)).toBeInTheDocument();
   });
 
-  it("should call onChange when a different option is selected", async () => {
-    const onChange = vi.fn();
+  it("should call onChange when a radio option is selected", async () => {
     const user = userEvent.setup();
-    render(<ExportTypeSelector value="daily" onChange={onChange} entityType="team" />);
-    await user.click(screen.getByText(/by team and key/));
+    const onChange = vi.fn();
+    renderWithProviders(
+      <ExportTypeSelector value="daily" onChange={onChange} entityType="team" />
+    );
+    await user.click(screen.getByRole("radio", { name: /Day-by-day breakdown by team and key/i }));
     expect(onChange).toHaveBeenCalledWith("daily_with_keys");
   });
 
-  it("should have the correct radio checked based on value prop", () => {
-    render(<ExportTypeSelector value="daily_with_models" onChange={vi.fn()} entityType="team" />);
-    const modelRadio = screen.getByRole("radio", { name: /by team and model/i });
-    expect(modelRadio).toBeChecked();
+  it("should have the correct radio checked", () => {
+    renderWithProviders(
+      <ExportTypeSelector value="daily_with_models" onChange={vi.fn()} entityType="team" />
+    );
+    expect(screen.getByRole("radio", { name: /Day-by-day by team and model/i })).toBeChecked();
   });
 });

ui/litellm-dashboard/src/components/GuardrailsMonitor/MetricCard.test.tsx

@@ -1,34 +1,49 @@
-import { render, screen } from "@testing-library/react";
+import { renderWithProviders, screen } from "../../../tests/test-utils";
+import React from "react";
 import { MetricCard } from "./MetricCard";
 
 describe("MetricCard", () => {
   it("should render", () => {
-    render(<MetricCard label="Total Requests" value={1234} />);
+    renderWithProviders(<MetricCard label="Total Requests" value={1234} />);
     expect(screen.getByText("Total Requests")).toBeInTheDocument();
   });
 
-  it("should display the numeric value", () => {
-    render(<MetricCard label="Total Requests" value={1234} />);
-    expect(screen.getByText("1234")).toBeInTheDocument();
+  it("should display the label and value", () => {
+    renderWithProviders(<MetricCard label="Success Rate" value="98.5%" />);
+    expect(screen.getByText("Success Rate")).toBeInTheDocument();
+    expect(screen.getByText("98.5%")).toBeInTheDocument();
   });
 
-  it("should display a string value", () => {
-    render(<MetricCard label="Pass Rate" value="95.2%" />);
-    expect(screen.getByText("95.2%")).toBeInTheDocument();
+  it("should display numeric values", () => {
+    renderWithProviders(<MetricCard label="Count" value={42} />);
+    expect(screen.getByText("42")).toBeInTheDocument();
   });
 
-  it("should display subtitle when provided", () => {
-    render(<MetricCard label="Blocked" value={42} subtitle="Last 7 days" />);
-    expect(screen.getByText("Last 7 days")).toBeInTheDocument();
+  it("should render icon when provided", () => {
+    renderWithProviders(
+      <MetricCard
+        label="Metric"
+        value={100}
+        icon={<span data-testid="test-icon">icon</span>}
+      />
+    );
+    expect(screen.getByTestId("test-icon")).toBeInTheDocument();
   });
 
-  it("should not display subtitle when not provided", () => {
-    render(<MetricCard label="Blocked" value={42} />);
-    expect(screen.queryByText(/days/)).not.toBeInTheDocument();
+  it("should not render icon container when no icon provided", () => {
+    renderWithProviders(<MetricCard label="Metric" value={100} />);
+    expect(screen.queryByTestId("test-icon")).not.toBeInTheDocument();
   });
 
-  it("should display icon when provided", () => {
-    render(<MetricCard label="Status" value="OK" icon={<span data-testid="icon">!</span>} />);
-    expect(screen.getByTestId("icon")).toBeInTheDocument();
+  it("should render subtitle when provided", () => {
+    renderWithProviders(
+      <MetricCard label="Metric" value={100} subtitle="Last 24 hours" />
+    );
+    expect(screen.getByText("Last 24 hours")).toBeInTheDocument();
+  });
+
+  it("should not render subtitle when not provided", () => {
+    renderWithProviders(<MetricCard label="Metric" value={100} />);
+    expect(screen.queryByText("Last 24 hours")).not.toBeInTheDocument();
   });
 });

ui/litellm-dashboard/src/components/HelpLink.test.tsx

@@ -23,6 +23,11 @@ describe("HelpLink", () => {
     expect(screen.getByText("Custom docs link")).toBeInTheDocument();
   });
 
+  it("should have the correct href", () => {
+    renderWithProviders(<HelpLink href="https://docs.example.com/test" />);
+    expect(screen.getByRole("link")).toHaveAttribute("href", "https://docs.example.com/test");
+  });
+
   it("should include a screen-reader-only label for accessibility", () => {
     renderWithProviders(<HelpLink href="https://docs.example.com" />);
 
@@ -46,7 +51,21 @@ describe("HelpIcon", () => {
     expect(screen.getByText("Tooltip help text")).toBeInTheDocument();
   });
 
+  it("should hide tooltip content when not hovered", () => {
+    renderWithProviders(<HelpIcon content="Hidden tooltip" />);
+    expect(screen.queryByText("Hidden tooltip")).not.toBeInTheDocument();
+  });
+
   it("should show learn more link when learnMoreHref is provided", async () => {
+    const user = userEvent.setup();
+    renderWithProviders(
+      <HelpIcon content="Help text" learnMoreHref="https://docs.example.com" />
+    );
+    await user.hover(screen.getByRole("button", { name: /help information/i }));
+    expect(screen.getByText("Learn more")).toBeInTheDocument();
+  });
+
+  it("should use custom learn more text when provided", async () => {
     const user = userEvent.setup();
     renderWithProviders(
       <HelpIcon
@@ -84,6 +103,11 @@ describe("DocsMenu", () => {
     expect(screen.getByRole("button", { name: /docs/i })).toBeInTheDocument();
   });
 
+  it("should hide menu items initially", () => {
+    renderWithProviders(<DocsMenu items={items} />);
+    expect(screen.queryByText("Custom pricing")).not.toBeInTheDocument();
+  });
+
   it("should show menu items when button is clicked", async () => {
     const user = userEvent.setup();
     renderWithProviders(<DocsMenu items={items} />);

ui/litellm-dashboard/src/components/ToolPolicies/PolicySelect.test.tsx

@@ -0,0 +1,84 @@
+import { renderWithProviders, screen } from "../../../tests/test-utils";
+import { vi } from "vitest";
+import {
+  PolicySelect,
+  policyStyle,
+  INPUT_POLICY_OPTIONS,
+  OUTPUT_POLICY_OPTIONS,
+} from "./PolicySelect";
+
+describe("policyStyle", () => {
+  it("should return the matching option for a known policy", () => {
+    expect(policyStyle("trusted")).toEqual(INPUT_POLICY_OPTIONS[1]);
+  });
+
+  it("should return the matching option for blocked", () => {
+    expect(policyStyle("blocked")).toEqual(INPUT_POLICY_OPTIONS[2]);
+  });
+
+  it("should return the first option as fallback for unknown policy", () => {
+    expect(policyStyle("unknown")).toEqual(INPUT_POLICY_OPTIONS[0]);
+  });
+});
+
+describe("PolicySelect", () => {
+  it("should render", () => {
+    renderWithProviders(
+      <PolicySelect
+        value="untrusted"
+        toolName="test-tool"
+        saving={false}
+        onChange={vi.fn()}
+      />
+    );
+    expect(screen.getByText("untrusted")).toBeInTheDocument();
+  });
+
+  it("should show the current policy value", () => {
+    renderWithProviders(
+      <PolicySelect
+        value="trusted"
+        toolName="test-tool"
+        saving={false}
+        onChange={vi.fn()}
+      />
+    );
+    expect(screen.getByText("trusted")).toBeInTheDocument();
+  });
+
+  it("should be disabled when saving is true", () => {
+    renderWithProviders(
+      <PolicySelect
+        value="untrusted"
+        toolName="test-tool"
+        saving={true}
+        onChange={vi.fn()}
+      />
+    );
+    expect(screen.getByRole("combobox")).toHaveAttribute("aria-expanded", "false");
+    expect(screen.getByRole("combobox").closest(".ant-select")).toHaveClass("ant-select-disabled");
+  });
+
+  it("should not be disabled when saving is false", () => {
+    renderWithProviders(
+      <PolicySelect
+        value="untrusted"
+        toolName="test-tool"
+        saving={false}
+        onChange={vi.fn()}
+      />
+    );
+    expect(screen.getByRole("combobox").closest(".ant-select")).not.toHaveClass("ant-select-disabled");
+  });
+});
+
+describe("Policy option constants", () => {
+  it("should have 3 input policy options", () => {
+    expect(INPUT_POLICY_OPTIONS).toHaveLength(3);
+  });
+
+  it("should have 2 output policy options (no blocked)", () => {
+    expect(OUTPUT_POLICY_OPTIONS).toHaveLength(2);
+    expect(OUTPUT_POLICY_OPTIONS.map((o) => o.value)).toEqual(["untrusted", "trusted"]);
+  });
+});

ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx

@@ -0,0 +1,82 @@
+import { renderWithProviders, screen } from "../../../tests/test-utils";
+import { vi } from "vitest";
+import ComplexityRouterConfig from "./ComplexityRouterConfig";
+
+const mockModelInfo = [
+  { model_group: "gpt-4" },
+  { model_group: "gpt-3.5-turbo" },
+  { model_group: "claude-3-opus" },
+] as any[];
+
+const defaultTiers = {
+  SIMPLE: "gpt-3.5-turbo",
+  MEDIUM: "gpt-3.5-turbo",
+  COMPLEX: "gpt-4",
+  REASONING: "claude-3-opus",
+};
+
+describe("ComplexityRouterConfig", () => {
+  it("should render", () => {
+    renderWithProviders(
+      <ComplexityRouterConfig
+        modelInfo={mockModelInfo}
+        value={defaultTiers}
+        onChange={vi.fn()}
+      />
+    );
+    expect(screen.getByText("Complexity Tier Configuration")).toBeInTheDocument();
+  });
+
+  it("should display all four tier labels", () => {
+    renderWithProviders(
+      <ComplexityRouterConfig
+        modelInfo={mockModelInfo}
+        value={defaultTiers}
+        onChange={vi.fn()}
+      />
+    );
+    expect(screen.getByText("Simple Tier")).toBeInTheDocument();
+    expect(screen.getByText("Medium Tier")).toBeInTheDocument();
+    expect(screen.getByText("Complex Tier")).toBeInTheDocument();
+    expect(screen.getByText("Reasoning Tier")).toBeInTheDocument();
+  });
+
+  it("should show example queries for each tier", () => {
+    renderWithProviders(
+      <ComplexityRouterConfig
+        modelInfo={mockModelInfo}
+        value={defaultTiers}
+        onChange={vi.fn()}
+      />
+    );
+    expect(screen.getByText(/Hello!/)).toBeInTheDocument();
+    expect(screen.getByText(/Explain how REST APIs work/)).toBeInTheDocument();
+    expect(screen.getByText(/Design a microservices architecture/)).toBeInTheDocument();
+    expect(screen.getByText(/Think step by step/)).toBeInTheDocument();
+  });
+
+  it("should display the how classification works section", () => {
+    renderWithProviders(
+      <ComplexityRouterConfig
+        modelInfo={mockModelInfo}
+        value={defaultTiers}
+        onChange={vi.fn()}
+      />
+    );
+    expect(screen.getByText("How Classification Works")).toBeInTheDocument();
+  });
+
+  it("should show score thresholds in the classification section", () => {
+    renderWithProviders(
+      <ComplexityRouterConfig
+        modelInfo={mockModelInfo}
+        value={defaultTiers}
+        onChange={vi.fn()}
+      />
+    );
+    expect(screen.getByText(/Score < 0.15/)).toBeInTheDocument();
+    expect(screen.getByText(/Score 0.15 - 0.35/)).toBeInTheDocument();
+    expect(screen.getByText(/Score 0.35 - 0.60/)).toBeInTheDocument();
+    expect(screen.getByText(/Score > 0.60/)).toBeInTheDocument();
+  });
+});

ui/litellm-dashboard/src/components/agents/agent_card_grid.test.tsx

@@ -0,0 +1,90 @@
+import { renderWithProviders, screen } from "../../../tests/test-utils";
+import { vi } from "vitest";
+import AgentCardGrid from "./agent_card_grid";
+import type { Agent, AgentKeyInfo } from "./types";
+
+vi.mock("./agent_card", () => ({
+  default: ({ agent, onAgentClick }: any) => (
+    <div data-testid={`agent-card-${agent.agent_id}`} onClick={() => onAgentClick(agent.agent_id)}>
+      {agent.agent_name}
+    </div>
+  ),
+}));
+
+const mockAgents: Agent[] = [
+  {
+    agent_id: "agent-1",
+    agent_name: "Test Agent 1",
+    litellm_params: { model: "gpt-4" },
+    agent_card_params: { description: "First agent" },
+  },
+  {
+    agent_id: "agent-2",
+    agent_name: "Test Agent 2",
+    litellm_params: { model: "claude-3" },
+    agent_card_params: { description: "Second agent" },
+  },
+];
+
+const mockKeyInfoMap: Record<string, AgentKeyInfo> = {
+  "agent-1": { has_key: true, key_alias: "key-1" },
+  "agent-2": { has_key: false },
+};
+
+const defaultProps = {
+  agentsList: mockAgents,
+  keyInfoMap: mockKeyInfoMap,
+  isLoading: false,
+  onDeleteClick: vi.fn(),
+  accessToken: "test-token",
+  onAgentUpdated: vi.fn(),
+  isAdmin: true,
+  onAgentClick: vi.fn(),
+};
+
+describe("AgentCardGrid", () => {
+  it("should render", () => {
+    renderWithProviders(<AgentCardGrid {...defaultProps} />);
+    expect(screen.getByText("Test Agent 1")).toBeInTheDocument();
+  });
+
+  it("should render all agent cards", () => {
+    renderWithProviders(<AgentCardGrid {...defaultProps} />);
+    expect(screen.getByText("Test Agent 1")).toBeInTheDocument();
+    expect(screen.getByText("Test Agent 2")).toBeInTheDocument();
+  });
+
+  it("should show loading skeletons when isLoading is true", () => {
+    renderWithProviders(<AgentCardGrid {...defaultProps} isLoading={true} />);
+    expect(screen.queryByText("Test Agent 1")).not.toBeInTheDocument();
+  });
+
+  it("should show admin empty state message when no agents and isAdmin", () => {
+    renderWithProviders(
+      <AgentCardGrid {...defaultProps} agentsList={[]} isAdmin={true} />
+    );
+    expect(
+      screen.getByText("No agents found. Create one to get started.")
+    ).toBeInTheDocument();
+  });
+
+  it("should show non-admin empty state message when no agents and not admin", () => {
+    renderWithProviders(
+      <AgentCardGrid {...defaultProps} agentsList={[]} isAdmin={false} />
+    );
+    expect(
+      screen.getByText("No agents found. Contact an admin to create agents.")
+    ).toBeInTheDocument();
+  });
+
+  it("should call onAgentClick when a card is clicked", async () => {
+    const onAgentClick = vi.fn();
+    renderWithProviders(
+      <AgentCardGrid {...defaultProps} onAgentClick={onAgentClick} />
+    );
+    const { default: userEvent } = await import("@testing-library/user-event");
+    const user = userEvent.setup();
+    await user.click(screen.getByTestId("agent-card-agent-1"));
+    expect(onAgentClick).toHaveBeenCalledWith("agent-1");
+  });
+});

ui/litellm-dashboard/src/components/common_components/RateLimitTypeFormItem.test.tsx

@@ -0,0 +1,61 @@
+import { renderWithProviders, screen } from "../../../tests/test-utils";
+import userEvent from "@testing-library/user-event";
+import { vi } from "vitest";
+import { Form } from "antd";
+import React from "react";
+import { RateLimitTypeFormItem } from "./RateLimitTypeFormItem";
+
+const Wrapper: React.FC<{ children: React.ReactNode }> = ({ children }) => (
+  <Form>{children}</Form>
+);
+
+describe("RateLimitTypeFormItem", () => {
+  it("should render", () => {
+    renderWithProviders(
+      <Wrapper>
+        <RateLimitTypeFormItem type="tpm" name="tpm_type" />
+      </Wrapper>
+    );
+    expect(screen.getByText(/TPM Rate Limit Type/)).toBeInTheDocument();
+  });
+
+  it("should display TPM label for tpm type", () => {
+    renderWithProviders(
+      <Wrapper>
+        <RateLimitTypeFormItem type="tpm" name="tpm_type" />
+      </Wrapper>
+    );
+    expect(screen.getByText(/TPM Rate Limit Type/)).toBeInTheDocument();
+  });
+
+  it("should display RPM label for rpm type", () => {
+    renderWithProviders(
+      <Wrapper>
+        <RateLimitTypeFormItem type="rpm" name="rpm_type" />
+      </Wrapper>
+    );
+    expect(screen.getByText(/RPM Rate Limit Type/)).toBeInTheDocument();
+  });
+
+  it("should show the select placeholder by default", () => {
+    renderWithProviders(
+      <Wrapper>
+        <RateLimitTypeFormItem type="tpm" name="tpm_type" />
+      </Wrapper>
+    );
+    expect(screen.getByText("Select rate limit type")).toBeInTheDocument();
+  });
+
+  it("should call onChange when provided", async () => {
+    const user = userEvent.setup();
+    const onChange = vi.fn();
+    renderWithProviders(
+      <Wrapper>
+        <RateLimitTypeFormItem type="tpm" name="tpm_type" onChange={onChange} />
+      </Wrapper>
+    );
+    await user.click(screen.getByRole("combobox"));
+    await user.click(screen.getByText("Guaranteed throughput"));
+    expect(onChange).toHaveBeenCalledWith("guaranteed_throughput");
+  });
+});

ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx

@@ -60,6 +60,7 @@ import CodeInterpreterOutput from "./CodeInterpreterOutput";
 import CodeInterpreterTool from "./CodeInterpreterTool";
 import { generateCodeSnippet } from "./CodeSnippets";
 import EndpointSelector from "./EndpointSelector";
+import FilePreviewCard from "./FilePreviewCard";
 import MCPEventsDisplay from "./MCPEventsDisplay";
 import type { MCPEvent } from "../../mcp_tools/types";
 import { EndpointType, getEndpointType } from "./mode_endpoint_mapping";
@@ -2231,67 +2232,19 @@ const ChatUI: React.FC<ChatUIProps> = ({
 
               {/* Show file previews above input when files are uploaded */}
               {endpointType === EndpointType.RESPONSES && responsesUploadedImage && (
-                <div className="mb-2">
-                  <div className="flex items-center gap-3 p-3 bg-gray-50 rounded-lg border border-gray-200">
-                    <div className="relative inline-block">
-                      {responsesUploadedImage.name.toLowerCase().endsWith(".pdf") ? (
-                        <div className="w-10 h-10 rounded-md bg-red-500 flex items-center justify-center">
-                          <FilePdfOutlined style={{ fontSize: "16px", color: "white" }} />
-                        </div>
-                      ) : (
-                        <img
-                          src={responsesImagePreviewUrl || ""}
-                          alt="Upload preview"
-                          className="w-10 h-10 rounded-md border border-gray-200 object-cover"
-                        />
-                      )}
-                    </div>
-                    <div className="flex-1 min-w-0">
-                      <div className="text-sm font-medium text-gray-900 truncate">{responsesUploadedImage.name}</div>
-                      <div className="text-xs text-gray-500">
-                        {responsesUploadedImage.name.toLowerCase().endsWith(".pdf") ? "PDF" : "Image"}
-                      </div>
-                    </div>
-                    <button
-                      className="flex items-center justify-center w-6 h-6 text-gray-400 hover:text-gray-600 hover:bg-gray-200 rounded-full transition-colors"
-                      onClick={handleRemoveResponsesImage}
-                    >
-                      <DeleteOutlined style={{ fontSize: "12px" }} />
-                    </button>
-                  </div>
-                </div>
+                <FilePreviewCard
+                  file={responsesUploadedImage}
+                  previewUrl={responsesImagePreviewUrl}
+                  onRemove={handleRemoveResponsesImage}
+                />
               )}
 
               {endpointType === EndpointType.CHAT && chatUploadedImage && (
-                <div className="mb-2">
-                  <div className="flex items-center gap-3 p-3 bg-gray-50 rounded-lg border border-gray-200">
-                    <div className="relative inline-block">
-                      {chatUploadedImage.name.toLowerCase().endsWith(".pdf") ? (
-                        <div className="w-10 h-10 rounded-md bg-red-500 flex items-center justify-center">
-                          <FilePdfOutlined style={{ fontSize: "16px", color: "white" }} />
-                        </div>
-                      ) : (
-                        <img
-                          src={chatImagePreviewUrl || ""}
-                          alt="Upload preview"
-                          className="w-10 h-10 rounded-md border border-gray-200 object-cover"
-                        />
-                      )}
-                    </div>
-                    <div className="flex-1 min-w-0">
-                      <div className="text-sm font-medium text-gray-900 truncate">{chatUploadedImage.name}</div>
-                      <div className="text-xs text-gray-500">
-                        {chatUploadedImage.name.toLowerCase().endsWith(".pdf") ? "PDF" : "Image"}
-                      </div>
-                    </div>
-                    <button
-                      className="flex items-center justify-center w-6 h-6 text-gray-400 hover:text-gray-600 hover:bg-gray-200 rounded-full transition-colors"
-                      onClick={handleRemoveChatImage}
-                    >
-                      <DeleteOutlined style={{ fontSize: "12px" }} />
-                    </button>
-                  </div>
-                </div>
+                <FilePreviewCard
+                  file={chatUploadedImage}
+                  previewUrl={chatImagePreviewUrl}
+                  onRemove={handleRemoveChatImage}
+                />
               )}
 
               {/* Code Interpreter indicator and sample prompts when enabled */}

ui/litellm-dashboard/src/components/playground/chat_ui/FilePreviewCard.test.tsx

@@ -0,0 +1,70 @@
+import { render, screen } from "@testing-library/react";
+import userEvent from "@testing-library/user-event";
+import { describe, expect, it, vi } from "vitest";
+import FilePreviewCard from "./FilePreviewCard";
+
+function makeFile(name: string): File {
+  return new File(["dummy"], name, { type: "application/octet-stream" });
+}
+
+describe("FilePreviewCard", () => {
+  it("should render", () => {
+    render(
+      <FilePreviewCard file={makeFile("photo.png")} previewUrl={null} onRemove={vi.fn()} />
+    );
+    expect(screen.getByText("photo.png")).toBeInTheDocument();
+  });
+
+  it("should display the file name", () => {
+    render(
+      <FilePreviewCard file={makeFile("my-screenshot.jpg")} previewUrl={null} onRemove={vi.fn()} />
+    );
+    expect(screen.getByText("my-screenshot.jpg")).toBeInTheDocument();
+  });
+
+  it("should show 'Image' label for non-PDF files", () => {
+    render(
+      <FilePreviewCard file={makeFile("photo.png")} previewUrl="blob:http://localhost/abc" onRemove={vi.fn()} />
+    );
+    expect(screen.getByText("Image")).toBeInTheDocument();
+  });
+
+  it("should show 'PDF' label for PDF files", () => {
+    render(
+      <FilePreviewCard file={makeFile("report.pdf")} previewUrl={null} onRemove={vi.fn()} />
+    );
+    expect(screen.getByText("PDF")).toBeInTheDocument();
+  });
+
+  it("should render an image preview when the file is not a PDF", () => {
+    render(
+      <FilePreviewCard file={makeFile("photo.png")} previewUrl="blob:http://localhost/abc" onRemove={vi.fn()} />
+    );
+    expect(screen.getByAltText("Upload preview")).toBeInTheDocument();
+  });
+
+  it("should not render an image preview when the file is a PDF", () => {
+    render(
+      <FilePreviewCard file={makeFile("doc.PDF")} previewUrl={null} onRemove={vi.fn()} />
+    );
+    expect(screen.queryByAltText("Upload preview")).not.toBeInTheDocument();
+  });
+
+  it("should call onRemove when the remove button is clicked", async () => {
+    const onRemove = vi.fn();
+    const user = userEvent.setup();
+    render(
+      <FilePreviewCard file={makeFile("photo.png")} previewUrl={null} onRemove={onRemove} />
+    );
+    await user.click(screen.getByRole("button"));
+    expect(onRemove).toHaveBeenCalledOnce();
+  });
+
+  it("should treat .PDF (uppercase) as a PDF file", () => {
+    render(
+      <FilePreviewCard file={makeFile("REPORT.PDF")} previewUrl={null} onRemove={vi.fn()} />
+    );
+    expect(screen.getByText("PDF")).toBeInTheDocument();
+    expect(screen.queryByAltText("Upload preview")).not.toBeInTheDocument();
+  });
+});

ui/litellm-dashboard/src/components/playground/chat_ui/FilePreviewCard.tsx

@@ -0,0 +1,45 @@
+import { DeleteOutlined, FilePdfOutlined } from "@ant-design/icons";
+
+interface FilePreviewCardProps {
+  file: File;
+  previewUrl: string | null;
+  onRemove: () => void;
+}
+
+function FilePreviewCard({ file, previewUrl, onRemove }: FilePreviewCardProps) {
+  const isPdf = file.name.toLowerCase().endsWith(".pdf");
+
+  return (
+    <div className="mb-2">
+      <div className="flex items-center gap-3 p-3 bg-gray-50 rounded-lg border border-gray-200">
+        <div className="relative inline-block">
+          {isPdf ? (
+            <div className="w-10 h-10 rounded-md bg-red-500 flex items-center justify-center">
+              <FilePdfOutlined style={{ fontSize: "16px", color: "white" }} />
+            </div>
+          ) : (
+            <img
+              src={previewUrl || ""}
+              alt="Upload preview"
+              className="w-10 h-10 rounded-md border border-gray-200 object-cover"
+            />
+          )}
+        </div>
+        <div className="flex-1 min-w-0">
+          <div className="text-sm font-medium text-gray-900 truncate">{file.name}</div>
+          <div className="text-xs text-gray-500">
+            {isPdf ? "PDF" : "Image"}
+          </div>
+        </div>
+        <button
+          className="flex items-center justify-center w-6 h-6 text-gray-400 hover:text-gray-600 hover:bg-gray-200 rounded-full transition-colors"
+          onClick={onRemove}
+        >
+          <DeleteOutlined style={{ fontSize: "12px" }} />
+        </button>
+      </div>
+    </div>
+  );
+}
+
+export default FilePreviewCard;