Merge branch 'main' into litellm_mcp_control_internet

.circleci/config.yml

@@ -1372,6 +1372,51 @@ jobs:
           paths:
             - mcp_coverage.xml
             - mcp_coverage
+  agent_testing:
+    docker:
+      - image: cimg/python:3.11
+        auth:
+          username: ${DOCKERHUB_USERNAME}
+          password: ${DOCKERHUB_PASSWORD}
+    working_directory: ~/project
+
+    steps:
+      - checkout
+      - setup_google_dns
+      - run:
+          name: Install Dependencies
+          command: |
+            python -m pip install --upgrade pip
+            python -m pip install -r requirements.txt
+            pip install "pytest==7.3.1"
+            pip install "pytest-retry==1.6.3"
+            pip install "pytest-cov==5.0.0"
+            pip install "pytest-asyncio==0.21.1"
+            pip install "respx==0.22.0"
+            pip install "pydantic==2.11.0"
+            pip install "a2a-sdk"
+      # Run pytest and generate JUnit XML report
+      - run:
+          name: Run tests
+          command: |
+            pwd
+            ls
+            python -m pytest -vv tests/agent_tests --ignore=tests/agent_tests/local_only_agent_tests --cov=litellm --cov-report=xml -x -s -v --junitxml=test-results/junit.xml --durations=5
+          no_output_timeout: 120m
+      - run:
+          name: Rename the coverage files
+          command: |
+            mv coverage.xml agent_coverage.xml
+            mv .coverage agent_coverage
+
+      # Store test results
+      - store_test_results:
+          path: test-results
+      - persist_to_workspace:
+          root: .
+          paths:
+            - agent_coverage.xml
+            - agent_coverage
   guardrails_testing:
     docker:
       - image: cimg/python:3.11
@@ -4264,6 +4309,12 @@ workflows:
               only:
                 - main
                 - /litellm_.*/
+      - agent_testing:
+          filters:
+            branches:
+              only:
+                - main
+                - /litellm_.*/
       - guardrails_testing:
           filters:
             branches:
@@ -4371,6 +4422,7 @@ workflows:
             - llm_translation_testing
             - realtime_translation_testing
             - mcp_testing
+            - agent_testing
             - google_generate_content_endpoint_testing
             - guardrails_testing
             - llm_responses_api_testing
@@ -4449,6 +4501,7 @@ workflows:
             - llm_translation_testing
             - realtime_translation_testing
             - mcp_testing
+            - agent_testing
             - google_generate_content_endpoint_testing
             - llm_responses_api_testing
             - ocr_testing

docs/my-website/docs/proxy/guardrails/custom_code_guardrail.md

@@ -61,15 +61,23 @@ curl -X POST http://localhost:4000/chat/completions \
 
 ### Function Signature
 
-Your code must define an `apply_guardrail` function:
+Your code must define an `apply_guardrail` function. It can be either sync or async:
 
 ```python
+# Sync version
 def apply_guardrail(inputs, request_data, input_type):
     # inputs: see table below
     # request_data: {"model": "...", "user_id": "...", "team_id": "...", "metadata": {...}}
     # input_type: "request" or "response"
     
     return allow()  # or block() or modify()
+
+# Async version (recommended when using HTTP primitives)
+async def apply_guardrail(inputs, request_data, input_type):
+    response = await http_post("https://api.example.com/check", body={"text": inputs["texts"][0]})
+    if response["success"] and response["body"].get("flagged"):
+        return block("Content flagged")
+    return allow()
 ```
 
 ### `inputs` Parameter
@@ -145,6 +153,29 @@ def apply_guardrail(inputs, request_data, input_type):
 | `char_count(text)` | Count characters |
 | `lower(text)` / `upper(text)` / `trim(text)` | String transforms |
 
+### HTTP Requests (Async)
+
+Make async HTTP requests to external APIs for additional validation or content moderation.
+
+| Function | Description |
+|----------|-------------|
+| `await http_request(url, method, headers, body, timeout)` | General async HTTP request |
+| `await http_get(url, headers, timeout)` | Async GET request |
+| `await http_post(url, body, headers, timeout)` | Async POST request |
+
+**Response format:**
+```python
+{
+    "status_code": 200,        # HTTP status code
+    "body": {...},             # Response body (parsed JSON or string)
+    "headers": {...},          # Response headers
+    "success": True,           # True if status code is 2xx
+    "error": None              # Error message if request failed
+}
+```
+
+**Note:** When using HTTP primitives, define your function as `async def apply_guardrail(...)` for non-blocking execution.
+
 ## Examples
 
 ### Block PII (SSN)
@@ -213,6 +244,29 @@ def apply_guardrail(inputs, request_data, input_type):
     return allow()
 ```
 
+### Call External Moderation API (Async)
+
+```python
+async def apply_guardrail(inputs, request_data, input_type):
+    # Call an external moderation API
+    for text in inputs["texts"]:
+        response = await http_post(
+            "https://api.example.com/moderate",
+            body={"text": text, "user_id": request_data["user_id"]},
+            headers={"Authorization": "Bearer YOUR_API_KEY"},
+            timeout=10
+        )
+        
+        if not response["success"]:
+            # API call failed - decide whether to allow or block
+            return allow()
+        
+        if response["body"].get("flagged"):
+            return block(response["body"].get("reason", "Content flagged"))
+    
+    return allow()
+```
+
 ### Combine Multiple Checks
 
 ```python
@@ -241,8 +295,8 @@ Custom code runs in a restricted environment:
 
 - ❌ No `import` statements
 - ❌ No file I/O
-- ❌ No network access
 - ❌ No `exec()` or `eval()`
+- ✅ HTTP requests via built-in `http_request`, `http_get`, `http_post` primitives
 - ✅ Only LiteLLM-provided primitives available
 
 ## Per-Request Usage

litellm/a2a_protocol/__init__.py

@@ -39,6 +39,12 @@ Example usage (class-based):
 """
 
 from litellm.a2a_protocol.client import A2AClient
+from litellm.a2a_protocol.exceptions import (
+    A2AAgentCardError,
+    A2AConnectionError,
+    A2AError,
+    A2ALocalhostURLError,
+)
 from litellm.a2a_protocol.main import (
     aget_agent_card,
     asend_message,
@@ -49,11 +55,19 @@ from litellm.a2a_protocol.main import (
 from litellm.types.agents import LiteLLMSendMessageResponse
 
 __all__ = [
+    # Client
     "A2AClient",
+    # Functions
     "asend_message",
     "send_message",
     "asend_message_streaming",
     "aget_agent_card",
     "create_a2a_client",
+    # Response types
     "LiteLLMSendMessageResponse",
+    # Exceptions
+    "A2AError",
+    "A2AConnectionError",
+    "A2AAgentCardError",
+    "A2ALocalhostURLError",
 ]

litellm/a2a_protocol/card_resolver.py

@@ -7,6 +7,7 @@ Extends the A2A SDK's card resolver to support multiple well-known paths.
 from typing import TYPE_CHECKING, Any, Dict, Optional
 
 from litellm._logging import verbose_logger
+from litellm.constants import LOCALHOST_URL_PATTERNS
 
 if TYPE_CHECKING:
     from a2a.types import AgentCard
@@ -26,15 +27,61 @@ except ImportError:
     pass
 
 
+def is_localhost_or_internal_url(url: Optional[str]) -> bool:
+    """
+    Check if a URL is a localhost or internal URL.
+
+    This detects common development URLs that are accidentally left in
+    agent cards when deploying to production.
+
+    Args:
+        url: The URL to check
+
+    Returns:
+        True if the URL is localhost/internal
+    """
+    if not url:
+        return False
+
+    url_lower = url.lower()
+
+    return any(pattern in url_lower for pattern in LOCALHOST_URL_PATTERNS)
+
+
+def fix_agent_card_url(agent_card: "AgentCard", base_url: str) -> "AgentCard":
+    """
+    Fix the agent card URL if it contains a localhost/internal address.
+
+    Many A2A agents are deployed with agent cards that contain internal URLs
+    like "http://0.0.0.0:8001/" or "http://localhost:8000/". This function
+    replaces such URLs with the provided base_url.
+
+    Args:
+        agent_card: The agent card to fix
+        base_url: The base URL to use as replacement
+
+    Returns:
+        The agent card with the URL fixed if necessary
+    """
+    card_url = getattr(agent_card, "url", None)
+
+    if card_url and is_localhost_or_internal_url(card_url):
+        # Normalize base_url to ensure it ends with /
+        fixed_url = base_url.rstrip("/") + "/"
+        agent_card.url = fixed_url
+
+    return agent_card
+
+
 class LiteLLMA2ACardResolver(_A2ACardResolver):  # type: ignore[misc]
     """
     Custom A2A card resolver that supports multiple well-known paths.
-    
+
     Extends the base A2ACardResolver to try both:
     - /.well-known/agent-card.json (standard)
     - /.well-known/agent.json (previous/alternative)
     """
-    
+
     async def get_agent_card(
         self,
         relative_card_path: Optional[str] = None,
@@ -42,17 +89,17 @@ class LiteLLMA2ACardResolver(_A2ACardResolver):  # type: ignore[misc]
     ) -> "AgentCard":
         """
         Fetch the agent card, trying multiple well-known paths.
-        
+
         First tries the standard path, then falls back to the previous path.
-        
+
         Args:
             relative_card_path: Optional path to the agent card endpoint.
                 If None, tries both well-known paths.
             http_kwargs: Optional dictionary of keyword arguments to pass to httpx.get
-            
+
         Returns:
             AgentCard from the A2A agent
-            
+
         Raises:
             A2AClientHTTPError or A2AClientJSONError if both paths fail
         """
@@ -62,13 +109,13 @@ class LiteLLMA2ACardResolver(_A2ACardResolver):  # type: ignore[misc]
                 relative_card_path=relative_card_path,
                 http_kwargs=http_kwargs,
             )
-        
+
         # Try both well-known paths
         paths = [
             AGENT_CARD_WELL_KNOWN_PATH,
             PREV_AGENT_CARD_WELL_KNOWN_PATH,
         ]
-        
+
         last_error = None
         for path in paths:
             try:
@@ -85,11 +132,11 @@ class LiteLLMA2ACardResolver(_A2ACardResolver):  # type: ignore[misc]
                 )
                 last_error = e
                 continue
-        
+
         # If we get here, all paths failed - re-raise the last error
         if last_error is not None:
             raise last_error
-        
+
         # This shouldn't happen, but just in case
         raise Exception(
             f"Failed to fetch agent card from {self.base_url}. "

litellm/a2a_protocol/exception_mapping_utils.py

@@ -0,0 +1,203 @@
+"""
+A2A Protocol Exception Mapping Utils.
+
+Maps A2A SDK exceptions to LiteLLM A2A exception types.
+"""
+
+from typing import TYPE_CHECKING, Any, Optional
+
+from litellm._logging import verbose_logger
+from litellm.a2a_protocol.card_resolver import (
+    fix_agent_card_url,
+    is_localhost_or_internal_url,
+)
+from litellm.a2a_protocol.exceptions import (
+    A2AAgentCardError,
+    A2AConnectionError,
+    A2AError,
+    A2ALocalhostURLError,
+)
+from litellm.constants import CONNECTION_ERROR_PATTERNS
+
+if TYPE_CHECKING:
+    from a2a.client import A2AClient as A2AClientType
+
+
+# Runtime import
+A2A_SDK_AVAILABLE = False
+try:
+    from a2a.client import A2AClient as _A2AClient  # type: ignore[no-redef]
+
+    A2A_SDK_AVAILABLE = True
+except ImportError:
+    _A2AClient = None  # type: ignore[misc]
+
+
+class A2AExceptionCheckers:
+    """
+    Helper class for checking various A2A error conditions.
+    """
+
+    @staticmethod
+    def is_connection_error(error_str: str) -> bool:
+        """
+        Check if an error string indicates a connection error.
+
+        Args:
+            error_str: The error string to check
+
+        Returns:
+            True if the error indicates a connection issue
+        """
+        if not isinstance(error_str, str):
+            return False
+
+        error_str_lower = error_str.lower()
+        return any(pattern in error_str_lower for pattern in CONNECTION_ERROR_PATTERNS)
+
+    @staticmethod
+    def is_localhost_url(url: Optional[str]) -> bool:
+        """
+        Check if a URL is a localhost/internal URL.
+
+        Args:
+            url: The URL to check
+
+        Returns:
+            True if the URL is localhost/internal
+        """
+        return is_localhost_or_internal_url(url)
+
+    @staticmethod
+    def is_agent_card_error(error_str: str) -> bool:
+        """
+        Check if an error string indicates an agent card error.
+
+        Args:
+            error_str: The error string to check
+
+        Returns:
+            True if the error is related to agent card fetching/parsing
+        """
+        if not isinstance(error_str, str):
+            return False
+
+        error_str_lower = error_str.lower()
+        agent_card_patterns = [
+            "agent card",
+            "agent-card",
+            ".well-known",
+            "card not found",
+            "invalid agent",
+        ]
+        return any(pattern in error_str_lower for pattern in agent_card_patterns)
+
+
+def map_a2a_exception(
+    original_exception: Exception,
+    card_url: Optional[str] = None,
+    api_base: Optional[str] = None,
+    model: Optional[str] = None,
+) -> Exception:
+    """
+    Map an A2A SDK exception to a LiteLLM A2A exception type.
+
+    Args:
+        original_exception: The original exception from the A2A SDK
+        card_url: The URL from the agent card (if available)
+        api_base: The original API base URL
+        model: The model/agent name
+
+    Returns:
+        A mapped LiteLLM A2A exception
+
+    Raises:
+        A2ALocalhostURLError: If the error is a connection error to a localhost URL
+        A2AConnectionError: If the error is a general connection error
+        A2AAgentCardError: If the error is related to agent card issues
+        A2AError: For other A2A-related errors
+    """
+    error_str = str(original_exception)
+
+    # Check for localhost URL connection error (special case - retryable)
+    if (
+        card_url
+        and api_base
+        and A2AExceptionCheckers.is_localhost_url(card_url)
+        and A2AExceptionCheckers.is_connection_error(error_str)
+    ):
+        raise A2ALocalhostURLError(
+            localhost_url=card_url,
+            base_url=api_base,
+            original_error=original_exception,
+            model=model,
+        )
+
+    # Check for agent card errors
+    if A2AExceptionCheckers.is_agent_card_error(error_str):
+        raise A2AAgentCardError(
+            message=error_str,
+            url=api_base,
+            model=model,
+        )
+
+    # Check for general connection errors
+    if A2AExceptionCheckers.is_connection_error(error_str):
+        raise A2AConnectionError(
+            message=error_str,
+            url=card_url or api_base,
+            model=model,
+        )
+
+    # Default: wrap in generic A2AError
+    raise A2AError(
+        message=error_str,
+        model=model,
+    )
+
+
+def handle_a2a_localhost_retry(
+    error: A2ALocalhostURLError,
+    agent_card: Any,
+    a2a_client: "A2AClientType",
+    is_streaming: bool = False,
+) -> "A2AClientType":
+    """
+    Handle A2ALocalhostURLError by fixing the URL and creating a new client.
+
+    This is called when we catch an A2ALocalhostURLError and want to retry
+    with the corrected URL.
+
+    Args:
+        error: The localhost URL error
+        agent_card: The agent card object to fix
+        a2a_client: The current A2A client
+        is_streaming: Whether this is a streaming request (for logging)
+
+    Returns:
+        A new A2A client with the fixed URL
+
+    Raises:
+        ImportError: If the A2A SDK is not installed
+    """
+    if not A2A_SDK_AVAILABLE or _A2AClient is None:
+        raise ImportError(
+            "A2A SDK is required for localhost retry handling. "
+            "Install it with: pip install a2a"
+        )
+
+    request_type = "streaming " if is_streaming else ""
+    verbose_logger.warning(
+        f"A2A {request_type}request to '{error.localhost_url}' failed: {error.original_error}. "
+        f"Agent card contains localhost/internal URL. "
+        f"Retrying with base_url '{error.base_url}'."
+    )
+
+    # Fix the agent card URL
+    fix_agent_card_url(agent_card, error.base_url)
+
+    # Create a new client with the fixed agent card (transport caches URL)
+    return _A2AClient(
+        httpx_client=a2a_client._transport.httpx_client,  # type: ignore[union-attr]
+        agent_card=agent_card,
+    )

litellm/a2a_protocol/exceptions.py

@@ -0,0 +1,150 @@
+"""
+A2A Protocol Exceptions.
+
+Custom exception types for A2A protocol operations, following LiteLLM's exception pattern.
+"""
+
+from typing import Optional
+
+import httpx
+
+
+class A2AError(Exception):
+    """
+    Base exception for A2A protocol errors.
+
+    Follows the same pattern as LiteLLM's main exceptions.
+    """
+
+    def __init__(
+        self,
+        message: str,
+        status_code: int = 500,
+        llm_provider: str = "a2a_agent",
+        model: Optional[str] = None,
+        response: Optional[httpx.Response] = None,
+        litellm_debug_info: Optional[str] = None,
+        max_retries: Optional[int] = None,
+        num_retries: Optional[int] = None,
+    ):
+        self.status_code = status_code
+        self.message = f"litellm.A2AError: {message}"
+        self.llm_provider = llm_provider
+        self.model = model
+        self.litellm_debug_info = litellm_debug_info
+        self.max_retries = max_retries
+        self.num_retries = num_retries
+        self.response = response or httpx.Response(
+            status_code=self.status_code,
+            request=httpx.Request(method="POST", url="https://litellm.ai"),
+        )
+        super().__init__(self.message)
+
+    def __str__(self) -> str:
+        _message = self.message
+        if self.num_retries:
+            _message += f" LiteLLM Retried: {self.num_retries} times"
+        if self.max_retries:
+            _message += f", LiteLLM Max Retries: {self.max_retries}"
+        return _message
+
+    def __repr__(self) -> str:
+        return self.__str__()
+
+
+class A2AConnectionError(A2AError):
+    """
+    Raised when connection to an A2A agent fails.
+
+    This typically occurs when:
+    - The agent is unreachable
+    - The agent card contains a localhost/internal URL
+    - Network issues prevent connection
+    """
+
+    def __init__(
+        self,
+        message: str,
+        url: Optional[str] = None,
+        model: Optional[str] = None,
+        response: Optional[httpx.Response] = None,
+        litellm_debug_info: Optional[str] = None,
+        max_retries: Optional[int] = None,
+        num_retries: Optional[int] = None,
+    ):
+        self.url = url
+        super().__init__(
+            message=message,
+            status_code=503,
+            llm_provider="a2a_agent",
+            model=model,
+            response=response,
+            litellm_debug_info=litellm_debug_info,
+            max_retries=max_retries,
+            num_retries=num_retries,
+        )
+
+
+class A2AAgentCardError(A2AError):
+    """
+    Raised when there's an issue with the agent card.
+
+    This includes:
+    - Failed to fetch agent card
+    - Invalid agent card format
+    - Missing required fields
+    """
+
+    def __init__(
+        self,
+        message: str,
+        url: Optional[str] = None,
+        model: Optional[str] = None,
+        response: Optional[httpx.Response] = None,
+        litellm_debug_info: Optional[str] = None,
+    ):
+        self.url = url
+        super().__init__(
+            message=message,
+            status_code=404,
+            llm_provider="a2a_agent",
+            model=model,
+            response=response,
+            litellm_debug_info=litellm_debug_info,
+        )
+
+
+class A2ALocalhostURLError(A2AConnectionError):
+    """
+    Raised when an agent card contains a localhost/internal URL.
+
+    Many A2A agents are deployed with agent cards that contain internal URLs
+    like "http://0.0.0.0:8001/" or "http://localhost:8000/". This error
+    indicates that the URL needs to be corrected and the request should be retried.
+
+    Attributes:
+        localhost_url: The localhost/internal URL found in the agent card
+        base_url: The public base URL that should be used instead
+        original_error: The original connection error that was raised
+    """
+
+    def __init__(
+        self,
+        localhost_url: str,
+        base_url: str,
+        original_error: Optional[Exception] = None,
+        model: Optional[str] = None,
+    ):
+        self.localhost_url = localhost_url
+        self.base_url = base_url
+        self.original_error = original_error
+
+        message = (
+            f"Agent card contains localhost/internal URL '{localhost_url}'. "
+            f"Retrying with base URL '{base_url}'."
+        )
+        super().__init__(
+            message=message,
+            url=localhost_url,
+            model=model,
+        )

litellm/a2a_protocol/main.py

@@ -44,6 +44,11 @@ except ImportError:
 
 # Import our custom card resolver that supports multiple well-known paths
 from litellm.a2a_protocol.card_resolver import LiteLLMA2ACardResolver
+from litellm.a2a_protocol.exception_mapping_utils import (
+    handle_a2a_localhost_retry,
+    map_a2a_exception,
+)
+from litellm.a2a_protocol.exceptions import A2ALocalhostURLError
 
 # Use our custom resolver instead of the default A2A SDK resolver
 A2ACardResolver = LiteLLMA2ACardResolver
@@ -244,10 +249,50 @@ async def asend_message(
 
     verbose_logger.info(f"A2A send_message request_id={request.id}, agent={agent_name}")
 
-    a2a_response = await a2a_client.send_message(request)
+    # Get agent card URL for localhost retry logic
+    agent_card = getattr(a2a_client, "_litellm_agent_card", None) or getattr(
+        a2a_client, "agent_card", None
+    )
+    card_url = getattr(agent_card, "url", None) if agent_card else None
+
+    # Retry loop: if connection fails due to localhost URL in agent card, retry with fixed URL
+    a2a_response = None
+    for _ in range(2):  # max 2 attempts: original + 1 retry
+        try:
+            a2a_response = await a2a_client.send_message(request)
+            break  # success, exit retry loop
+        except A2ALocalhostURLError as e:
+            # Localhost URL error - fix and retry
+            a2a_client = handle_a2a_localhost_retry(
+                error=e,
+                agent_card=agent_card,
+                a2a_client=a2a_client,
+                is_streaming=False,
+            )
+            card_url = agent_card.url if agent_card else None
+        except Exception as e:
+            # Map exception - will raise A2ALocalhostURLError if applicable
+            try:
+                map_a2a_exception(e, card_url, api_base, model=agent_name)
+            except A2ALocalhostURLError as localhost_err:
+                # Localhost URL error - fix and retry
+                a2a_client = handle_a2a_localhost_retry(
+                    error=localhost_err,
+                    agent_card=agent_card,
+                    a2a_client=a2a_client,
+                    is_streaming=False,
+                )
+                card_url = agent_card.url if agent_card else None
+                continue
+            except Exception:
+                # Re-raise the mapped exception
+                raise
 
     verbose_logger.info(f"A2A send_message completed, request_id={request.id}")
 
+    # a2a_response is guaranteed to be set if we reach here (loop breaks on success or raises)
+    assert a2a_response is not None
+
     # Wrap in LiteLLM response type for _hidden_params support
     response = LiteLLMSendMessageResponse.from_a2a_response(a2a_response)
 
@@ -307,6 +352,48 @@ def send_message(
         )
 
 
+def _build_streaming_logging_obj(
+    request: "SendStreamingMessageRequest",
+    agent_name: str,
+    agent_id: Optional[str],
+    litellm_params: Optional[Dict[str, Any]],
+    metadata: Optional[Dict[str, Any]],
+    proxy_server_request: Optional[Dict[str, Any]],
+) -> Logging:
+    """Build logging object for streaming A2A requests."""
+    start_time = datetime.datetime.now()
+    model = f"a2a_agent/{agent_name}"
+
+    logging_obj = Logging(
+        model=model,
+        messages=[{"role": "user", "content": "streaming-request"}],
+        stream=False,
+        call_type="asend_message_streaming",
+        start_time=start_time,
+        litellm_call_id=str(request.id),
+        function_id=str(request.id),
+    )
+    logging_obj.model = model
+    logging_obj.custom_llm_provider = "a2a_agent"
+    logging_obj.model_call_details["model"] = model
+    logging_obj.model_call_details["custom_llm_provider"] = "a2a_agent"
+    if agent_id:
+        logging_obj.model_call_details["agent_id"] = agent_id
+
+    _litellm_params = litellm_params.copy() if litellm_params else {}
+    if metadata:
+        _litellm_params["metadata"] = metadata
+    if proxy_server_request:
+        _litellm_params["proxy_server_request"] = proxy_server_request
+
+    logging_obj.litellm_params = _litellm_params
+    logging_obj.optional_params = _litellm_params
+    logging_obj.model_call_details["litellm_params"] = _litellm_params
+    logging_obj.model_call_details["metadata"] = metadata or {}
+
+    return logging_obj
+
+
 async def asend_message_streaming(
     a2a_client: Optional["A2AClientType"] = None,
     request: Optional["SendStreamingMessageRequest"] = None,
@@ -403,55 +490,72 @@ async def asend_message_streaming(
 
     verbose_logger.info(f"A2A send_message_streaming request_id={request.id}")
 
-    # Track for logging
-    start_time = datetime.datetime.now()
-    stream = a2a_client.send_message_streaming(request)
-
     # Build logging object for streaming completion callbacks
     agent_card = getattr(a2a_client, "_litellm_agent_card", None) or getattr(
         a2a_client, "agent_card", None
     )
+    card_url = getattr(agent_card, "url", None) if agent_card else None
     agent_name = getattr(agent_card, "name", "unknown") if agent_card else "unknown"
-    model = f"a2a_agent/{agent_name}"
 
-    logging_obj = Logging(
-        model=model,
-        messages=[{"role": "user", "content": "streaming-request"}],
-        stream=False,  # complete response logging after stream ends
-        call_type="asend_message_streaming",
-        start_time=start_time,
-        litellm_call_id=str(request.id),
-        function_id=str(request.id),
-    )
-    logging_obj.model = model
-    logging_obj.custom_llm_provider = "a2a_agent"
-    logging_obj.model_call_details["model"] = model
-    logging_obj.model_call_details["custom_llm_provider"] = "a2a_agent"
-    if agent_id:
-        logging_obj.model_call_details["agent_id"] = agent_id
-
-    # Propagate litellm_params for spend logging (includes cost_per_query, etc.)
-    _litellm_params = litellm_params.copy() if litellm_params else {}
-    # Merge metadata into litellm_params.metadata (required for proxy cost tracking)
-    if metadata:
-        _litellm_params["metadata"] = metadata
-    if proxy_server_request:
-        _litellm_params["proxy_server_request"] = proxy_server_request
-
-    logging_obj.litellm_params = _litellm_params
-    logging_obj.optional_params = _litellm_params  # used by cost calc
-    logging_obj.model_call_details["litellm_params"] = _litellm_params
-    logging_obj.model_call_details["metadata"] = metadata or {}
-
-    iterator = A2AStreamingIterator(
-        stream=stream,
+    logging_obj = _build_streaming_logging_obj(
         request=request,
-        logging_obj=logging_obj,
         agent_name=agent_name,
+        agent_id=agent_id,
+        litellm_params=litellm_params,
+        metadata=metadata,
+        proxy_server_request=proxy_server_request,
     )
 
-    async for chunk in iterator:
-        yield chunk
+    # Retry loop: if connection fails due to localhost URL in agent card, retry with fixed URL
+    # Connection errors in streaming typically occur on first chunk iteration
+    first_chunk = True
+    for attempt in range(2):  # max 2 attempts: original + 1 retry
+        stream = a2a_client.send_message_streaming(request)
+        iterator = A2AStreamingIterator(
+            stream=stream,
+            request=request,
+            logging_obj=logging_obj,
+            agent_name=agent_name,
+        )
+
+        try:
+            first_chunk = True
+            async for chunk in iterator:
+                if first_chunk:
+                    first_chunk = False  # connection succeeded
+                yield chunk
+            return  # stream completed successfully
+        except A2ALocalhostURLError as e:
+            # Only retry on first chunk, not mid-stream
+            if first_chunk and attempt == 0:
+                a2a_client = handle_a2a_localhost_retry(
+                    error=e,
+                    agent_card=agent_card,
+                    a2a_client=a2a_client,
+                    is_streaming=True,
+                )
+                card_url = agent_card.url if agent_card else None
+            else:
+                raise
+        except Exception as e:
+            # Only map exception on first chunk
+            if first_chunk and attempt == 0:
+                try:
+                    map_a2a_exception(e, card_url, api_base, model=agent_name)
+                except A2ALocalhostURLError as localhost_err:
+                    # Localhost URL error - fix and retry
+                    a2a_client = handle_a2a_localhost_retry(
+                        error=localhost_err,
+                        agent_card=agent_card,
+                        a2a_client=a2a_client,
+                        is_streaming=True,
+                    )
+                    card_url = agent_card.url if agent_card else None
+                    continue
+                except Exception:
+                    # Re-raise the mapped exception
+                    raise
+            raise
 
 
 async def create_a2a_client(

litellm/constants.py

@@ -306,6 +306,22 @@ DEFAULT_MAX_TOKENS_FOR_TRITON = int(os.getenv("DEFAULT_MAX_TOKENS_FOR_TRITON", 2
 #### Networking settings ####
 request_timeout: float = float(os.getenv("REQUEST_TIMEOUT", 6000))  # time in seconds
 DEFAULT_A2A_AGENT_TIMEOUT: float = float(os.getenv("DEFAULT_A2A_AGENT_TIMEOUT", 6000))  # 10 minutes
+# Patterns that indicate a localhost/internal URL in A2A agent cards that should be
+# replaced with the original base_url. This is a common misconfiguration where
+# developers deploy agents with development URLs in their agent cards.
+LOCALHOST_URL_PATTERNS: List[str] = [
+    "localhost",
+    "127.0.0.1",
+    "0.0.0.0",
+    "[::1]",  # IPv6 localhost
+]
+# Patterns in error messages that indicate a connection failure
+CONNECTION_ERROR_PATTERNS: List[str] = [
+    "connect",
+    "connection",
+    "network",
+    "refused",
+]
 STREAM_SSE_DONE_STRING: str = "[DONE]"
 STREAM_SSE_DATA_PREFIX: str = "data: "
 ### SPEND TRACKING ###

litellm/experimental_mcp_client/client.py

@@ -209,6 +209,8 @@ class MCPClient:
                     headers["X-API-Key"] = self._mcp_auth_value
                 elif self.auth_type == MCPAuth.authorization:
                     headers["Authorization"] = self._mcp_auth_value
+                elif self.auth_type == MCPAuth.oauth2:
+                    headers["Authorization"] = f"Bearer {self._mcp_auth_value}"
             elif isinstance(self._mcp_auth_value, dict):
                 headers.update(self._mcp_auth_value)
 

litellm/integrations/custom_guardrail.py

@@ -268,6 +268,7 @@ class CustomGuardrail(CustomLogger):
         """
         Returns the guardrail(s) to be run from the metadata or root
         """
+
         if "guardrails" in data:
             return data["guardrails"]
         metadata = data.get("litellm_metadata") or data.get("metadata", {})

litellm/litellm_core_utils/token_counter.py

@@ -706,7 +706,7 @@ def _count_content_list(
             if isinstance(c, str):
                 num_tokens += count_function(c)
             elif c["type"] == "text":
-                num_tokens += count_function(c.get("text", ""))
+                num_tokens += count_function(str(c.get("text", "")))
             elif c["type"] == "image_url":
                 image_url = c.get("image_url")
                 num_tokens += _count_image_tokens(
@@ -722,7 +722,7 @@ def _count_content_list(
             elif c["type"] == "thinking":
                 # Claude extended thinking content block
                 # Count the thinking text and skip signature (opaque signature blob)
-                thinking_text = c.get("thinking", "")
+                thinking_text = str(c.get("thinking", ""))
                 if thinking_text:
                     num_tokens += count_function(thinking_text)
             else:

litellm/llms/a2a/chat/guardrail_translation/README.md

@@ -0,0 +1,155 @@
+# A2A Protocol Guardrail Translation Handler
+
+Handler for processing A2A (Agent-to-Agent) Protocol messages with guardrails.
+
+## Overview
+
+This handler processes A2A JSON-RPC 2.0 input/output by:
+1. Extracting text from message parts (`kind: "text"`)
+2. Applying guardrails to text content
+3. Mapping guardrailed text back to original structure
+
+## A2A Protocol Format
+
+### Input Format (JSON-RPC 2.0)
+
+```json
+{
+  "jsonrpc": "2.0",
+  "id": "request-id",
+  "method": "message/send",
+  "params": {
+    "message": {
+      "kind": "message",
+      "messageId": "...",
+      "role": "user",
+      "parts": [
+        {"kind": "text", "text": "Hello, my SSN is 123-45-6789"}
+      ]
+    },
+    "metadata": {
+      "guardrails": ["block-ssn"]
+    }
+  }
+}
+```
+
+### Output Formats
+
+The handler supports multiple A2A response formats:
+
+**Direct message:**
+```json
+{
+  "result": {
+    "kind": "message",
+    "parts": [{"kind": "text", "text": "Response text"}]
+  }
+}
+```
+
+**Nested message:**
+```json
+{
+  "result": {
+    "message": {
+      "parts": [{"kind": "text", "text": "Response text"}]
+    }
+  }
+}
+```
+
+**Task with artifacts:**
+```json
+{
+  "result": {
+    "kind": "task",
+    "artifacts": [
+      {"parts": [{"kind": "text", "text": "Artifact text"}]}
+    ]
+  }
+}
+```
+
+**Task with status message:**
+```json
+{
+  "result": {
+    "kind": "task",
+    "status": {
+      "message": {
+        "parts": [{"kind": "text", "text": "Status message"}]
+      }
+    }
+  }
+}
+```
+
+**Streaming artifact-update:**
+```json
+{
+  "result": {
+    "kind": "artifact-update",
+    "artifact": {
+      "parts": [{"kind": "text", "text": "Streaming text"}]
+    }
+  }
+}
+```
+
+## Usage
+
+The handler is automatically discovered and applied when guardrails are used with A2A endpoints.
+
+### Via LiteLLM Proxy
+
+```bash
+curl -X POST 'http://localhost:4000/a2a/my-agent' \
+-H 'Content-Type: application/json' \
+-H 'Authorization: Bearer your-api-key' \
+-d '{
+  "jsonrpc": "2.0",
+  "id": "1",
+  "method": "message/send",
+  "params": {
+    "message": {
+      "kind": "message",
+      "messageId": "msg-1",
+      "role": "user",
+      "parts": [{"kind": "text", "text": "Hello, my SSN is 123-45-6789"}]
+    },
+    "metadata": {
+      "guardrails": ["block-ssn"]
+    }
+  }
+}'
+```
+
+### Specifying Guardrails
+
+Guardrails can be specified in the A2A request via the `metadata.guardrails` field:
+
+```json
+{
+  "params": {
+    "message": {...},
+    "metadata": {
+      "guardrails": ["block-ssn", "pii-filter"]
+    }
+  }
+}
+```
+
+## Extension
+
+Override these methods to customize behavior:
+
+- `_extract_texts_from_result()`: Custom text extraction from A2A responses
+- `_extract_texts_from_parts()`: Custom text extraction from message parts
+- `_apply_text_to_path()`: Custom application of guardrailed text
+
+## Call Types
+
+This handler is registered for:
+- `CallTypes.send_message`: Synchronous A2A message sending
+- `CallTypes.asend_message`: Asynchronous A2A message sending

litellm/llms/a2a/chat/guardrail_translation/__init__.py

@@ -0,0 +1,11 @@
+"""A2A Protocol handler for Unified Guardrails."""
+
+from litellm.llms.a2a.chat.guardrail_translation.handler import A2AGuardrailHandler
+from litellm.types.utils import CallTypes
+
+guardrail_translation_mappings = {
+    CallTypes.send_message: A2AGuardrailHandler,
+    CallTypes.asend_message: A2AGuardrailHandler,
+}
+
+__all__ = ["guardrail_translation_mappings"]

litellm/llms/a2a/chat/guardrail_translation/handler.py

@@ -0,0 +1,315 @@
+"""
+A2A Protocol Handler for Unified Guardrails
+
+This module provides guardrail translation support for A2A (Agent-to-Agent) Protocol.
+It handles both JSON-RPC 2.0 input requests and output responses, extracting text
+from message parts and applying guardrails.
+
+A2A Protocol Format:
+- Input: JSON-RPC 2.0 with params.message.parts containing text parts
+- Output: JSON-RPC 2.0 with result containing message/artifact parts
+"""
+
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
+
+from litellm._logging import verbose_proxy_logger
+from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
+from litellm.types.utils import GenericGuardrailAPIInputs
+
+if TYPE_CHECKING:
+    from litellm.integrations.custom_guardrail import CustomGuardrail
+    from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
+    from litellm.proxy._types import UserAPIKeyAuth
+
+
+class A2AGuardrailHandler(BaseTranslation):
+    """
+    Handler for processing A2A Protocol messages with guardrails.
+
+    This class provides methods to:
+    1. Process input messages (pre-call hook) - extracts text from A2A message parts
+    2. Process output responses (post-call hook) - extracts text from A2A response parts
+
+    A2A Message Format:
+    - Input: params.message.parts[].text (where kind == "text")
+    - Output: result.message.parts[].text or result.artifacts[].parts[].text
+    """
+
+    async def process_input_messages(
+        self,
+        data: dict,
+        guardrail_to_apply: "CustomGuardrail",
+        litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None,
+    ) -> Any:
+        """
+        Process A2A input messages by applying guardrails to text content.
+
+        Extracts text from A2A message parts and applies guardrails.
+
+        Args:
+            data: The A2A JSON-RPC 2.0 request data
+            guardrail_to_apply: The guardrail instance to apply
+            litellm_logging_obj: Optional logging object
+
+        Returns:
+            Modified data with guardrails applied to text content
+        """
+        # A2A request format: { "params": { "message": { "parts": [...] } } }
+        params = data.get("params", {})
+        message = params.get("message", {})
+        parts = message.get("parts", [])
+
+        if not parts:
+            verbose_proxy_logger.debug("A2A: No parts in message, skipping guardrail")
+            return data
+
+        texts_to_check: List[str] = []
+        text_part_indices: List[int] = []  # Track which parts contain text
+
+        # Step 1: Extract text from all text parts
+        for part_idx, part in enumerate(parts):
+            if part.get("kind") == "text":
+                text = part.get("text", "")
+                if text:
+                    texts_to_check.append(text)
+                    text_part_indices.append(part_idx)
+
+        # Step 2: Apply guardrail to all texts in batch
+        if texts_to_check:
+            inputs = GenericGuardrailAPIInputs(texts=texts_to_check)
+
+            # Pass the structured A2A message to guardrails
+            inputs["structured_messages"] = [message]
+
+            # Include agent model info if available
+            model = data.get("model")
+            if model:
+                inputs["model"] = model
+
+            guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
+                inputs=inputs,
+                request_data=data,
+                input_type="request",
+                logging_obj=litellm_logging_obj,
+            )
+
+            guardrailed_texts = guardrailed_inputs.get("texts", [])
+
+            # Step 3: Apply guardrailed text back to original parts
+            if guardrailed_texts and len(guardrailed_texts) == len(text_part_indices):
+                for task_idx, part_idx in enumerate(text_part_indices):
+                    parts[part_idx]["text"] = guardrailed_texts[task_idx]
+
+        verbose_proxy_logger.debug("A2A: Processed input message: %s", message)
+
+        return data
+
+    async def process_output_response(
+        self,
+        response: Any,
+        guardrail_to_apply: "CustomGuardrail",
+        litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None,
+        user_api_key_dict: Optional["UserAPIKeyAuth"] = None,
+    ) -> Any:
+        """
+        Process A2A output response by applying guardrails to text content.
+
+        Handles multiple A2A response formats:
+        - Direct message: {"result": {"kind": "message", "parts": [...]}}
+        - Nested message: {"result": {"message": {"parts": [...]}}}
+        - Task with artifacts: {"result": {"kind": "task", "artifacts": [{"parts": [...]}]}}
+        - Task with status message: {"result": {"kind": "task", "status": {"message": {"parts": [...]}}}}
+
+        Args:
+            response: A2A JSON-RPC 2.0 response dict or object
+            guardrail_to_apply: The guardrail instance to apply
+            litellm_logging_obj: Optional logging object
+            user_api_key_dict: User API key metadata
+
+        Returns:
+            Modified response with guardrails applied to text content
+        """
+        # Handle both dict and Pydantic model responses
+        if hasattr(response, "model_dump"):
+            response_dict = response.model_dump()
+            is_pydantic = True
+        elif isinstance(response, dict):
+            response_dict = response
+            is_pydantic = False
+        else:
+            verbose_proxy_logger.warning(
+                "A2A: Unknown response type %s, skipping guardrail", type(response)
+            )
+            return response
+
+        result = response_dict.get("result", {})
+        if not result or not isinstance(result, dict):
+            verbose_proxy_logger.debug("A2A: No result in response, skipping guardrail")
+            return response
+
+        # Find all text-containing parts in the response
+        texts_to_check: List[str] = []
+        # Each mapping is (path_to_parts_list, part_index)
+        # path_to_parts_list is a tuple of keys to navigate to the parts list
+        task_mappings: List[Tuple[Tuple[str, ...], int]] = []
+
+        # Extract texts from all possible locations
+        self._extract_texts_from_result(
+            result=result,
+            texts_to_check=texts_to_check,
+            task_mappings=task_mappings,
+        )
+
+        if not texts_to_check:
+            verbose_proxy_logger.debug("A2A: No text content in response")
+            return response
+
+        # Step 2: Apply guardrail to all texts in batch
+        # Create a request_data dict with response info and user API key metadata
+        request_data: dict = {"response": response_dict}
+
+        # Add user API key metadata with prefixed keys
+        user_metadata = self.transform_user_api_key_dict_to_metadata(user_api_key_dict)
+        if user_metadata:
+            request_data["litellm_metadata"] = user_metadata
+
+        inputs = GenericGuardrailAPIInputs(texts=texts_to_check)
+
+        guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
+            inputs=inputs,
+            request_data=request_data,
+            input_type="response",
+            logging_obj=litellm_logging_obj,
+        )
+
+        guardrailed_texts = guardrailed_inputs.get("texts", [])
+
+        # Step 3: Apply guardrailed text back to original response
+        if guardrailed_texts and len(guardrailed_texts) == len(task_mappings):
+            for task_idx, (path, part_idx) in enumerate(task_mappings):
+                self._apply_text_to_path(
+                    result=result,
+                    path=path,
+                    part_idx=part_idx,
+                    text=guardrailed_texts[task_idx],
+                )
+
+        verbose_proxy_logger.debug("A2A: Processed output response")
+
+        # Update the original response
+        if is_pydantic:
+            # For Pydantic models, we need to update the underlying dict
+            # and the model will reflect the changes
+            response_dict["result"] = result
+            return response
+        else:
+            response["result"] = result
+            return response
+
+    def _extract_texts_from_result(
+        self,
+        result: Dict[str, Any],
+        texts_to_check: List[str],
+        task_mappings: List[Tuple[Tuple[str, ...], int]],
+    ) -> None:
+        """
+        Extract text from all possible locations in an A2A result.
+
+        Handles multiple response formats:
+        1. Direct message with parts: {"parts": [...]}
+        2. Nested message: {"message": {"parts": [...]}}
+        3. Task with artifacts: {"artifacts": [{"parts": [...]}]}
+        4. Task with status message: {"status": {"message": {"parts": [...]}}}
+        5. Streaming artifact-update: {"artifact": {"parts": [...]}}
+        """
+        # Case 1: Direct parts in result (direct message)
+        if "parts" in result:
+            self._extract_texts_from_parts(
+                parts=result["parts"],
+                path=("parts",),
+                texts_to_check=texts_to_check,
+                task_mappings=task_mappings,
+            )
+
+        # Case 2: Nested message
+        message = result.get("message")
+        if message and isinstance(message, dict) and "parts" in message:
+            self._extract_texts_from_parts(
+                parts=message["parts"],
+                path=("message", "parts"),
+                texts_to_check=texts_to_check,
+                task_mappings=task_mappings,
+            )
+
+        # Case 3: Streaming artifact-update (singular artifact)
+        artifact = result.get("artifact")
+        if artifact and isinstance(artifact, dict) and "parts" in artifact:
+            self._extract_texts_from_parts(
+                parts=artifact["parts"],
+                path=("artifact", "parts"),
+                texts_to_check=texts_to_check,
+                task_mappings=task_mappings,
+            )
+
+        # Case 4: Task with status message
+        status = result.get("status", {})
+        if isinstance(status, dict):
+            status_message = status.get("message")
+            if (
+                status_message
+                and isinstance(status_message, dict)
+                and "parts" in status_message
+            ):
+                self._extract_texts_from_parts(
+                    parts=status_message["parts"],
+                    path=("status", "message", "parts"),
+                    texts_to_check=texts_to_check,
+                    task_mappings=task_mappings,
+                )
+
+        # Case 5: Task with artifacts (plural, array)
+        artifacts = result.get("artifacts", [])
+        if artifacts and isinstance(artifacts, list):
+            for artifact_idx, art in enumerate(artifacts):
+                if isinstance(art, dict) and "parts" in art:
+                    self._extract_texts_from_parts(
+                        parts=art["parts"],
+                        path=("artifacts", str(artifact_idx), "parts"),
+                        texts_to_check=texts_to_check,
+                        task_mappings=task_mappings,
+                    )
+
+    def _extract_texts_from_parts(
+        self,
+        parts: List[Dict[str, Any]],
+        path: Tuple[str, ...],
+        texts_to_check: List[str],
+        task_mappings: List[Tuple[Tuple[str, ...], int]],
+    ) -> None:
+        """Extract text from message parts."""
+        for part_idx, part in enumerate(parts):
+            if part.get("kind") == "text":
+                text = part.get("text", "")
+                if text:
+                    texts_to_check.append(text)
+                    task_mappings.append((path, part_idx))
+
+    def _apply_text_to_path(
+        self,
+        result: Dict[Union[str, int], Any],
+        path: Tuple[str, ...],
+        part_idx: int,
+        text: str,
+    ) -> None:
+        """Apply guardrailed text back to the specified path in the result."""
+        # Navigate to the parts list
+        current = result
+        for key in path:
+            if key.isdigit():
+                # Array index
+                current = current[int(key)]
+            else:
+                current = current[key]
+
+        # Update the text in the part
+        current[part_idx]["text"] = text

litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py

@@ -1,11 +1,12 @@
 from typing import Dict, List, Optional, Set, Tuple
 
+from fastapi import HTTPException
 from starlette.datastructures import Headers
 from starlette.requests import Request
 from starlette.types import Scope
 
 from litellm._logging import verbose_logger
-from litellm.proxy._types import LiteLLM_TeamTable, SpecialHeaders, UserAPIKeyAuth
+from litellm.proxy._types import LiteLLM_TeamTable, ProxyException, SpecialHeaders, UserAPIKeyAuth
 from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
 
 
@@ -63,6 +64,13 @@ class MCPRequestHandler:
             HTTPException: If headers are invalid or missing required headers
         """
         headers = MCPRequestHandler._safe_get_headers_from_scope(scope)
+
+        # Check if there is an explicit LiteLLM API key (primary header)
+        has_explicit_litellm_key = (
+            headers.get(MCPRequestHandler.LITELLM_API_KEY_HEADER_NAME_PRIMARY)
+            is not None
+        )
+
         litellm_api_key = (
             MCPRequestHandler.get_litellm_api_key_from_headers(headers) or ""
         )
@@ -106,16 +114,38 @@ class MCPRequestHandler:
         request.body = mock_body  # type: ignore
         if ".well-known" in str(request.url):  # public routes
             validated_user_api_key_auth = UserAPIKeyAuth()
-        # elif litellm_api_key == "":
-        #     from fastapi import HTTPException
-
-        #     raise HTTPException(
-        #         status_code=401,
-        #         detail="LiteLLM API key is missing. Please add it or use OAuth authentication.",
-        #         headers={
-        #             "WWW-Authenticate": f'Bearer resource_metadata=f"{request.base_url}/.well-known/oauth-protected-resource"',
-        #         },
-        #     )
+        elif has_explicit_litellm_key:
+            # Explicit x-litellm-api-key provided - always validate normally
+            validated_user_api_key_auth = await user_api_key_auth(
+                api_key=litellm_api_key, request=request
+            )
+        elif oauth2_headers:
+            # No x-litellm-api-key, but Authorization header present.
+            # Could be a LiteLLM key (backward compat) OR an OAuth2 token
+            # from an upstream MCP provider (e.g. Atlassian).
+            # Try LiteLLM auth first; on auth failure, treat as OAuth2 passthrough.
+            try:
+                validated_user_api_key_auth = await user_api_key_auth(
+                    api_key=litellm_api_key, request=request
+                )
+            except HTTPException as e:
+                if e.status_code in (401, 403):
+                    verbose_logger.debug(
+                        "MCP OAuth2: Authorization header is not a valid LiteLLM key, "
+                        "treating as OAuth2 token passthrough"
+                    )
+                    validated_user_api_key_auth = UserAPIKeyAuth()
+                else:
+                    raise
+            except ProxyException as e:
+                if str(e.code) in ("401", "403"):
+                    verbose_logger.debug(
+                        "MCP OAuth2: Authorization header is not a valid LiteLLM key, "
+                        "treating as OAuth2 token passthrough"
+                    )
+                    validated_user_api_key_auth = UserAPIKeyAuth()
+                else:
+                    raise
         else:
             validated_user_api_key_auth = await user_api_key_auth(
                 api_key=litellm_api_key, request=request

litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py

@@ -1,26 +1,37 @@
 """
 MCP Guardrail Handler for Unified Guardrails.
 
-This handler works with the synthetic "messages" payload generated by
-`ProxyLogging._convert_mcp_to_llm_format`, which always produces a single user
-message whose `content` string encodes the MCP tool name and arguments. The
-handler simply feeds that text through the configured guardrail and writes the
-result back onto the message.
+Converts an MCP call_tool (name + arguments) into a single OpenAI-compatible
+tool_call and passes it to apply_guardrail. Works with the synthetic payload
+from ProxyLogging._convert_mcp_to_llm_format.
+
+Note: For MCP tool definitions (schema) -> OpenAI tools=[], see
+litellm.experimental_mcp_client.tools.transform_mcp_tool_to_openai_tool
+when you have a full MCP Tool from list_tools. Here we only have the call
+payload (name + arguments) so we just build the tool_call.
 """
 
 from typing import TYPE_CHECKING, Any, Dict, Optional
 
+from mcp.types import Tool as MCPTool
+
 from litellm._logging import verbose_proxy_logger
+from litellm.experimental_mcp_client.tools import transform_mcp_tool_to_openai_tool
 from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
+from litellm.types.llms.openai import (
+    ChatCompletionToolParam,
+    ChatCompletionToolParamFunctionChunk,
+)
 from litellm.types.utils import GenericGuardrailAPIInputs
 
 if TYPE_CHECKING:
-    from litellm.integrations.custom_guardrail import CustomGuardrail
     from mcp.types import CallToolResult
 
+    from litellm.integrations.custom_guardrail import CustomGuardrail
+
 
 class MCPGuardrailTranslationHandler(BaseTranslation):
-    """Guardrail translation handler for MCP tool calls."""
+    """Guardrail translation handler for MCP tool calls (passes a single tool_call to guardrail)."""
 
     async def process_input_messages(
         self,
@@ -28,56 +39,51 @@ class MCPGuardrailTranslationHandler(BaseTranslation):
         guardrail_to_apply: "CustomGuardrail",
         litellm_logging_obj: Optional[Any] = None,
     ) -> Dict[str, Any]:
-        messages = data.get("messages")
-        if not isinstance(messages, list) or not messages:
-            verbose_proxy_logger.debug("MCP Guardrail: No messages to process")
+        mcp_tool_name = data.get("mcp_tool_name") or data.get("name")
+        mcp_arguments = data.get("mcp_arguments") or data.get("arguments")
+        mcp_tool_description = data.get("mcp_tool_description") or data.get(
+            "description"
+        )
+        if mcp_arguments is None or not isinstance(mcp_arguments, dict):
+            mcp_arguments = {}
+
+        if not mcp_tool_name:
+            verbose_proxy_logger.debug("MCP Guardrail: mcp_tool_name missing")
             return data
 
-        first_message = messages[0]
-        content: Optional[str] = None
-        if isinstance(first_message, dict):
-            content = first_message.get("content")
-        else:
-            content = getattr(first_message, "content", None)
+        # Convert MCP input via transform_mcp_tool_to_openai_tool, then map to litellm
+        # ChatCompletionToolParam (openai SDK type has incompatible strict/cache_control).
+        mcp_tool = MCPTool(
+            name=mcp_tool_name,
+            description=mcp_tool_description or "",
+            inputSchema={},  # Call payload has no schema; guardrail gets args from request_data
+        )
+        openai_tool = transform_mcp_tool_to_openai_tool(mcp_tool)
+        fn = openai_tool["function"]
+        tool_def: ChatCompletionToolParam = {
+            "type": "function",
+            "function": ChatCompletionToolParamFunctionChunk(
+                name=fn["name"],
+                description=fn.get("description") or "",
+                parameters=fn.get("parameters")
+                or {
+                    "type": "object",
+                    "properties": {},
+                    "additionalProperties": False,
+                },
+                strict=fn.get("strict", False) or False,  # Default to False if None
+            ),
+        }
+        inputs: GenericGuardrailAPIInputs = GenericGuardrailAPIInputs(
+            tools=[tool_def],
+        )
 
-        if not isinstance(content, str):
-            verbose_proxy_logger.debug(
-                "MCP Guardrail: Message content missing or not a string",
-            )
-            return data
-
-        inputs = GenericGuardrailAPIInputs(texts=[content])
-        # Include model information if available
-        model = data.get("model")
-        if model:
-            inputs["model"] = model
-        guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
+        await guardrail_to_apply.apply_guardrail(
             inputs=inputs,
             request_data=data,
             input_type="request",
             logging_obj=litellm_logging_obj,
         )
-        guardrailed_texts = (
-            guardrailed_inputs.get("texts", []) if guardrailed_inputs else []
-        )
-
-        if guardrailed_texts:
-            new_content = guardrailed_texts[0]
-            if isinstance(first_message, dict):
-                first_message["content"] = new_content
-            else:
-                setattr(first_message, "content", new_content)
-
-            verbose_proxy_logger.debug(
-                "MCP Guardrail: Updated content for tool %s",
-                data.get("mcp_tool_name"),
-            )
-        else:
-            verbose_proxy_logger.debug(
-                "MCP Guardrail: Guardrail returned no text updates for tool %s",
-                data.get("mcp_tool_name"),
-            )
-
         return data
 
     async def process_output_response(
@@ -87,7 +93,6 @@ class MCPGuardrailTranslationHandler(BaseTranslation):
         litellm_logging_obj: Optional[Any] = None,
         user_api_key_dict: Optional[Any] = None,
     ) -> Any:
-        # Not implemented: MCP guardrail translation never calls this path today.
         verbose_proxy_logger.debug(
             "MCP Guardrail: Output processing not implemented for MCP tools",
         )

litellm/proxy/_experimental/mcp_server/mcp_server_manager.py

@@ -72,8 +72,10 @@ try:
     from mcp.shared.tool_name_validation import (
         SEP_986_URL,
     )
+    from mcp.shared.tool_name_validation import SEP_986_URL
 except ImportError:
     from pydantic import BaseModel
+
     SEP_986_URL = "https://github.com/modelcontextprotocol/protocol/blob/main/proposals/0001-tool-name-validation.md"
 
     class _ToolNameValidationResult(BaseModel):
@@ -475,12 +477,12 @@ class MCPServerManager:
                     )
 
                     # Update tool name to server name mapping (for both prefixed and base names)
-                    self.tool_name_to_mcp_server_name_mapping[
-                        base_tool_name
-                    ] = server_prefix
-                    self.tool_name_to_mcp_server_name_mapping[
-                        prefixed_tool_name
-                    ] = server_prefix
+                    self.tool_name_to_mcp_server_name_mapping[base_tool_name] = (
+                        server_prefix
+                    )
+                    self.tool_name_to_mcp_server_name_mapping[prefixed_tool_name] = (
+                        server_prefix
+                    )
 
                     registered_count += 1
                     verbose_logger.debug(
@@ -1955,7 +1957,9 @@ class MCPServerManager:
         )
 
         async def _call_tool_via_client(client, params):
-            return await client.call_tool(params, host_progress_callback=host_progress_callback)
+            return await client.call_tool(
+                params, host_progress_callback=host_progress_callback
+            )
 
         tasks.append(
             asyncio.create_task(_call_tool_via_client(client, call_tool_params))
@@ -1993,7 +1997,6 @@ class MCPServerManager:
         oauth2_headers: Optional[Dict[str, str]] = None,
         raw_headers: Optional[Dict[str, str]] = None,
         host_progress_callback: Optional[Callable] = None,
-
     ) -> CallToolResult:
         """
         Call a tool with the given name and arguments

litellm/proxy/_experimental/mcp_server/rest_endpoints.py

@@ -13,6 +13,7 @@ from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.auth.ip_address_utils import IPAddressUtils
 from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
 from litellm.types.mcp import MCPAuth
+from litellm.types.utils import CallTypes
 
 MCP_AVAILABLE: bool = True
 try:
@@ -118,6 +119,35 @@ if MCP_AVAILABLE:
 
         return _create_tool_response_objects(tools, server.mcp_info)
 
+    async def _resolve_allowed_mcp_servers_for_tool_call(
+        user_api_key_dict: UserAPIKeyAuth,
+        server_id: str,
+    ) -> List[MCPServer]:
+        """Resolve allowed MCP servers for the given user and validate server_id access."""
+        auth_contexts = await build_effective_auth_contexts(user_api_key_dict)
+        allowed_server_ids_set = set()
+        for auth_context in auth_contexts:
+            servers = await global_mcp_server_manager.get_allowed_mcp_servers(
+                user_api_key_auth=auth_context
+            )
+            allowed_server_ids_set.update(servers)
+        if server_id not in allowed_server_ids_set:
+            raise HTTPException(
+                status_code=403,
+                detail={
+                    "error": "access_denied",
+                    "message": f"The key is not allowed to access server {server_id}",
+                },
+            )
+        allowed_mcp_servers: List[MCPServer] = []
+        for allowed_server_id in allowed_server_ids_set:
+            server = global_mcp_server_manager.get_mcp_server_by_id(
+                allowed_server_id
+            )
+            if server is not None:
+                allowed_mcp_servers.append(server)
+        return allowed_mcp_servers
+
     ########################################################
     @router.get("/tools/list", dependencies=[Depends(user_api_key_auth)])
     async def list_tool_rest_api(
@@ -289,7 +319,14 @@ if MCP_AVAILABLE:
         from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
             MCPRequestHandler,
         )
-        from litellm.proxy.proxy_server import add_litellm_data_to_request, proxy_config
+        from litellm.proxy.common_request_processing import (
+            ProxyBaseLLMRequestProcessing,
+        )
+        from litellm.proxy.proxy_server import (
+            general_settings,
+            proxy_config,
+            proxy_logging_obj,
+        )
 
         try:
             data = await request.json()
@@ -317,11 +354,16 @@ if MCP_AVAILABLE:
 
             tool_arguments = data.get("arguments")
 
-            data = await add_litellm_data_to_request(
-                data=data,
-                request=request,
-                user_api_key_dict=user_api_key_dict,
-                proxy_config=proxy_config,
+            proxy_base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data)
+            data, logging_obj = (
+                await proxy_base_llm_response_processor.common_processing_pre_call_logic(
+                    request=request,
+                    user_api_key_dict=user_api_key_dict,
+                    proxy_config=proxy_config,
+                    route_type=CallTypes.call_mcp_tool.value,
+                    proxy_logging_obj=proxy_logging_obj,
+                    general_settings=general_settings,
+                )
             )
 
             # Extract MCP auth headers from request and add to data dict

litellm/proxy/_types.py

@@ -2162,6 +2162,7 @@ class LiteLLM_VerificationToken(LiteLLMPydanticObjectBase):
     rotation_interval: Optional[str] = None  # How often to rotate (e.g., "30d", "90d")
     last_rotation_at: Optional[datetime] = None  # When this key was last rotated
     key_rotation_at: Optional[datetime] = None  # When this key should next be rotated
+    router_settings: Optional[dict] = None
     model_config = ConfigDict(protected_namespaces=())
 
 

litellm/proxy/agent_endpoints/a2a_endpoints.py

@@ -14,6 +14,7 @@ from fastapi.responses import JSONResponse, StreamingResponse
 from litellm._logging import verbose_proxy_logger
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
+from litellm.types.utils import all_litellm_params
 
 router = APIRouter()
 
@@ -75,10 +76,7 @@ async def _handle_stream_message(
 
         return StreamingResponse(_error_stream(), media_type="application/x-ndjson")
 
-    from a2a.types import (
-        MessageSendParams,
-        SendStreamingMessageRequest,
-    )
+    from a2a.types import MessageSendParams, SendStreamingMessageRequest
 
     async def stream_response():
         try:
@@ -208,16 +206,17 @@ async def invoke_agent_a2a(
     from litellm.proxy.agent_endpoints.auth.agent_permission_handler import (
         AgentRequestHandler,
     )
-    from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
     from litellm.proxy.proxy_server import (
         general_settings,
         proxy_config,
+        proxy_logging_obj,
         version,
     )
 
     body = {}
     try:
         body = await request.json()
+
         verbose_proxy_logger.debug(f"A2A request for agent '{agent_id}': {body}")
 
         # Validate JSON-RPC format
@@ -230,6 +229,16 @@ async def invoke_agent_a2a(
         method = body.get("method")
         params = body.get("params", {})
 
+        if params:
+            # extract any litellm params from the params - eg. 'guardrails'
+            params_to_remove = []
+            for key, value in params.items():
+                if key in all_litellm_params:
+                    params_to_remove.append(key)
+                    body[key] = value
+            for key in params_to_remove:
+                params.pop(key)
+
         if not A2A_SDK_AVAILABLE:
             return _jsonrpc_error(
                 request_id,
@@ -283,12 +292,18 @@ async def invoke_agent_a2a(
         )
 
         # Add litellm data (user_api_key, user_id, team_id, etc.)
-        data = await add_litellm_data_to_request(
-            data=body,
+        from litellm.proxy.common_request_processing import (
+            ProxyBaseLLMRequestProcessing,
+        )
+
+        processor = ProxyBaseLLMRequestProcessing(data=body)
+        data, logging_obj = await processor.common_processing_pre_call_logic(
             request=request,
-            user_api_key_dict=user_api_key_dict,
-            proxy_config=proxy_config,
             general_settings=general_settings,
+            user_api_key_dict=user_api_key_dict,
+            proxy_logging_obj=proxy_logging_obj,
+            proxy_config=proxy_config,
+            route_type="asend_message",
             version=version,
         )
 

litellm/proxy/common_request_processing.py

@@ -267,7 +267,9 @@ def _override_openai_response_model(
     hidden_params = getattr(response_obj, "_hidden_params", {}) or {}
     if isinstance(hidden_params, dict):
         fallback_headers = hidden_params.get("additional_headers", {}) or {}
-        attempted_fallbacks = fallback_headers.get("x-litellm-attempted-fallbacks", None)
+        attempted_fallbacks = fallback_headers.get(
+            "x-litellm-attempted-fallbacks", None
+        )
         if attempted_fallbacks is not None and attempted_fallbacks > 0:
             # A fallback occurred - preserve the actual model that was used
             verbose_proxy_logger.debug(
@@ -517,6 +519,8 @@ class ProxyBaseLLMRequestProcessing:
             "aget_interaction",
             "adelete_interaction",
             "acancel_interaction",
+            "asend_message",
+            "call_mcp_tool",
         ],
         version: Optional[str] = None,
         user_model: Optional[str] = None,
@@ -616,6 +620,23 @@ class ProxyBaseLLMRequestProcessing:
             user_api_key_dict=user_api_key_dict, data=self.data, call_type=route_type  # type: ignore
         )
 
+        # Apply hierarchical router_settings (Key > Team)
+        # Global router_settings are already on the Router object itself.
+        if llm_router is not None and proxy_config is not None:
+            from litellm.proxy.proxy_server import prisma_client
+
+            router_settings = await proxy_config._get_hierarchical_router_settings(
+                user_api_key_dict=user_api_key_dict,
+                prisma_client=prisma_client,
+                proxy_logging_obj=proxy_logging_obj,
+            )
+
+            # If router_settings found (from key or team), apply them
+            # Pass settings as per-request overrides instead of creating a new Router
+            # This avoids expensive Router instantiation on each request
+            if router_settings is not None:
+                self.data["router_settings_override"] = router_settings
+
         if "messages" in self.data and self.data["messages"]:
             logging_obj.update_messages(self.data["messages"])
 
@@ -820,7 +841,9 @@ class ProxyBaseLLMRequestProcessing:
             # aliasing/routing, but the OpenAI-compatible response `model` field should reflect
             # what the client sent.
             if requested_model_from_client:
-                self.data["_litellm_client_requested_model"] = requested_model_from_client
+                self.data["_litellm_client_requested_model"] = (
+                    requested_model_from_client
+                )
             if route_type == "allm_passthrough_route":
                 # Check if response is an async generator
                 if self._is_streaming_response(response):
@@ -1392,9 +1415,9 @@ class ProxyBaseLLMRequestProcessing:
 
             # Add cache-related fields to **params (handled by Usage.__init__)
             if cache_creation_input_tokens is not None:
-                usage_kwargs[
-                    "cache_creation_input_tokens"
-                ] = cache_creation_input_tokens
+                usage_kwargs["cache_creation_input_tokens"] = (
+                    cache_creation_input_tokens
+                )
             if cache_read_input_tokens is not None:
                 usage_kwargs["cache_read_input_tokens"] = cache_read_input_tokens
 

litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py

@@ -5,7 +5,7 @@ This module provides a guardrail that executes user-defined Python-like code
 to implement custom guardrail logic. The code runs in a sandboxed environment
 with access to LiteLLM-provided primitives for common guardrail operations.
 
-Example custom code:
+Example custom code (sync):
 
     def apply_guardrail(inputs, request_data, input_type):
         '''Block messages containing SSNs'''
@@ -13,8 +13,22 @@ Example custom code:
             if regex_match(text, r"\\d{3}-\\d{2}-\\d{4}"):
                 return block("Social Security Number detected")
         return allow()
+
+Example custom code (async with HTTP):
+
+    async def apply_guardrail(inputs, request_data, input_type):
+        '''Call external moderation API'''
+        for text in inputs["texts"]:
+            response = await http_post(
+                "https://api.example.com/moderate",
+                body={"text": text}
+            )
+            if response["success"] and response["body"].get("flagged"):
+                return block("Content flagged by moderation API")
+        return allow()
 """
 
+import asyncio
 import threading
 from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Type, cast
 
@@ -101,6 +115,9 @@ class CustomCodeGuardrail(CustomGuardrail):
             GuardrailEventHooks.pre_call,
             GuardrailEventHooks.during_call,
             GuardrailEventHooks.post_call,
+            GuardrailEventHooks.pre_mcp_call,
+            GuardrailEventHooks.during_mcp_call,
+            GuardrailEventHooks.logging_only,
         ]
 
         super().__init__(
@@ -175,6 +192,13 @@ class CustomCodeGuardrail(CustomGuardrail):
         This method calls the user-defined apply_guardrail function and
         processes its result to determine the appropriate action.
 
+        The user-defined function can be either sync or async:
+        - Sync: def apply_guardrail(inputs, request_data, input_type): ...
+        - Async: async def apply_guardrail(inputs, request_data, input_type): ...
+
+        Async functions are recommended when using http_request, http_get, or
+        http_post primitives to avoid blocking the event loop.
+
         Args:
             inputs: Dictionary containing texts, images, tool_calls
             request_data: The original request data with metadata
@@ -188,6 +212,7 @@ class CustomCodeGuardrail(CustomGuardrail):
             HTTPException: If content is blocked
             CustomCodeExecutionError: If execution fails
         """
+
         if self._compiled_function is None:
             if self._compile_error:
                 raise CustomCodeExecutionError(
@@ -201,9 +226,13 @@ class CustomCodeGuardrail(CustomGuardrail):
             # Prepare request_data with safe subset of information
             safe_request_data = self._prepare_safe_request_data(request_data)
 
-            # Execute the custom function
+            # Execute the custom function - handle both sync and async functions
             result = self._compiled_function(inputs, safe_request_data, input_type)
 
+            # If the function is async (returns a coroutine), await it
+            if asyncio.iscoroutine(result):
+                result = await result
+
             # Process the result
             return self._process_result(
                 result=result,

litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py

@@ -10,7 +10,11 @@ import re
 from typing import Any, Dict, List, Optional, Tuple, Type, Union
 from urllib.parse import urlparse
 
+import httpx
+
 from litellm._logging import verbose_proxy_logger
+from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
+from litellm.types.llms.custom_http import httpxSpecialProvider
 
 # =============================================================================
 # Result Types - Used by Starlark code to return guardrail decisions
@@ -349,6 +353,228 @@ def get_url_domain(url: str) -> Optional[str]:
         return None
 
 
+# =============================================================================
+# HTTP Request Primitives (Async)
+# =============================================================================
+
+# Default timeout for HTTP requests (in seconds)
+_HTTP_DEFAULT_TIMEOUT = 30.0
+
+# Maximum allowed timeout (in seconds)
+_HTTP_MAX_TIMEOUT = 60.0
+
+
+def _http_error_response(error: str) -> Dict[str, Any]:
+    """Create a standardized error response for HTTP requests."""
+    return {
+        "status_code": 0,
+        "body": None,
+        "headers": {},
+        "success": False,
+        "error": error,
+    }
+
+
+def _http_success_response(response: httpx.Response) -> Dict[str, Any]:
+    """Create a standardized success response from an httpx Response."""
+    parsed_body: Any
+    try:
+        parsed_body = response.json()
+    except (json.JSONDecodeError, ValueError):
+        parsed_body = response.text
+
+    return {
+        "status_code": response.status_code,
+        "body": parsed_body,
+        "headers": dict(response.headers),
+        "success": 200 <= response.status_code < 300,
+        "error": None,
+    }
+
+
+def _prepare_http_body(
+    body: Optional[Any],
+) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
+    """Prepare body arguments for HTTP request - returns (json_body, data_body)."""
+    if body is None:
+        return None, None
+    if isinstance(body, dict):
+        return body, None
+    if isinstance(body, list):
+        return None, json.dumps(body)
+    if isinstance(body, str):
+        return None, body
+    return None, str(body)
+
+
+async def http_request(
+    url: str,
+    method: str = "GET",
+    headers: Optional[Dict[str, str]] = None,
+    body: Optional[Any] = None,
+    timeout: Optional[float] = None,
+) -> Dict[str, Any]:
+    """
+    Make an async HTTP request to an external service.
+
+    This function allows custom guardrails to call external APIs for
+    additional validation, content moderation, or data enrichment.
+
+    Uses LiteLLM's global cached AsyncHTTPHandler for connection pooling
+    and better performance.
+
+    Args:
+        url: The URL to request
+        method: HTTP method (GET, POST, PUT, DELETE, PATCH). Defaults to GET.
+        headers: Optional dict of HTTP headers
+        body: Optional request body (will be JSON-encoded if dict/list)
+        timeout: Optional timeout in seconds (default: 30, max: 60)
+
+    Returns:
+        Dict containing:
+            - status_code: HTTP status code
+            - body: Response body (parsed as JSON if possible, otherwise string)
+            - headers: Response headers as dict
+            - success: True if status code is 2xx
+            - error: Error message if request failed, None otherwise
+
+    Example:
+        # Simple GET request
+        response = await http_request("https://api.example.com/check")
+        if response["success"]:
+            data = response["body"]
+
+        # POST request with JSON body
+        response = await http_request(
+            "https://api.example.com/moderate",
+            method="POST",
+            headers={"Authorization": "Bearer token"},
+            body={"text": "content to check"}
+        )
+    """
+    # Validate URL
+    if not is_valid_url(url):
+        return _http_error_response(f"Invalid URL: {url}")
+
+    # Validate and normalize method
+    method = method.upper()
+    allowed_methods = {"GET", "POST", "PUT", "DELETE", "PATCH"}
+    if method not in allowed_methods:
+        return _http_error_response(
+            f"Invalid HTTP method: {method}. Allowed: {', '.join(allowed_methods)}"
+        )
+
+    # Apply timeout limits
+    if timeout is None:
+        timeout = _HTTP_DEFAULT_TIMEOUT
+    else:
+        timeout = min(max(0.1, timeout), _HTTP_MAX_TIMEOUT)
+
+    # Get the global cached async HTTP client
+    client = get_async_httpx_client(
+        llm_provider=httpxSpecialProvider.GuardrailCallback,
+        params={"timeout": httpx.Timeout(timeout=timeout, connect=5.0)},
+    )
+
+    try:
+        response = await _execute_http_request(
+            client, method, url, headers, body, timeout
+        )
+        return _http_success_response(response)
+
+    except httpx.TimeoutException as e:
+        verbose_proxy_logger.warning(f"Custom code http_request timeout: {e}")
+        return _http_error_response(f"Request timeout after {timeout}s")
+    except httpx.HTTPStatusError as e:
+        # Return the response even for non-2xx status codes
+        return _http_success_response(e.response)
+    except httpx.RequestError as e:
+        verbose_proxy_logger.warning(f"Custom code http_request error: {e}")
+        return _http_error_response(f"Request failed: {str(e)}")
+    except Exception as e:
+        verbose_proxy_logger.warning(f"Custom code http_request unexpected error: {e}")
+        return _http_error_response(f"Unexpected error: {str(e)}")
+
+
+async def _execute_http_request(
+    client: Any,
+    method: str,
+    url: str,
+    headers: Optional[Dict[str, str]],
+    body: Optional[Any],
+    timeout: float,
+) -> httpx.Response:
+    """Execute the HTTP request using the appropriate client method."""
+    json_body, data_body = _prepare_http_body(body)
+
+    if method == "GET":
+        return await client.get(url=url, headers=headers)
+    elif method == "POST":
+        return await client.post(
+            url=url, headers=headers, json=json_body, data=data_body, timeout=timeout
+        )
+    elif method == "PUT":
+        return await client.put(
+            url=url, headers=headers, json=json_body, data=data_body, timeout=timeout
+        )
+    elif method == "DELETE":
+        return await client.delete(
+            url=url, headers=headers, json=json_body, data=data_body, timeout=timeout
+        )
+    elif method == "PATCH":
+        return await client.patch(
+            url=url, headers=headers, json=json_body, data=data_body, timeout=timeout
+        )
+    else:
+        raise ValueError(f"Unsupported HTTP method: {method}")
+
+
+async def http_get(
+    url: str,
+    headers: Optional[Dict[str, str]] = None,
+    timeout: Optional[float] = None,
+) -> Dict[str, Any]:
+    """
+    Make an async HTTP GET request.
+
+    Convenience wrapper around http_request for GET requests.
+
+    Args:
+        url: The URL to request
+        headers: Optional dict of HTTP headers
+        timeout: Optional timeout in seconds
+
+    Returns:
+        Same as http_request
+    """
+    return await http_request(url=url, method="GET", headers=headers, timeout=timeout)
+
+
+async def http_post(
+    url: str,
+    body: Optional[Any] = None,
+    headers: Optional[Dict[str, str]] = None,
+    timeout: Optional[float] = None,
+) -> Dict[str, Any]:
+    """
+    Make an async HTTP POST request.
+
+    Convenience wrapper around http_request for POST requests.
+
+    Args:
+        url: The URL to request
+        body: Optional request body (will be JSON-encoded if dict/list)
+        headers: Optional dict of HTTP headers
+        timeout: Optional timeout in seconds
+
+    Returns:
+        Same as http_request
+    """
+    return await http_request(
+        url=url, method="POST", headers=headers, body=body, timeout=timeout
+    )
+
+
 # =============================================================================
 # Code Detection Primitives
 # =============================================================================
@@ -575,6 +801,10 @@ def get_custom_code_primitives() -> Dict[str, Any]:
         "is_valid_url": is_valid_url,
         "all_urls_valid": all_urls_valid,
         "get_url_domain": get_url_domain,
+        # HTTP (async)
+        "http_request": http_request,
+        "http_get": http_get,
+        "http_post": http_post,
         # Code detection
         "detect_code": detect_code,
         "detect_code_languages": detect_code_languages,

litellm/proxy/guardrails/guardrail_hooks/pangea/pangea.py

@@ -1,6 +1,6 @@
 # litellm/proxy/guardrails/guardrail_hooks/pangea.py
 import os
-from typing import TYPE_CHECKING, Any, Optional, Type
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type, cast
 
 from fastapi import HTTPException
 
@@ -250,10 +250,13 @@ class PangeaHandler(CustomGuardrail):
         if isinstance(response, TextCompletionResponse):
             # Assume the earlier call type as well
             input_messages = _TextCompletionRequest(data).get_messages()
-        if not isinstance(response, ModelResponse):
-            return
+        elif isinstance(response, ModelResponse):
+            messages = data.get("messages")
+            if messages is None:
+                return  # No messages to check
+            input_messages = cast(List[Dict[Any, Any]], messages)
         else:
-            input_messages = data.get("messages")
+            return
 
         if choices := response.get("choices"):
             if isinstance(choices, list):

litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py

@@ -51,6 +51,7 @@ class UnifiedLLMGuardrails(CustomLogger):
         Runs on only Input
         Use this if you want to MODIFY the input
         """
+
         global endpoint_guardrail_translation_mappings
         from litellm.proxy.common_utils.callback_utils import (
             add_guardrail_to_applied_guardrails_header,
@@ -66,6 +67,7 @@ class UnifiedLLMGuardrails(CustomLogger):
         if call_type == CallTypes.call_mcp_tool.value:
             event_type = GuardrailEventHooks.pre_mcp_call
 
+
         if (
             guardrail_to_apply.should_run_guardrail(data=data, event_type=event_type)
             is not True

litellm/proxy/management_endpoints/types.py

@@ -28,17 +28,24 @@ def is_valid_litellm_user_role(role_str: str) -> bool:
         return False
 
 
-def get_litellm_user_role(role_str: str) -> Optional[LitellmUserRoles]:
+def get_litellm_user_role(role_str) -> Optional[LitellmUserRoles]:
     """
-    Convert a string to a LitellmUserRoles enum if valid (case-insensitive).
+    Convert a string (or list of strings) to a LitellmUserRoles enum if valid (case-insensitive).
+
+    Handles list inputs since some SSO providers (e.g., Keycloak) return roles
+    as arrays like ["proxy_admin"] instead of plain strings.
 
     Args:
-        role_str: String to convert (e.g., "proxy_admin", "PROXY_ADMIN", "internal_user")
+        role_str: String or list to convert (e.g., "proxy_admin", ["proxy_admin"])
 
     Returns:
         LitellmUserRoles enum if valid, None otherwise
     """
     try:
+        if isinstance(role_str, list):
+            if len(role_str) == 0:
+                return None
+            role_str = role_str[0]
         # Use _value2member_map_ for O(1) lookup, case-insensitive
         result = LitellmUserRoles._value2member_map_.get(role_str.lower())
         return cast(Optional[LitellmUserRoles], result)

litellm/proxy/management_endpoints/ui_sso.py

@@ -171,36 +171,86 @@ def process_sso_jwt_access_token(
     access_token_str: Optional[str],
     sso_jwt_handler: Optional[JWTHandler],
     result: Union[OpenID, dict, None],
+    role_mappings: Optional["RoleMappings"] = None,
 ) -> None:
     """
-    Process SSO JWT access token and extract team IDs if available.
+    Process SSO JWT access token and extract team IDs and user role if available.
 
-    This function decodes the JWT access token and extracts team IDs using the
-    sso_jwt_handler, then sets the team_ids attribute on the result object.
+    This function decodes the JWT access token and extracts team IDs and user
+    role, then sets them on the result object. Role extraction from the access
+    token is needed because some SSO providers (e.g., Keycloak) do not include
+    role claims in the UserInfo endpoint response.
 
     Args:
         access_token_str: The JWT access token string
         sso_jwt_handler: SSO-specific JWT handler for team ID extraction
-        result: The SSO result object to update with team IDs
+        result: The SSO result object to update with team IDs and role
+        role_mappings: Optional role mappings configuration for group-based role determination
     """
-    if access_token_str and sso_jwt_handler and result:
+    if access_token_str and result:
         import jwt
 
         access_token_payload = jwt.decode(
             access_token_str, options={"verify_signature": False}
         )
 
-        # Handle both dict and object result types
-        if isinstance(result, dict):
-            result_team_ids: Optional[List[str]] = result.get("team_ids", [])
-            if not result_team_ids:
-                team_ids = sso_jwt_handler.get_team_ids_from_jwt(access_token_payload)
-                result["team_ids"] = team_ids
-        else:
-            result_team_ids = getattr(result, "team_ids", []) if result else []
-            if not result_team_ids:
-                team_ids = sso_jwt_handler.get_team_ids_from_jwt(access_token_payload)
-                setattr(result, "team_ids", team_ids)
+        # Extract team IDs from access token if sso_jwt_handler is available
+        if sso_jwt_handler:
+            if isinstance(result, dict):
+                result_team_ids: Optional[List[str]] = result.get("team_ids", [])
+                if not result_team_ids:
+                    team_ids = sso_jwt_handler.get_team_ids_from_jwt(access_token_payload)
+                    result["team_ids"] = team_ids
+            else:
+                result_team_ids = getattr(result, "team_ids", []) if result else []
+                if not result_team_ids:
+                    team_ids = sso_jwt_handler.get_team_ids_from_jwt(access_token_payload)
+                    setattr(result, "team_ids", team_ids)
+
+        # Extract user role from access token if not already set from UserInfo
+        existing_role = result.get("user_role") if isinstance(result, dict) else getattr(result, "user_role", None)
+        if existing_role is None:
+            user_role: Optional[LitellmUserRoles] = None
+
+            # Try role_mappings first (group-based role determination)
+            if role_mappings is not None and role_mappings.roles:
+                group_claim = role_mappings.group_claim
+                user_groups_raw: Any = get_nested_value(access_token_payload, group_claim)
+
+                user_groups: List[str] = []
+                if isinstance(user_groups_raw, list):
+                    user_groups = [str(g) for g in user_groups_raw]
+                elif isinstance(user_groups_raw, str):
+                    user_groups = [g.strip() for g in user_groups_raw.split(",") if g.strip()]
+                elif user_groups_raw is not None:
+                    user_groups = [str(user_groups_raw)]
+
+                if user_groups:
+                    user_role = determine_role_from_groups(user_groups, role_mappings)
+                    verbose_proxy_logger.debug(
+                        f"Determined role '{user_role}' from access token groups '{user_groups}' using role_mappings"
+                    )
+                elif role_mappings.default_role:
+                    user_role = role_mappings.default_role
+
+            # Fallback: try GENERIC_USER_ROLE_ATTRIBUTE on the access token payload
+            if user_role is None:
+                generic_user_role_attribute_name = os.getenv("GENERIC_USER_ROLE_ATTRIBUTE", "role")
+                user_role_from_token = get_nested_value(access_token_payload, generic_user_role_attribute_name)
+                if user_role_from_token is not None:
+                    user_role = get_litellm_user_role(user_role_from_token)
+                    verbose_proxy_logger.debug(
+                        f"Extracted role '{user_role}' from access token field '{generic_user_role_attribute_name}'"
+                    )
+
+            if user_role is not None:
+                if isinstance(result, dict):
+                    result["user_role"] = user_role
+                else:
+                    setattr(result, "user_role", user_role)
+                verbose_proxy_logger.debug(
+                    f"Set user_role='{user_role}' from JWT access token"
+                )
 
 
 @router.get("/sso/key/generate", tags=["experimental"], include_in_schema=False)
@@ -688,7 +738,7 @@ async def get_generic_sso_response(
         )
 
         access_token_str: Optional[str] = generic_sso.access_token
-        process_sso_jwt_access_token(access_token_str, sso_jwt_handler, result)
+        process_sso_jwt_access_token(access_token_str, sso_jwt_handler, result, role_mappings=role_mappings)
 
     except Exception as e:
         verbose_proxy_logger.exception(

litellm/proxy/proxy_server.py

@@ -3249,17 +3249,18 @@ class ProxyConfig:
                 _model_list: list = self.decrypt_model_list_from_db(
                     new_models=models_list
                 )
-                if len(_model_list) > 0:
-                    verbose_proxy_logger.debug(f"_model_list: {_model_list}")
-                    llm_router = litellm.Router(
-                        model_list=_model_list,
-                        router_general_settings=RouterGeneralSettings(
-                            async_only_mode=True  # only init async clients
-                        ),
-                        search_tools=search_tools,
-                        ignore_invalid_deployments=True,
-                    )
-                    verbose_proxy_logger.debug(f"updated llm_router: {llm_router}")
+                # Create router even with empty model list to support search_tools
+                # Router can function with model_list=[] and only search_tools
+                verbose_proxy_logger.debug(f"_model_list: {_model_list}")
+                llm_router = litellm.Router(
+                    model_list=_model_list,
+                    router_general_settings=RouterGeneralSettings(
+                        async_only_mode=True  # only init async clients
+                    ),
+                    search_tools=search_tools,
+                    ignore_invalid_deployments=True,
+                )
+                verbose_proxy_logger.debug(f"updated llm_router: {llm_router}")
             else:
                 verbose_proxy_logger.debug(f"len new_models: {len(models_list)}")
                 if search_tools is not None and llm_router is not None:
@@ -3402,6 +3403,86 @@ class ProxyConfig:
             decrypted_variables[k] = decrypted_value
         return decrypted_variables
 
+    @staticmethod
+    def _parse_router_settings_value(value: Any) -> Optional[dict]:
+        """
+        Parse a router_settings value that may be a dict or a JSON/YAML string.
+
+        Returns a non-empty dict if valid, otherwise None.
+        """
+        if value is None:
+            return None
+
+        parsed: Optional[dict] = None
+        if isinstance(value, dict):
+            parsed = value
+        elif isinstance(value, str):
+            import json
+
+            import yaml
+
+            try:
+                parsed = yaml.safe_load(value)
+            except (yaml.YAMLError, json.JSONDecodeError):
+                try:
+                    parsed = json.loads(value)
+                except json.JSONDecodeError:
+                    pass
+
+        if isinstance(parsed, dict) and parsed:
+            return parsed
+        return None
+
+    async def _get_hierarchical_router_settings(
+        self,
+        user_api_key_dict: Optional["UserAPIKeyAuth"],
+        prisma_client: Optional[PrismaClient],
+        proxy_logging_obj: Optional["ProxyLogging"] = None,
+    ) -> Optional[dict]:
+        """
+        Get router_settings in priority order: Key > Team
+
+        Uses the already-cached key object and the cached team lookup
+        (get_team_object) to avoid direct DB queries on the hot path.
+
+        Global router_settings are NOT looked up here — they are already
+        applied to the Router object at config-load / DB-sync time.
+
+        Returns:
+            dict: router_settings, or None if no settings found
+        """
+        # 1. Try key-level router_settings
+        # user_api_key_dict is already the cached/authenticated key object —
+        # no DB call needed.
+        if user_api_key_dict is not None:
+            key_settings = self._parse_router_settings_value(
+                getattr(user_api_key_dict, "router_settings", None)
+            )
+            if key_settings is not None:
+                return key_settings
+
+        # 2. Try team-level router_settings using cached team lookup
+        # get_team_object checks in-memory cache / Redis first, only falls
+        # back to DB on a cache miss.
+        if user_api_key_dict is not None and user_api_key_dict.team_id is not None:
+            try:
+                team_obj = await get_team_object(
+                    team_id=user_api_key_dict.team_id,
+                    prisma_client=prisma_client,
+                    user_api_key_cache=user_api_key_cache,
+                    proxy_logging_obj=proxy_logging_obj,
+                )
+                team_settings = self._parse_router_settings_value(
+                    getattr(team_obj, "router_settings", None)
+                )
+                if team_settings is not None:
+                    return team_settings
+            except Exception:
+                # If team lookup fails, no team-level settings available
+                pass
+
+        return None
+
     async def _add_router_settings_from_db_config(
         self,
         config_data: dict,

litellm/proxy/route_llm_request.py

@@ -210,18 +210,33 @@ async def route_request(
             models = [model.strip() for model in data.pop("model").split(",")]
             return llm_router.abatch_completion(models=models, **data)
 
-    elif "user_config" in data:
-        router_config = data.pop("user_config")
+    elif "router_settings_override" in data:
+        # Apply per-request router settings overrides from key/team config
+        # Instead of creating a new Router (expensive), merge settings into kwargs
+        # The Router already supports per-request overrides for these settings
+        override_settings = data.pop("router_settings_override")
 
-        # Filter router_config to only include valid Router.__init__ arguments
-        # This prevents TypeError when invalid parameters are stored in the database
-        valid_args = litellm.Router.get_valid_args()
-        filtered_config = {k: v for k, v in router_config.items() if k in valid_args}
+        # Settings that the Router accepts as per-request kwargs
+        # These override the global router settings for this specific request
+        per_request_settings = [
+            "fallbacks",
+            "context_window_fallbacks",
+            "content_policy_fallbacks",
+            "num_retries",
+            "timeout",
+            "model_group_retry_policy",
+        ]
 
-        user_router = litellm.Router(**filtered_config)
-        ret_val = getattr(user_router, f"{route_type}")(**data)
-        user_router.discard()
-        return ret_val
+        # Merge override settings into data (only if not already set in request)
+        for key in per_request_settings:
+            if key in override_settings and key not in data:
+                data[key] = override_settings[key]
+
+        # Use main router with overridden kwargs
+        if llm_router is not None:
+            return getattr(llm_router, f"{route_type}")(**data)
+        else:
+            return getattr(litellm, f"{route_type}")(**data)
     elif llm_router is not None:
         # Skip model-based routing for container operations
         if route_type in [

litellm/responses/litellm_completion_transformation/handler.py

@@ -90,6 +90,7 @@ class LiteLLMCompletionTransformationHandler:
                 custom_llm_provider=custom_llm_provider,
                 litellm_metadata=kwargs.get("litellm_metadata", {}),
             )
+        raise ValueError(f"Unexpected response type: {type(litellm_completion_response)}")
 
     async def async_response_api_handler(
         self,
@@ -140,3 +141,4 @@ class LiteLLMCompletionTransformationHandler:
                 ),
                 litellm_metadata=kwargs.get("litellm_metadata", {}),
             )
+        raise ValueError(f"Unexpected response type: {type(litellm_completion_response)}")

litellm/router.py

@@ -4893,6 +4893,10 @@ class Router:
         content_policy_fallbacks = kwargs.pop(
             "content_policy_fallbacks", self.content_policy_fallbacks
         )
+        # Support per-request model_group_retry_policy override (from key/team settings)
+        model_group_retry_policy = kwargs.pop(
+            "model_group_retry_policy", self.model_group_retry_policy
+        )
         model_group: Optional[str] = kwargs.get("model")
         num_retries = kwargs.pop("num_retries")
 
@@ -4941,7 +4945,7 @@ class Router:
             _retry_policy_applies = False
             if (
                 self.retry_policy is not None
-                or self.model_group_retry_policy is not None
+                or model_group_retry_policy is not None
             ):
                 # get num_retries from retry policy
                 # Use the model_group captured at the start of the function, or get it from metadata
@@ -4949,9 +4953,12 @@ class Router:
                 _model_group_for_retry_policy = (
                     model_group or _metadata.get("model_group") or kwargs.get("model")
                 )
-                _retry_policy_retries = self.get_num_retries_from_retry_policy(
+                # Use per-request model_group_retry_policy if provided, otherwise use self
+                _retry_policy_retries = _get_num_retries_from_retry_policy(
                     exception=original_exception,
                     model_group=_model_group_for_retry_policy,
+                    model_group_retry_policy=model_group_retry_policy,
+                    retry_policy=self.retry_policy,
                 )
                 if _retry_policy_retries is not None:
                     num_retries = _retry_policy_retries

model_prices_and_context_window.json

@@ -1263,6 +1263,36 @@
         "supports_vision": true,
         "tool_use_system_prompt_tokens": 346
     },
+    "au.anthropic.claude-opus-4-6-v1:0": {
+        "cache_creation_input_token_cost": 6.875e-06,
+        "cache_creation_input_token_cost_above_200k_tokens": 1.375e-05,
+        "cache_read_input_token_cost": 5.5e-07,
+        "cache_read_input_token_cost_above_200k_tokens": 1.1e-06,
+        "input_cost_per_token": 5.5e-06,
+        "input_cost_per_token_above_200k_tokens": 1.1e-05,
+        "litellm_provider": "bedrock_converse",
+        "max_input_tokens": 200000,
+        "max_output_tokens": 128000,
+        "max_tokens": 128000,
+        "mode": "chat",
+        "output_cost_per_token": 2.75e-05,
+        "output_cost_per_token_above_200k_tokens": 4.125e-05,
+        "search_context_cost_per_query": {
+            "search_context_size_high": 0.01,
+            "search_context_size_low": 0.01,
+            "search_context_size_medium": 0.01
+        },
+        "supports_assistant_prefill": false,
+        "supports_computer_use": true,
+        "supports_function_calling": true,
+        "supports_pdf_input": true,
+        "supports_prompt_caching": true,
+        "supports_reasoning": true,
+        "supports_response_schema": true,
+        "supports_tool_choice": true,
+        "supports_vision": true,
+        "tool_use_system_prompt_tokens": 346
+    },
     "anthropic.claude-sonnet-4-20250514-v1:0": {
         "cache_creation_input_token_cost": 3.75e-06,
         "cache_read_input_token_cost": 3e-07,

tests/agent_tests/test_a2a_agent.py

@@ -0,0 +1,83 @@
+"""
+Simple A2A agent tests - non-streaming and streaming.
+
+These tests validate the localhost URL retry logic: if an A2A agent's card
+contains a localhost/internal URL (e.g., http://0.0.0.0:8001/), the request
+will fail with a connection error. LiteLLM detects this and automatically
+retries using the original api_base URL instead.
+
+Requires A2A_AGENT_URL environment variable to be set.
+
+Run with:
+    A2A_AGENT_URL=https://your-agent.example.com pytest tests/agent_tests/test_a2a_agent.py -v -s
+"""
+
+import os
+
+import pytest
+from uuid import uuid4
+
+
+def get_a2a_agent_url():
+    """Get A2A agent URL from environment, skip test if not set."""
+    url = os.environ.get("A2A_AGENT_URL")
+    return url
+
+
+@pytest.mark.asyncio
+async def test_a2a_non_streaming():
+    """Test non-streaming A2A request."""
+    from a2a.types import MessageSendParams, SendMessageRequest
+    from litellm.a2a_protocol import asend_message
+
+    api_base = get_a2a_agent_url()
+
+    request = SendMessageRequest(
+        id=str(uuid4()),
+        params=MessageSendParams(
+            message={
+                "role": "user",
+                "parts": [{"kind": "text", "text": "Say hello in one word"}],
+                "messageId": uuid4().hex,
+            }
+        ),
+    )
+
+    response = await asend_message(
+        request=request,
+        api_base=api_base,
+    )
+
+    assert response is not None
+    print(f"\nNon-streaming response: {response}")
+
+
+@pytest.mark.asyncio
+async def test_a2a_streaming():
+    """Test streaming A2A request."""
+    from a2a.types import MessageSendParams, SendStreamingMessageRequest
+    from litellm.a2a_protocol import asend_message_streaming
+
+    api_base = get_a2a_agent_url()
+
+    request = SendStreamingMessageRequest(
+        id=str(uuid4()),
+        params=MessageSendParams(
+            message={
+                "role": "user",
+                "parts": [{"kind": "text", "text": "Say hello in one word"}],
+                "messageId": uuid4().hex,
+            }
+        ),
+    )
+
+    chunks = []
+    async for chunk in asend_message_streaming(
+        request=request,
+        api_base=api_base,
+    ):
+        chunks.append(chunk)
+        print(f"\nStreaming chunk: {chunk}")
+
+    assert len(chunks) > 0, "Should receive at least one chunk"
+    print(f"\nTotal chunks received: {len(chunks)}")

tests/mcp_tests/test_mcp_client_unit.py

@@ -77,6 +77,27 @@ class TestMCPClientUnitTests:
             "Authorization": "Token custom_token",
         }
 
+        # OAuth2
+        client = MCPClient(
+            "http://example.com",
+            auth_type=MCPAuth.oauth2,
+            auth_value="oauth2-access-token-xyz",
+        )
+        headers = client._get_auth_headers()
+        assert headers == {
+            "Authorization": "Bearer oauth2-access-token-xyz",
+        }
+
+        # OAuth2 with extra_headers (per-user flow overrides auth_value)
+        client = MCPClient(
+            "http://example.com",
+            auth_type=MCPAuth.oauth2,
+            auth_value="static-server-token",
+            extra_headers={"Authorization": "Bearer per-user-token"},
+        )
+        headers = client._get_auth_headers()
+        assert headers["Authorization"] == "Bearer per-user-token"
+
         # No auth
         client = MCPClient("http://example.com")
         headers = client._get_auth_headers()

tests/test_litellm/a2a_protocol/test_card_resolver.py

@@ -4,11 +4,16 @@ Mock tests for LiteLLMA2ACardResolver.
 Tests that the card resolver tries both old and new well-known paths.
 """
 
-import sys
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 
+from litellm.a2a_protocol.card_resolver import (
+    LiteLLMA2ACardResolver,
+    fix_agent_card_url,
+    is_localhost_or_internal_url,
+)
+
 
 @pytest.mark.asyncio
 async def test_card_resolver_fallback_from_new_to_old_path():
@@ -24,39 +29,31 @@ async def test_card_resolver_fallback_from_new_to_old_path():
     # Track which paths were called
     paths_called = []
 
-    # Create a mock base class
-    class MockA2ACardResolver:
-        def __init__(self, base_url):
-            self.base_url = base_url
-
-        async def get_agent_card(self, relative_card_path=None, http_kwargs=None):
-            paths_called.append(relative_card_path)
-            if relative_card_path == "/.well-known/agent-card.json":
-                # First call (new path) fails
-                raise Exception("404 Not Found")
-            else:
-                # Second call (old path) succeeds
-                return mock_agent_card
-
-    # Create mock A2A module
-    mock_a2a_module = MagicMock()
-    mock_a2a_client = MagicMock()
-    mock_a2a_constants = MagicMock()
-    mock_a2a_constants.AGENT_CARD_WELL_KNOWN_PATH = "/.well-known/agent-card.json"
-    mock_a2a_constants.PREV_AGENT_CARD_WELL_KNOWN_PATH = "/.well-known/agent.json"
-
-    with patch.dict(
-        sys.modules,
-        {
-            "a2a": mock_a2a_module,
-            "a2a.client": MagicMock(A2ACardResolver=MockA2ACardResolver),
-            "a2a.utils.constants": mock_a2a_constants,
-        },
+    # Create a mock for the parent's get_agent_card method
+    async def mock_parent_get_agent_card(
+        self, relative_card_path=None, http_kwargs=None
     ):
-        # Import after patching
-        from litellm.a2a_protocol.card_resolver import LiteLLMA2ACardResolver
+        paths_called.append(relative_card_path)
+        if relative_card_path == "/.well-known/agent-card.json":
+            # First call (new path) fails
+            raise Exception("404 Not Found")
+        else:
+            # Second call (old path) succeeds
+            return mock_agent_card
 
-        resolver = LiteLLMA2ACardResolver(base_url="http://test-agent:8000")
+    # Create a mock httpx client
+    mock_httpx_client = MagicMock()
+
+    # Patch the parent class's get_agent_card method
+    # We need to patch the actual parent class method that super() calls
+    with patch.object(
+        LiteLLMA2ACardResolver.__bases__[0],
+        "get_agent_card",
+        mock_parent_get_agent_card,
+    ):
+        resolver = LiteLLMA2ACardResolver(
+            httpx_client=mock_httpx_client, base_url="http://test-agent:8000"
+        )
         result = await resolver.get_agent_card()
 
         # Verify both paths were tried in correct order
@@ -67,3 +64,27 @@ async def test_card_resolver_fallback_from_new_to_old_path():
         # Verify the result
         assert result == mock_agent_card
         assert result.name == "Test Agent"
+
+
+def test_is_localhost_or_internal_url():
+    """Test that localhost/internal URLs are correctly detected."""
+    # Should return True for localhost variants
+    assert is_localhost_or_internal_url("http://localhost:8000/") is True
+    assert is_localhost_or_internal_url("http://0.0.0.0:8001/") is True
+
+    # Should return False for public URLs
+    assert is_localhost_or_internal_url("https://my-agent.example.com/") is False
+    assert is_localhost_or_internal_url(None) is False
+
+
+def test_fix_agent_card_url_replaces_localhost():
+    """Test that fix_agent_card_url replaces localhost URLs with base_url."""
+    # Create mock agent card with localhost URL
+    mock_card = MagicMock()
+    mock_card.url = "http://0.0.0.0:8001/"
+
+    # Fix the URL
+    result = fix_agent_card_url(mock_card, "https://my-public-agent.example.com")
+
+    # Verify localhost URL was replaced with base_url
+    assert result.url == "https://my-public-agent.example.com/"

tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py

@@ -544,6 +544,235 @@ class TestMCPRequestHandler:
             assert mcp_server_auth_headers == {}
 
 
+@pytest.mark.asyncio
+class TestMCPOAuth2AuthFlow:
+    """Test suite for OAuth2 authentication flow in MCP requests.
+
+    Tests the fix for the 'Capabilities: none' bug where OAuth2 tokens
+    from upstream MCP providers (e.g., Atlassian) were mistakenly validated
+    as LiteLLM API keys, causing auth failures and empty tool listings.
+    """
+
+    async def test_oauth2_token_in_authorization_header_fallback(self):
+        """
+        When only Authorization header is present with a non-LiteLLM OAuth2 token,
+        auth should fall back to permissive mode (OAuth2 passthrough).
+        """
+        from fastapi import HTTPException
+
+        scope = {
+            "type": "http",
+            "method": "POST",
+            "path": "/mcp/atlassian_mcp",
+            "headers": [
+                (b"authorization", b"Bearer atlassian-oauth2-access-token-xyz"),
+            ],
+        }
+
+        async def mock_user_api_key_auth_fails(api_key, request):
+            raise HTTPException(status_code=401, detail="Invalid API key")
+
+        with patch(
+            "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth",
+            side_effect=mock_user_api_key_auth_fails,
+        ):
+            (
+                auth_result,
+                mcp_auth_header,
+                mcp_servers,
+                mcp_server_auth_headers,
+                oauth2_headers,
+                raw_headers,
+            ) = await MCPRequestHandler.process_mcp_request(scope)
+
+            # Should succeed with default UserAPIKeyAuth (OAuth2 fallback)
+            assert auth_result is not None
+            assert isinstance(auth_result, UserAPIKeyAuth)
+            # OAuth2 headers should contain the token for upstream forwarding
+            assert (
+                oauth2_headers.get("Authorization")
+                == "Bearer atlassian-oauth2-access-token-xyz"
+            )
+
+    async def test_explicit_litellm_key_with_oauth2_authorization(self):
+        """
+        When both x-litellm-api-key AND Authorization header are present,
+        LiteLLM key should be used for auth and Authorization preserved for OAuth2.
+        """
+        scope = {
+            "type": "http",
+            "method": "POST",
+            "path": "/mcp/atlassian_mcp",
+            "headers": [
+                (b"x-litellm-api-key", b"sk-litellm-valid-key"),
+                (b"authorization", b"Bearer atlassian-oauth2-token"),
+            ],
+        }
+
+        async def mock_user_api_key_auth(api_key, request):
+            return UserAPIKeyAuth(api_key=api_key, user_id="test-user")
+
+        with patch(
+            "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth",
+            side_effect=mock_user_api_key_auth,
+        ) as mock_auth:
+            (
+                auth_result,
+                mcp_auth_header,
+                mcp_servers,
+                mcp_server_auth_headers,
+                oauth2_headers,
+                raw_headers,
+            ) = await MCPRequestHandler.process_mcp_request(scope)
+
+            # LiteLLM key should be used for auth
+            mock_auth.assert_called_once()
+            call_args = mock_auth.call_args
+            assert call_args.kwargs["api_key"] == "sk-litellm-valid-key"
+
+            # OAuth2 headers should still contain the Authorization token
+            assert (
+                oauth2_headers.get("Authorization")
+                == "Bearer atlassian-oauth2-token"
+            )
+
+    async def test_litellm_key_in_authorization_backward_compat(self):
+        """
+        Backward compatibility: when only Authorization header is present
+        with a valid LiteLLM key (not OAuth2), auth should succeed normally.
+        """
+        scope = {
+            "type": "http",
+            "method": "POST",
+            "path": "/mcp/some_server",
+            "headers": [
+                (b"authorization", b"Bearer sk-litellm-valid-key"),
+            ],
+        }
+
+        async def mock_user_api_key_auth(api_key, request):
+            return UserAPIKeyAuth(api_key=api_key, user_id="test-user")
+
+        with patch(
+            "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth",
+            side_effect=mock_user_api_key_auth,
+        ) as mock_auth:
+            (
+                auth_result,
+                _,
+                _,
+                _,
+                _,
+                _,
+            ) = await MCPRequestHandler.process_mcp_request(scope)
+
+            # Should succeed with the LiteLLM key from Authorization header
+            assert auth_result.api_key == "Bearer sk-litellm-valid-key"
+            mock_auth.assert_called_once()
+
+    async def test_non_auth_http_exception_still_raises(self):
+        """
+        If user_api_key_auth raises a non-401/403 HTTPException (e.g., 500),
+        it should NOT be caught by the OAuth2 fallback.
+        """
+        from fastapi import HTTPException
+
+        scope = {
+            "type": "http",
+            "method": "POST",
+            "path": "/mcp/some_server",
+            "headers": [
+                (b"authorization", b"Bearer some-token"),
+            ],
+        }
+
+        async def mock_user_api_key_auth_server_error(api_key, request):
+            raise HTTPException(status_code=500, detail="Internal server error")
+
+        with patch(
+            "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth",
+            side_effect=mock_user_api_key_auth_server_error,
+        ):
+            with pytest.raises(HTTPException) as exc_info:
+                await MCPRequestHandler.process_mcp_request(scope)
+            assert exc_info.value.status_code == 500
+
+    async def test_proxy_exception_oauth2_fallback(self):
+        """
+        user_api_key_auth raises ProxyException (not HTTPException) in production.
+        The OAuth2 fallback must catch ProxyException with code 401/403 too.
+        """
+        from litellm.proxy._types import ProxyException
+
+        scope = {
+            "type": "http",
+            "method": "POST",
+            "path": "/mcp/atlassian_mcp",
+            "headers": [
+                (b"authorization", b"Bearer atlassian-oauth2-access-token-xyz"),
+            ],
+        }
+
+        async def mock_user_api_key_auth_proxy_exception(api_key, request):
+            raise ProxyException(
+                message="Authentication Error: Invalid API key",
+                type="auth_error",
+                param="api_key",
+                code=401,
+            )
+
+        with patch(
+            "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth",
+            side_effect=mock_user_api_key_auth_proxy_exception,
+        ):
+            (
+                auth_result,
+                mcp_auth_header,
+                mcp_servers,
+                mcp_server_auth_headers,
+                oauth2_headers,
+                raw_headers,
+            ) = await MCPRequestHandler.process_mcp_request(scope)
+
+            # Should succeed with default UserAPIKeyAuth (OAuth2 fallback)
+            assert auth_result is not None
+            assert isinstance(auth_result, UserAPIKeyAuth)
+            assert (
+                oauth2_headers.get("Authorization")
+                == "Bearer atlassian-oauth2-access-token-xyz"
+            )
+
+    async def test_proxy_exception_non_auth_still_raises(self):
+        """
+        ProxyException with non-401/403 code should NOT be caught.
+        """
+        from litellm.proxy._types import ProxyException
+
+        scope = {
+            "type": "http",
+            "method": "POST",
+            "path": "/mcp/some_server",
+            "headers": [
+                (b"authorization", b"Bearer some-token"),
+            ],
+        }
+
+        async def mock_user_api_key_auth_500(api_key, request):
+            raise ProxyException(
+                message="Internal error",
+                type="server_error",
+                param=None,
+                code=500,
+            )
+
+        with patch(
+            "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth",
+            side_effect=mock_user_api_key_auth_500,
+        ):
+            with pytest.raises(ProxyException):
+                await MCPRequestHandler.process_mcp_request(scope)
+
+
 class TestMCPCustomHeaderName:
     """Test suite for custom MCP authentication header name functionality"""
 

tests/test_litellm/proxy/management_endpoints/test_ui_sso.py

@@ -25,6 +25,8 @@ from litellm.proxy.management_endpoints.ui_sso import (
     MicrosoftSSOHandler,
     SSOAuthenticationHandler,
     normalize_email,
+    process_sso_jwt_access_token,
+    determine_role_from_groups,
     _setup_team_mappings,
 )
 from litellm.types.proxy.management_endpoints.ui_sso import (
@@ -1298,6 +1300,7 @@ async def test_get_generic_sso_response_with_additional_headers():
     # Mock the SSO provider and its methods
     mock_sso_instance = MagicMock()
     mock_sso_instance.verify_and_process = AsyncMock(return_value=mock_sso_response)
+    mock_sso_instance.access_token = None  # Avoid triggering JWT decode in process_sso_jwt_access_token
 
     mock_sso_class = MagicMock(return_value=mock_sso_instance)
 
@@ -1359,6 +1362,7 @@ async def test_get_generic_sso_response_with_empty_headers():
     # Mock the SSO provider and its methods
     mock_sso_instance = MagicMock()
     mock_sso_instance.verify_and_process = AsyncMock(return_value=mock_sso_response)
+    mock_sso_instance.access_token = None  # Avoid triggering JWT decode in process_sso_jwt_access_token
 
     mock_sso_class = MagicMock(return_value=mock_sso_instance)
 
@@ -2546,22 +2550,25 @@ class TestProcessSSOJWTAccessToken:
             assert result.team_ids == []
 
     def test_process_sso_jwt_access_token_no_sso_jwt_handler(self, sample_jwt_token):
-        """Test that nothing happens when sso_jwt_handler is None"""
+        """Test that JWT is decoded for role extraction even when sso_jwt_handler is None,
+        but team_ids are not extracted (team extraction requires sso_jwt_handler)."""
         from litellm.proxy.management_endpoints.ui_sso import (
             process_sso_jwt_access_token,
         )
 
         result = CustomOpenID(id="test_user", email="test@example.com", team_ids=[])
 
-        with patch("jwt.decode") as mock_jwt_decode:
+        mock_payload = {"sub": "test_user", "email": "test@example.com"}
+        with patch("jwt.decode", return_value=mock_payload) as mock_jwt_decode:
             # Act
             process_sso_jwt_access_token(
                 access_token_str=sample_jwt_token, sso_jwt_handler=None, result=result
             )
 
-            # Assert nothing was processed
-            mock_jwt_decode.assert_not_called()
+            # JWT is decoded (for role extraction) but team_ids are not extracted
+            mock_jwt_decode.assert_called_once()
             assert result.team_ids == []
+            assert result.user_role is None
 
     def test_process_sso_jwt_access_token_no_result(
         self, mock_jwt_handler, sample_jwt_token
@@ -3848,3 +3855,219 @@ async def test_setup_team_mappings():
         mock_prisma.db.litellm_ssoconfig.find_unique.assert_called_once_with(
             where={"id": "sso_config"}
         )
+
+
+# ============================================================================
+# Tests for get_litellm_user_role with list inputs (Keycloak returns lists)
+# ============================================================================
+
+
+def test_get_litellm_user_role_with_string():
+    """Test that get_litellm_user_role works with a plain string."""
+    from litellm.proxy._types import LitellmUserRoles
+    from litellm.proxy.management_endpoints.types import get_litellm_user_role
+
+    result = get_litellm_user_role("proxy_admin")
+    assert result == LitellmUserRoles.PROXY_ADMIN
+
+
+def test_get_litellm_user_role_with_list():
+    """
+    Test that get_litellm_user_role handles list inputs.
+    Keycloak returns roles as arrays like ["proxy_admin"] instead of strings.
+    """
+    from litellm.proxy._types import LitellmUserRoles
+    from litellm.proxy.management_endpoints.types import get_litellm_user_role
+
+    result = get_litellm_user_role(["proxy_admin"])
+    assert result == LitellmUserRoles.PROXY_ADMIN
+
+
+def test_get_litellm_user_role_with_empty_list():
+    """Test that get_litellm_user_role returns None for empty lists."""
+    from litellm.proxy.management_endpoints.types import get_litellm_user_role
+
+    result = get_litellm_user_role([])
+    assert result is None
+
+
+def test_get_litellm_user_role_with_invalid_role():
+    """Test that get_litellm_user_role returns None for invalid roles."""
+    from litellm.proxy.management_endpoints.types import get_litellm_user_role
+
+    result = get_litellm_user_role("not_a_real_role")
+    assert result is None
+
+
+def test_get_litellm_user_role_with_list_multiple_roles():
+    """Test that get_litellm_user_role takes the first element from a multi-element list."""
+    from litellm.proxy._types import LitellmUserRoles
+    from litellm.proxy.management_endpoints.types import get_litellm_user_role
+
+    result = get_litellm_user_role(["proxy_admin", "internal_user"])
+    assert result == LitellmUserRoles.PROXY_ADMIN
+
+
+# ============================================================================
+# Tests for process_sso_jwt_access_token role extraction
+# ============================================================================
+
+
+def test_process_sso_jwt_access_token_extracts_role_from_access_token():
+    """
+    Test that process_sso_jwt_access_token extracts user role from the JWT
+    access token when the UserInfo response did not include it.
+
+    This is the core fix for the Keycloak SSO role mapping bug: Keycloak's
+    UserInfo endpoint does not return role claims, but the JWT access token
+    contains them.
+    """
+    import jwt as pyjwt
+
+    from litellm.proxy._types import LitellmUserRoles
+
+    # Create a JWT access token with role claims (as Keycloak would)
+    access_token_payload = {
+        "sub": "user-123",
+        "email": "admin@test.com",
+        "litellm_role": ["proxy_admin"],
+    }
+    access_token_str = pyjwt.encode(access_token_payload, "secret", algorithm="HS256")
+
+    # Result object with no role set (simulating UserInfo response without roles)
+    result = CustomOpenID(
+        id="user-123",
+        email="admin@test.com",
+        display_name="Admin User",
+        team_ids=[],
+        user_role=None,
+    )
+
+    # Call with GENERIC_USER_ROLE_ATTRIBUTE pointing to litellm_role
+    with patch.dict(os.environ, {"GENERIC_USER_ROLE_ATTRIBUTE": "litellm_role"}):
+        process_sso_jwt_access_token(
+            access_token_str=access_token_str,
+            sso_jwt_handler=None,
+            result=result,
+            role_mappings=None,
+        )
+
+    assert result.user_role == LitellmUserRoles.PROXY_ADMIN
+
+
+def test_process_sso_jwt_access_token_does_not_override_existing_role():
+    """
+    Test that process_sso_jwt_access_token does NOT override a role that was
+    already extracted from the UserInfo response.
+    """
+    import jwt as pyjwt
+
+    from litellm.proxy._types import LitellmUserRoles
+
+    access_token_payload = {
+        "sub": "user-123",
+        "litellm_role": ["internal_user"],
+    }
+    access_token_str = pyjwt.encode(access_token_payload, "secret", algorithm="HS256")
+
+    # Result already has a role (e.g., set from UserInfo)
+    result = CustomOpenID(
+        id="user-123",
+        email="admin@test.com",
+        display_name="Admin User",
+        team_ids=[],
+        user_role=LitellmUserRoles.PROXY_ADMIN,
+    )
+
+    with patch.dict(os.environ, {"GENERIC_USER_ROLE_ATTRIBUTE": "litellm_role"}):
+        process_sso_jwt_access_token(
+            access_token_str=access_token_str,
+            sso_jwt_handler=None,
+            result=result,
+            role_mappings=None,
+        )
+
+    # Should keep the original role
+    assert result.user_role == LitellmUserRoles.PROXY_ADMIN
+
+
+def test_process_sso_jwt_access_token_extracts_role_from_nested_field():
+    """
+    Test role extraction from a nested JWT field like resource_access.client.roles.
+    """
+    import jwt as pyjwt
+
+    from litellm.proxy._types import LitellmUserRoles
+
+    access_token_payload = {
+        "sub": "user-123",
+        "resource_access": {
+            "my-client": {
+                "roles": ["proxy_admin"]
+            }
+        },
+    }
+    access_token_str = pyjwt.encode(access_token_payload, "secret", algorithm="HS256")
+
+    result = CustomOpenID(
+        id="user-123",
+        email="admin@test.com",
+        display_name="Admin User",
+        team_ids=[],
+        user_role=None,
+    )
+
+    with patch.dict(os.environ, {"GENERIC_USER_ROLE_ATTRIBUTE": "resource_access.my-client.roles"}):
+        process_sso_jwt_access_token(
+            access_token_str=access_token_str,
+            sso_jwt_handler=None,
+            result=result,
+            role_mappings=None,
+        )
+
+    assert result.user_role == LitellmUserRoles.PROXY_ADMIN
+
+
+def test_process_sso_jwt_access_token_with_role_mappings():
+    """
+    Test role extraction using role_mappings (group-based role determination)
+    from the JWT access token.
+    """
+    import jwt as pyjwt
+
+    from litellm.proxy._types import LitellmUserRoles
+    from litellm.types.proxy.management_endpoints.ui_sso import RoleMappings
+
+    access_token_payload = {
+        "sub": "user-123",
+        "groups": ["keycloak-admins", "developers"],
+    }
+    access_token_str = pyjwt.encode(access_token_payload, "secret", algorithm="HS256")
+
+    result = CustomOpenID(
+        id="user-123",
+        email="admin@test.com",
+        display_name="Admin User",
+        team_ids=[],
+        user_role=None,
+    )
+
+    role_mappings = RoleMappings(
+        provider="generic",
+        group_claim="groups",
+        default_role=LitellmUserRoles.INTERNAL_USER,
+        roles={
+            LitellmUserRoles.PROXY_ADMIN: ["keycloak-admins"],
+            LitellmUserRoles.INTERNAL_USER: ["developers"],
+        },
+    )
+
+    process_sso_jwt_access_token(
+        access_token_str=access_token_str,
+        sso_jwt_handler=None,
+        result=result,
+        role_mappings=role_mappings,
+    )
+
+    # Should get highest privilege role
+    assert result.user_role == LitellmUserRoles.PROXY_ADMIN

tests/test_litellm/proxy/test_common_request_processing.py

@@ -77,6 +77,93 @@ class TestProxyBaseLLMRequestProcessing:
             pytest.fail("litellm_call_id is not a valid UUID")
         assert data_passed["litellm_call_id"] == returned_data["litellm_call_id"]
 
+    @pytest.mark.asyncio
+    async def test_should_apply_hierarchical_router_settings_as_override(
+        self, monkeypatch
+    ):
+        """
+        Test that hierarchical router settings are stored as router_settings_override
+        instead of creating a full user_config with model_list.
+        
+        This approach avoids expensive per-request Router instantiation by passing
+        settings as kwargs overrides to the main router.
+        """
+        processing_obj = ProxyBaseLLMRequestProcessing(data={})
+        mock_request = MagicMock(spec=Request)
+        mock_request.headers = {}
+
+        async def mock_add_litellm_data_to_request(*args, **kwargs):
+            return {}
+
+        async def mock_common_processing_pre_call_logic(
+            user_api_key_dict, data, call_type
+        ):
+            data_copy = copy.deepcopy(data)
+            return data_copy
+
+        mock_proxy_logging_obj = MagicMock(spec=ProxyLogging)
+        mock_proxy_logging_obj.pre_call_hook = AsyncMock(
+            side_effect=mock_common_processing_pre_call_logic
+        )
+        monkeypatch.setattr(
+            litellm.proxy.common_request_processing,
+            "add_litellm_data_to_request",
+            mock_add_litellm_data_to_request,
+        )
+
+        mock_general_settings = {}
+        mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth)
+        mock_proxy_config = MagicMock(spec=ProxyConfig)
+        
+        mock_router_settings = {
+            "routing_strategy": "least-busy",
+            "timeout": 30.0,
+            "num_retries": 3,
+        }
+        mock_proxy_config._get_hierarchical_router_settings = AsyncMock(
+            return_value=mock_router_settings
+        )
+
+        mock_llm_router = MagicMock()
+
+        mock_prisma_client = MagicMock()
+        monkeypatch.setattr(
+            "litellm.proxy.proxy_server.prisma_client",
+            mock_prisma_client,
+        )
+
+        route_type = "acompletion"
+
+        returned_data, logging_obj = await processing_obj.common_processing_pre_call_logic(
+            request=mock_request,
+            general_settings=mock_general_settings,
+            user_api_key_dict=mock_user_api_key_dict,
+            proxy_logging_obj=mock_proxy_logging_obj,
+            proxy_config=mock_proxy_config,
+            route_type=route_type,
+            llm_router=mock_llm_router,
+        )
+
+        mock_proxy_config._get_hierarchical_router_settings.assert_called_once_with(
+            user_api_key_dict=mock_user_api_key_dict,
+            prisma_client=mock_prisma_client,
+            proxy_logging_obj=mock_proxy_logging_obj,
+        )
+        # get_model_list should NOT be called - we no longer copy model list for per-request routers
+        mock_llm_router.get_model_list.assert_not_called()
+
+        # Settings should be stored as router_settings_override (not user_config)
+        # This allows passing them as kwargs to the main router instead of creating a new one
+        assert "router_settings_override" in returned_data
+        assert "user_config" not in returned_data
+        
+        router_settings_override = returned_data["router_settings_override"]
+        assert router_settings_override["routing_strategy"] == "least-busy"
+        assert router_settings_override["timeout"] == 30.0
+        assert router_settings_override["num_retries"] == 3
+        # model_list should NOT be in the override settings
+        assert "model_list" not in router_settings_override
+
     @pytest.mark.asyncio
     async def test_stream_timeout_header_processing(self):
         """

tests/test_litellm/proxy/test_route_llm_request.py

@@ -137,62 +137,103 @@ async def test_route_request_no_model_required_with_router_settings_and_no_route
 
 
 @pytest.mark.asyncio
-async def test_route_request_with_invalid_router_params():
+async def test_route_request_with_router_settings_override():
     """
-    Test that route_request filters out invalid Router init params from 'user_config'.
-    This covers the fix for https://github.com/BerriAI/litellm/issues/19693
+    Test that route_request handles router_settings_override by merging settings into kwargs
+    instead of creating a new Router (which is expensive and was the old behavior).
     """
-    import litellm
-    from litellm.router import Router
-    from unittest.mock import AsyncMock
-
-    # Mock data with user_config containing invalid keys (simulating DB entry)
+    # Mock data with router_settings_override containing per-request settings
     data = {
         "model": "gpt-3.5-turbo",
-        "user_config": {
-            "model_list": [
-                {
-                    "model_name": "gpt-3.5-turbo",
-                    "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "test"},
-                }
-            ],
-            "model_alias_map": {"alias": "real_model"},  # INVALID PARAM
-            "invalid_garbage_key": "crash_me",  # INVALID PARAM
+        "messages": [{"role": "user", "content": "Hello"}],
+        "router_settings_override": {
+            "fallbacks": [{"gpt-3.5-turbo": ["gpt-4"]}],
+            "num_retries": 5,
+            "timeout": 30,
+            "model_group_retry_policy": {"gpt-3.5-turbo": {"RateLimitErrorRetries": 3}},
+            # These settings should be ignored (not in per_request_settings list)
+            "routing_strategy": "least-busy",
+            "model_group_alias": {"alias": "real_model"},
         },
     }
 
-    # We expect Router(**config) to succeed because of the filtering.
-    # If filtering fails, this will raise TypeError and fail the test.
+    llm_router = MagicMock()
+    llm_router.acompletion.return_value = "success"
+
+    response = await route_request(data, llm_router, None, "acompletion")
+
+    assert response == "success"
+    # Verify the router method was called with merged settings
+    call_kwargs = llm_router.acompletion.call_args[1]
+    assert call_kwargs["fallbacks"] == [{"gpt-3.5-turbo": ["gpt-4"]}]
+    assert call_kwargs["num_retries"] == 5
+    assert call_kwargs["timeout"] == 30
+    assert call_kwargs["model_group_retry_policy"] == {"gpt-3.5-turbo": {"RateLimitErrorRetries": 3}}
+    # Verify unsupported settings were NOT merged
+    assert "routing_strategy" not in call_kwargs
+    assert "model_group_alias" not in call_kwargs
+    # Verify router_settings_override was removed from data
+    assert "router_settings_override" not in call_kwargs
+
+
+@pytest.mark.asyncio
+async def test_route_request_with_router_settings_override_no_router():
+    """
+    Test that router_settings_override works when no router is provided,
+    falling back to litellm module directly.
+    """
+    import litellm
+
+    data = {
+        "model": "gpt-3.5-turbo",
+        "messages": [{"role": "user", "content": "Hello"}],
+        "router_settings_override": {
+            "fallbacks": [{"gpt-3.5-turbo": ["gpt-4"]}],
+            "num_retries": 3,
+        },
+    }
+
+    # Use MagicMock explicitly to avoid auto-AsyncMock behavior in Python 3.12+
+    mock_completion = MagicMock(return_value="success")
+    original_acompletion = litellm.acompletion
+    litellm.acompletion = mock_completion
+
     try:
-        # route_request calls getattr(user_router, route_type)(**data)
-        # We'll mock the internal call to avoid making real network requests
-        with pytest.MonkeyPatch.context() as m:
-            # Mock the method that gets called on the router instance
-            # We don't easily have access to the instance created INSIDE existing route_request
-            # So we will wrap litellm.Router to spy on it or verify it doesn't crash
+        response = await route_request(data, None, None, "acompletion")
 
-            original_router_init = litellm.Router.__init__
+        assert response == "success"
+        # Verify litellm.acompletion was called with merged settings
+        call_kwargs = mock_completion.call_args[1]
+        assert call_kwargs["fallbacks"] == [{"gpt-3.5-turbo": ["gpt-4"]}]
+        assert call_kwargs["num_retries"] == 3
+    finally:
+        litellm.acompletion = original_acompletion
 
-            def safe_router_init(self, **kwargs):
-                # Verify that invalid keys are NOT present in kwargs
-                assert "model_alias_map" not in kwargs
-                assert "invalid_garbage_key" not in kwargs
-                # Call original init (which would raise TypeError if invalid keys were present)
-                original_router_init(self, **kwargs)
 
-            m.setattr(litellm.Router, "__init__", safe_router_init)
+@pytest.mark.asyncio
+async def test_route_request_with_router_settings_override_preserves_existing():
+    """
+    Test that router_settings_override does not override settings already in the request.
+    Request-level settings take precedence over key/team settings.
+    """
+    data = {
+        "model": "gpt-3.5-turbo",
+        "messages": [{"role": "user", "content": "Hello"}],
+        "num_retries": 10,  # Request-level setting
+        "router_settings_override": {
+            "num_retries": 3,  # Key/team setting - should NOT override
+            "timeout": 30,  # Key/team setting - should be applied
+        },
+    }
 
-            # Use 'acompletion' as the route_type
-            # We also need to mock the completion method to avoid real calls
-            m.setattr(Router, "acompletion", AsyncMock(return_value="success"))
+    llm_router = MagicMock()
+    llm_router.acompletion.return_value = "success"
 
-            response = await route_request(data, None, None, "acompletion")
-            assert response == "success"
+    response = await route_request(data, llm_router, None, "acompletion")
 
-    except TypeError as e:
-        pytest.fail(
-            f"route_request raised TypeError, implying invalid params were passed to Router: {e}"
-        )
-    except Exception:
-        # Other exceptions might happen (e.g. valid config issues) but we care about TypeError here
-        pass
+    assert response == "success"
+    call_kwargs = llm_router.acompletion.call_args[1]
+    # Request-level num_retries should take precedence
+    assert call_kwargs["num_retries"] == 10
+    # Key/team timeout should be applied since not in request
+    assert call_kwargs["timeout"] == 30

ui/litellm-dashboard/src/app/layout.tsx

@@ -2,6 +2,8 @@ import type { Metadata } from "next";
 import { Inter } from "next/font/google";
 import "./globals.css";
 
+import AntdGlobalProvider from "@/contexts/AntdGlobalProvider";
+
 const inter = Inter({ subsets: ["latin"] });
 
 export const metadata: Metadata = {
@@ -17,7 +19,9 @@ export default function RootLayout({
 }>) {
   return (
     <html lang="en">
-      <body className={inter.className}>{children}</body>
+      <body className={inter.className}>
+        <AntdGlobalProvider>{children}</AntdGlobalProvider>
+      </body>
     </html>
   );
 }

ui/litellm-dashboard/src/components/guardrails/custom_code/CustomCodeModal.tsx

@@ -9,7 +9,7 @@ import {
   CaretRightOutlined,
   SaveOutlined,
 } from "@ant-design/icons";
-import { createGuardrailCall, testCustomCodeGuardrail } from "../../networking";
+import { createGuardrailCall, updateGuardrailCall, testCustomCodeGuardrail } from "../../networking";
 import NotificationsManager from "../../molecules/notifications_manager";
 
 const { Panel } = Collapse;
@@ -19,7 +19,7 @@ const { TextArea } = Input;
 const CODE_TEMPLATES = {
   empty: {
     name: "Empty Template",
-    code: `def apply_guardrail(inputs, request_data, input_type):
+    code: `async def apply_guardrail(inputs, request_data, input_type):
     # inputs: {texts, images, tools, tool_calls, structured_messages, model}
     # request_data: {model, user_id, team_id, end_user_id, metadata}
     # input_type: "request" or "response"
@@ -68,6 +68,27 @@ const CODE_TEMPLATES = {
             return block("Response missing required fields")
     return allow()`,
   },
+  externalAPI: {
+    name: "External API Check (async)",
+    code: `async def apply_guardrail(inputs, request_data, input_type):
+    # Call an external moderation API (async for non-blocking)
+    for text in inputs["texts"]:
+        response = await http_post(
+            "https://api.example.com/moderate",
+            body={"text": text, "user_id": request_data["user_id"]},
+            headers={"Authorization": "Bearer YOUR_API_KEY"},
+            timeout=10
+        )
+        
+        if not response["success"]:
+            # API call failed, allow by default or block
+            return allow()
+        
+        if response["body"].get("flagged"):
+            return block(response["body"].get("reason", "Content flagged"))
+    
+    return allow()`,
+  },
 };
 
 // Available primitives organized by category
@@ -77,6 +98,11 @@ const PRIMITIVES = {
     { name: "block(reason)", desc: "Reject with message" },
     { name: "modify(texts=[], images=[], tool_calls=[])", desc: "Transform content" },
   ],
+  "HTTP Requests (async)": [
+    { name: "await http_request(url, method, headers, body)", desc: "Make async HTTP request" },
+    { name: "await http_get(url, headers)", desc: "Async GET request" },
+    { name: "await http_post(url, body, headers)", desc: "Async POST request" },
+  ],
   "Regex Functions": [
     { name: "regex_match(text, pattern)", desc: "Returns True if pattern found" },
     { name: "regex_replace(text, pattern, replacement)", desc: "Replace all matches" },
@@ -111,13 +137,30 @@ const MODE_OPTIONS = [
   { value: "post_call", label: "post_call (Response)" },
   { value: "during_call", label: "during_call (Parallel)" },
   { value: "logging_only", label: "logging_only" },
+  { value: "pre_mcp_call", label: "pre_mcp_call (Before MCP Tool Call)" },
+  { value: "post_mcp_call", label: "post_mcp_call (After MCP Tool Call)" },
+  { value: "during_mcp_call", label: "during_mcp_call (During MCP Tool Call)" },
 ];
 
+// Data for editing an existing guardrail
+export interface EditGuardrailData {
+  guardrail_id: string;
+  guardrail_name: string;
+  litellm_params: {
+    mode?: string | string[];
+    default_on?: boolean;
+    custom_code?: string;
+    [key: string]: any;
+  };
+}
+
 interface CustomCodeModalProps {
   visible: boolean;
   onClose: () => void;
   onSuccess: () => void;
   accessToken: string | null;
+  /** If provided, the modal will be in edit mode */
+  editData?: EditGuardrailData | null;
 }
 
 const CustomCodeModal: React.FC<CustomCodeModalProps> = ({
@@ -125,16 +168,72 @@ const CustomCodeModal: React.FC<CustomCodeModalProps> = ({
   onClose,
   onSuccess,
   accessToken,
+  editData,
 }) => {
+  const isEditMode = !!editData;
   const [guardrailName, setGuardrailName] = useState("");
-  const [mode, setMode] = useState<string>("pre_call");
+  const [mode, setMode] = useState<string[]>(["pre_call"]);
   const [defaultOn, setDefaultOn] = useState(false);
   const [selectedTemplate, setSelectedTemplate] = useState<string>("empty");
   const [code, setCode] = useState(CODE_TEMPLATES.empty.code);
   const [isSaving, setIsSaving] = useState(false);
   const [isTesting, setIsTesting] = useState(false);
   const [testExpanded, setTestExpanded] = useState(false);
-  const [testInput, setTestInput] = useState('{"texts": ["Hello, my SSN is 123-45-6789"], "images": [], "tools": [], "tool_calls": [], "structured_messages": [], "model": "gpt-4"}');
+  
+  // Test input examples for pre_call and post_call
+  const TEST_INPUT_EXAMPLES = {
+    pre_call: {
+      name: "Pre-call (Request)",
+      data: {
+        texts: ["Hello, my SSN is 123-45-6789"],
+        images: [],
+        tools: [
+          {
+            type: "function",
+            function: {
+              name: "get_weather",
+              description: "Get the current weather in a location",
+              parameters: {
+                type: "object",
+                properties: {
+                  location: { type: "string", description: "City name" }
+                },
+                required: ["location"]
+              }
+            }
+          }
+        ],
+        tool_calls: [],
+        structured_messages: [
+          { role: "system", content: "You are a helpful assistant." },
+          { role: "user", content: "Hello, my SSN is 123-45-6789" }
+        ],
+        model: "gpt-4"
+      }
+    },
+    post_call: {
+      name: "Post-call (Response)",
+      data: {
+        texts: ["The weather in San Francisco is 72°F and sunny."],
+        images: [],
+        tools: [],
+        tool_calls: [
+          {
+            id: "call_abc123",
+            type: "function",
+            function: {
+              name: "get_weather",
+              arguments: "{\"location\": \"San Francisco\"}"
+            }
+          }
+        ],
+        structured_messages: [],
+        model: "gpt-4"
+      }
+    }
+  };
+  
+  const [testInput, setTestInput] = useState(JSON.stringify(TEST_INPUT_EXAMPLES.pre_call.data, null, 2));
   const [testResult, setTestResult] = useState<any>(null);
   const [copiedPrimitive, setCopiedPrimitive] = useState<string | null>(null);
   const textareaRef = useRef<HTMLTextAreaElement>(null);
@@ -145,18 +244,35 @@ const CustomCodeModal: React.FC<CustomCodeModalProps> = ({
     setCode(CODE_TEMPLATES[templateKey as keyof typeof CODE_TEMPLATES].code);
   };
 
-  // Reset form when modal opens
+  // Normalize mode from API (string or string[]) to string[]
+  const normalizeMode = (m: string | string[] | undefined): string[] => {
+    if (m === undefined || m === null) return ["pre_call"];
+    if (Array.isArray(m)) return m.length ? m : ["pre_call"];
+    return [m];
+  };
+
+  // Reset form when modal opens or editData changes
   useEffect(() => {
     if (visible) {
-      setGuardrailName("");
-      setMode("pre_call");
-      setDefaultOn(false);
-      setSelectedTemplate("empty");
-      setCode(CODE_TEMPLATES.empty.code);
+      if (editData) {
+        // Edit mode: populate with existing data
+        setGuardrailName(editData.guardrail_name || "");
+        setMode(normalizeMode(editData.litellm_params?.mode));
+        setDefaultOn(editData.litellm_params?.default_on || false);
+        setCode(editData.litellm_params?.custom_code || CODE_TEMPLATES.empty.code);
+        setSelectedTemplate(""); // No template selected in edit mode
+      } else {
+        // Create mode: reset to defaults
+        setGuardrailName("");
+        setMode(["pre_call"]);
+        setDefaultOn(false);
+        setSelectedTemplate("empty");
+        setCode(CODE_TEMPLATES.empty.code);
+      }
       setTestResult(null);
       setTestExpanded(false);
     }
-  }, [visible]);
+  }, [visible, editData]);
 
   // Copy primitive to clipboard
   const copyPrimitive = async (primitive: string) => {
@@ -184,7 +300,7 @@ const CustomCodeModal: React.FC<CustomCodeModalProps> = ({
     }
   };
 
-  // Save guardrail
+  // Save guardrail (create or update)
   const handleSave = async () => {
     if (!guardrailName.trim()) {
       NotificationsManager.fromBackend("Please enter a guardrail name");
@@ -201,25 +317,53 @@ const CustomCodeModal: React.FC<CustomCodeModalProps> = ({
 
     setIsSaving(true);
     try {
-      const guardrailData = {
-        guardrail_name: guardrailName,
-        litellm_params: {
-          guardrail: "custom_code",
-          mode: mode,
-          default_on: defaultOn,
-          custom_code: code,
-        },
-        guardrail_info: {},
-      };
+      if (isEditMode && editData) {
+        // Update existing guardrail
+        const updateData: any = {
+          litellm_params: {
+            custom_code: code,
+          },
+        };
 
-      await createGuardrailCall(accessToken, guardrailData);
-      NotificationsManager.success("Custom code guardrail created successfully");
+        // Only include changed fields
+        if (guardrailName !== editData.guardrail_name) {
+          updateData.guardrail_name = guardrailName;
+        }
+        const existingMode = normalizeMode(editData.litellm_params?.mode);
+        const modeChanged =
+          mode.length !== existingMode.length ||
+          mode.some((m, i) => m !== existingMode[i]);
+        if (modeChanged) {
+          updateData.litellm_params.mode = mode;
+        }
+        if (defaultOn !== editData.litellm_params?.default_on) {
+          updateData.litellm_params.default_on = defaultOn;
+        }
+
+        await updateGuardrailCall(accessToken, editData.guardrail_id, updateData);
+        NotificationsManager.success("Custom code guardrail updated successfully");
+      } else {
+        // Create new guardrail
+        const guardrailData = {
+          guardrail_name: guardrailName,
+          litellm_params: {
+            guardrail: "custom_code",
+            mode: mode,
+            default_on: defaultOn,
+            custom_code: code,
+          },
+          guardrail_info: {},
+        };
+
+        await createGuardrailCall(accessToken, guardrailData);
+        NotificationsManager.success("Custom code guardrail created successfully");
+      }
       onSuccess();
       onClose();
     } catch (error) {
-      console.error("Failed to create guardrail:", error);
+      console.error("Failed to save guardrail:", error);
       NotificationsManager.fromBackend(
-        "Failed to create guardrail: " + (error instanceof Error ? error.message : String(error))
+        `Failed to ${isEditMode ? "update" : "create"} guardrail: ` + (error instanceof Error ? error.message : String(error))
       );
     } finally {
       setIsSaving(false);
@@ -252,10 +396,20 @@ const CustomCodeModal: React.FC<CustomCodeModalProps> = ({
         parsedInput.texts = [];
       }
 
+      // Use first request-like or response-like mode for test input_type
+      const requestModes = ["pre_call", "pre_mcp_call"];
+      const responseModes = ["post_call", "post_mcp_call"];
+      const testInputType: "request" | "response" =
+        mode.some((m) => requestModes.includes(m))
+          ? "request"
+          : mode.some((m) => responseModes.includes(m))
+            ? "response"
+            : "request";
+
       const response = await testCustomCodeGuardrail(accessToken, {
         custom_code: code,
         test_input: parsedInput,
-        input_type: mode as "request" | "response",
+        input_type: testInputType,
         request_data: {
           model: "test-model",
           metadata: {},
@@ -289,7 +443,7 @@ const CustomCodeModal: React.FC<CustomCodeModalProps> = ({
       open={visible}
       onCancel={onClose}
       footer={null}
-      width={1200}
+      width={1400}
       className="custom-code-modal"
       closable={true}
       destroyOnClose
@@ -297,7 +451,9 @@ const CustomCodeModal: React.FC<CustomCodeModalProps> = ({
       <div className="flex flex-col h-[80vh]">
         {/* Header */}
         <div className="pb-4 border-b border-gray-200">
-          <h2 className="text-xl font-semibold text-gray-900">Create Custom Guardrail</h2>
+          <h2 className="text-xl font-semibold text-gray-900">
+            {isEditMode ? "Edit Custom Guardrail" : "Create Custom Guardrail"}
+          </h2>
           <p className="text-sm text-gray-500 mt-1">Define custom logic using Python-like syntax</p>
         </div>
 
@@ -311,14 +467,16 @@ const CustomCodeModal: React.FC<CustomCodeModalProps> = ({
               placeholder="e.g., block-pii-custom"
             />
           </div>
-          <div className="w-[180px]">
-            <label className="block text-xs font-medium text-gray-600 mb-1">Mode</label>
+          <div className="w-[280px]">
+            <label className="block text-xs font-medium text-gray-600 mb-1">Mode (can select multiple)</label>
             <Select
+              mode="multiple"
               value={mode}
               onChange={setMode}
               options={MODE_OPTIONS}
               className="w-full"
               size="middle"
+              placeholder="Select modes"
             />
           </div>
           <div className="w-[180px]">
@@ -343,21 +501,21 @@ const CustomCodeModal: React.FC<CustomCodeModalProps> = ({
         </div>
 
         {/* Main Content */}
-        <div className="flex flex-1 overflow-hidden mt-4 gap-4">
+        <div className="flex flex-1 overflow-hidden mt-4 gap-6">
           {/* Code Editor */}
-          <div className="flex-1 flex flex-col min-w-0">
-            <div className="flex items-center justify-between mb-2">
+          <div className="flex-[2] flex flex-col min-w-0 overflow-y-auto">
+            <div className="flex items-center justify-between mb-2 flex-shrink-0">
               <span className="text-xs font-semibold text-gray-500 uppercase tracking-wide">Python Logic</span>
               <span className="text-xs text-gray-400">Restricted environment (no imports)</span>
             </div>
-            <div className="flex-1 relative rounded-lg overflow-hidden border border-gray-700 bg-[#1e1e1e]">
+            <div className="relative rounded-lg overflow-hidden border border-gray-700 bg-[#1e1e1e] flex-shrink-0" style={{ minHeight: "300px", maxHeight: "400px" }}>
               {/* Line numbers */}
               <div 
-                className="absolute left-0 top-0 bottom-0 w-10 bg-[#1e1e1e] border-r border-gray-700 text-right pr-2 pt-3 select-none overflow-hidden"
-                style={{ fontFamily: "monospace", fontSize: "13px", lineHeight: "1.5" }}
+                className="absolute left-0 top-0 bottom-0 w-12 bg-[#1e1e1e] border-r border-gray-700 text-right pr-3 pt-3 select-none overflow-hidden"
+                style={{ fontFamily: "'Fira Code', 'Monaco', 'Consolas', monospace", fontSize: "14px", lineHeight: "1.6" }}
               >
                 {Array.from({ length: Math.max(lineCount, 20) }, (_, i) => (
-                  <div key={i + 1} className="text-gray-500 h-[19.5px]">{i + 1}</div>
+                  <div key={i + 1} className="text-gray-500 h-[22.4px]">{i + 1}</div>
                 ))}
               </div>
               {/* Code textarea */}
@@ -367,8 +525,8 @@ const CustomCodeModal: React.FC<CustomCodeModalProps> = ({
                 onChange={(e) => setCode(e.target.value)}
                 onKeyDown={handleKeyDown}
                 spellCheck={false}
-                className="w-full h-full pl-12 pr-4 pt-3 pb-3 font-mono text-sm resize-none focus:outline-none bg-transparent text-gray-200"
-                style={{ lineHeight: "1.5", tabSize: 4 }}
+                className="w-full h-full pl-14 pr-4 pt-3 pb-3 resize-none focus:outline-none bg-transparent text-gray-200"
+                style={{ fontFamily: "'Fira Code', 'Monaco', 'Consolas', monospace", fontSize: "14px", lineHeight: "1.6", tabSize: 4 }}
               />
             </div>
 
@@ -376,7 +534,7 @@ const CustomCodeModal: React.FC<CustomCodeModalProps> = ({
             <Collapse
               activeKey={testExpanded ? ["test"] : []}
               onChange={(keys) => setTestExpanded(keys.includes("test"))}
-              className="mt-3 bg-white border border-gray-200 rounded-lg"
+              className="mt-3 bg-white border border-gray-200 rounded-lg flex-shrink-0"
               expandIcon={({ isActive }) => <CaretRightOutlined rotate={isActive ? 90 : 0} />}
             >
               <Panel
@@ -390,11 +548,40 @@ const CustomCodeModal: React.FC<CustomCodeModalProps> = ({
               >
                 <div className="space-y-3">
                   <div>
-                    <label className="block text-xs font-medium text-gray-600 mb-1">Test Input (JSON)</label>
+                    <div className="flex items-center justify-between mb-2">
+                      <label className="block text-xs font-medium text-gray-600">Test Input (JSON)</label>
+                      <div className="flex items-center gap-2">
+                        <span className="text-xs text-gray-500">Load example:</span>
+                        <button
+                          type="button"
+                          onClick={() => setTestInput(JSON.stringify(TEST_INPUT_EXAMPLES.pre_call.data, null, 2))}
+                          className="px-2 py-1 text-xs rounded border border-orange-200 bg-orange-50 text-orange-700 hover:bg-orange-100 transition-colors"
+                        >
+                          Pre-call
+                        </button>
+                        <button
+                          type="button"
+                          onClick={() => setTestInput(JSON.stringify(TEST_INPUT_EXAMPLES.post_call.data, null, 2))}
+                          className="px-2 py-1 text-xs rounded border border-green-200 bg-green-50 text-green-700 hover:bg-green-100 transition-colors"
+                        >
+                          Post-call
+                        </button>
+                      </div>
+                    </div>
+                    <div className="mb-2 p-2 bg-gray-50 rounded text-xs text-gray-600 border border-gray-200">
+                      <div className="grid grid-cols-2 gap-x-4 gap-y-1">
+                        <div><strong>texts</strong>: Message content (always)</div>
+                        <div><strong>images</strong>: Base64 images (vision)</div>
+                        <div><strong>tools</strong>: Tool definitions <span className="text-orange-600">(pre_call)</span></div>
+                        <div><strong>tool_calls</strong>: LLM tool calls <span className="text-green-600">(post_call)</span></div>
+                        <div><strong>structured_messages</strong>: Full messages <span className="text-orange-600">(pre_call)</span></div>
+                        <div><strong>model</strong>: Model name (always)</div>
+                      </div>
+                    </div>
                     <TextArea
                       value={testInput}
                       onChange={(e) => setTestInput(e.target.value)}
-                      rows={4}
+                      rows={8}
                       className="font-mono text-xs"
                       placeholder='{"texts": ["test message"], ...}'
                     />
@@ -448,7 +635,7 @@ const CustomCodeModal: React.FC<CustomCodeModalProps> = ({
           </div>
 
           {/* Primitives Panel */}
-          <div className="w-[280px] flex-shrink-0 overflow-auto">
+          <div className="w-[300px] flex-shrink-0 overflow-auto border-l border-gray-200 pl-6">
             <div className="flex items-center gap-2 mb-3">
               <CodeOutlined className="text-blue-500" />
               <span className="font-semibold text-gray-700">Available Primitives</span>
@@ -509,7 +696,7 @@ const CustomCodeModal: React.FC<CustomCodeModalProps> = ({
               disabled={isSaving || !guardrailName.trim()}
               icon={SaveOutlined}
             >
-              Save Guardrail
+              {isEditMode ? "Update Guardrail" : "Save Guardrail"}
             </Button>
           </div>
         </div>

ui/litellm-dashboard/src/components/guardrails/guardrail_info.tsx

@@ -14,7 +14,7 @@ import {
   TextInput,
 } from "@tremor/react";
 import { Button, Form, Input, Select, Divider, Tooltip } from "antd";
-import { InfoCircleOutlined, EyeInvisibleOutlined, StopOutlined } from "@ant-design/icons";
+import { InfoCircleOutlined, EyeInvisibleOutlined, StopOutlined, CodeOutlined } from "@ant-design/icons";
 import {
   getGuardrailInfo,
   updateGuardrailCall,
@@ -29,6 +29,7 @@ import ContentFilterManager, { formatContentFilterDataForAPI } from "./content_f
 import ToolPermissionRulesEditor, {
   ToolPermissionConfig,
 } from "./tool_permission/ToolPermissionRulesEditor";
+import CustomCodeModal, { EditGuardrailData } from "./custom_code/CustomCodeModal";
 import { ArrowLeftIcon } from "@heroicons/react/outline";
 import { copyToClipboard as utilCopyToClipboard } from "@/utils/dataUtils";
 import { CheckIcon, CopyIcon } from "lucide-react";
@@ -94,6 +95,7 @@ const GuardrailInfoView: React.FC<GuardrailInfoProps> = ({ guardrailId, onClose,
   };
   const [toolPermissionConfig, setToolPermissionConfig] = useState<ToolPermissionConfig>(emptyToolPermissionConfig);
   const [toolPermissionDirty, setToolPermissionDirty] = useState(false);
+  const [customCodeModalVisible, setCustomCodeModalVisible] = useState(false);
 
   // Content Filter data ref (managed by ContentFilterManager)
   const contentFilterDataRef = React.useRef<{ patterns: any[]; blockedWords: any[] }>({
@@ -552,6 +554,33 @@ const GuardrailInfoView: React.FC<GuardrailInfoProps> = ({ guardrailId, onClose,
               </Card>
             )}
 
+            {/* Custom Code Display */}
+            {guardrailData.litellm_params?.guardrail === "custom_code" && guardrailData.litellm_params?.custom_code && (
+              <Card className="mt-6">
+                <div className="flex justify-between items-center mb-4">
+                  <div className="flex items-center gap-2">
+                    <CodeOutlined className="text-blue-500" />
+                    <Text className="font-medium text-lg">Custom Code</Text>
+                  </div>
+                  {isAdmin && !isConfigGuardrail && (
+                    <TremorButton 
+                      size="xs"
+                      variant="secondary"
+                      icon={CodeOutlined}
+                      onClick={() => setCustomCodeModalVisible(true)}
+                    >
+                      Edit Code
+                    </TremorButton>
+                  )}
+                </div>
+                <div className="relative rounded-lg overflow-hidden border border-gray-700 bg-[#1e1e1e]">
+                  <pre className="p-4 text-sm text-gray-200 overflow-x-auto" style={{ fontFamily: "'Fira Code', 'Monaco', 'Consolas', monospace" }}>
+                    <code>{guardrailData.litellm_params.custom_code}</code>
+                  </pre>
+                </div>
+              </Card>
+            )}
+
             {/* Content Filter Configuration Display */}
             <ContentFilterManager
               guardrailData={guardrailData}
@@ -573,7 +602,16 @@ const GuardrailInfoView: React.FC<GuardrailInfoProps> = ({ guardrailId, onClose,
                     </Tooltip>
                   )}
                   {!isEditing && !isConfigGuardrail && (
-                    <TremorButton onClick={() => setIsEditing(true)}>Edit Settings</TremorButton>
+                    guardrailData.litellm_params?.guardrail === "custom_code" ? (
+                      <TremorButton 
+                        icon={CodeOutlined}
+                        onClick={() => setCustomCodeModalVisible(true)}
+                      >
+                        Edit Code
+                      </TremorButton>
+                    ) : (
+                      <TremorButton onClick={() => setIsEditing(true)}>Edit Settings</TremorButton>
+                    )
                   )}
                 </div>
 
@@ -757,6 +795,22 @@ const GuardrailInfoView: React.FC<GuardrailInfoProps> = ({ guardrailId, onClose,
           )}
         </TabPanels>
       </TabGroup>
+
+      {/* Custom Code Editor Modal */}
+      <CustomCodeModal
+        visible={customCodeModalVisible}
+        onClose={() => setCustomCodeModalVisible(false)}
+        onSuccess={() => {
+          setCustomCodeModalVisible(false);
+          fetchGuardrailInfo();
+        }}
+        accessToken={accessToken}
+        editData={guardrailData ? {
+          guardrail_id: guardrailData.guardrail_id,
+          guardrail_name: guardrailData.guardrail_name,
+          litellm_params: guardrailData.litellm_params,
+        } as EditGuardrailData : null}
+      />
     </div>
   );
 };

ui/litellm-dashboard/src/components/mcp_tools/MCPToolArgumentsForm.tsx

@@ -0,0 +1,345 @@
+import React, { forwardRef, useImperativeHandle, useMemo } from "react";
+import { Form, Input, InputNumber, Select, Tooltip } from "antd";
+import { InfoCircleOutlined } from "@ant-design/icons";
+import { MCPTool, InputSchema, InputSchemaProperty } from "./types";
+
+const isPlainObject = (value: unknown): value is Record<string, any> =>
+  typeof value === "object" && value !== null && !Array.isArray(value);
+
+function buildArrayItems(items?: InputSchemaProperty | InputSchemaProperty[]): any[] {
+  if (!items) return [];
+  if (Array.isArray(items)) {
+    return items
+      .map((item) => buildDefaultValue(item))
+      .filter((value) => value !== undefined);
+  }
+  const itemDefault = buildDefaultValue(items);
+  return itemDefault !== undefined ? [itemDefault] : [];
+}
+
+function buildDefaultValue(prop?: InputSchemaProperty, overrideDefault?: any): any {
+  if (!prop) return undefined;
+  const effectiveDefault = overrideDefault !== undefined ? overrideDefault : prop.default;
+
+  if (prop.type === "object") {
+    const base = isPlainObject(effectiveDefault) ? { ...effectiveDefault } : {};
+    if (prop.properties) {
+      Object.entries(prop.properties).forEach(([childKey, childProp]) => {
+        base[childKey] = buildDefaultValue(childProp, base[childKey]);
+      });
+    }
+    return base;
+  }
+
+  if (prop.type === "array") {
+    if (Array.isArray(effectiveDefault)) {
+      const itemSchema = prop.items;
+      if (!itemSchema) return effectiveDefault;
+      if (effectiveDefault.length === 0) {
+        const sample = buildArrayItems(itemSchema);
+        return sample.length ? sample : effectiveDefault;
+      }
+      if (Array.isArray(itemSchema)) {
+        return effectiveDefault.map((value, index) => {
+          const schema = itemSchema[index] ?? itemSchema[itemSchema.length - 1];
+          return buildDefaultValue(schema, value);
+        });
+      }
+      return effectiveDefault.map((value) => buildDefaultValue(itemSchema, value));
+    }
+    if (effectiveDefault !== undefined) return effectiveDefault;
+    return buildArrayItems(prop.items);
+  }
+
+  if (effectiveDefault !== undefined) return effectiveDefault;
+  switch (prop.type) {
+    case "integer":
+    case "number":
+      return 0;
+    case "boolean":
+      return false;
+    case "string":
+    default:
+      return "";
+  }
+}
+
+const getInitialValueForField = (prop: InputSchemaProperty): any => {
+  const defaultValue = buildDefaultValue(prop);
+  if (prop.type === "object" || prop.type === "array") {
+    const fallback = prop.type === "array" ? [] : {};
+    return JSON.stringify(defaultValue ?? fallback, null, 2);
+  }
+  return defaultValue;
+};
+
+function convertFormValues(
+  values: Record<string, any>,
+  actualSchema: InputSchema,
+  schema: InputSchema,
+): Record<string, any> {
+  const convertedValues: Record<string, any> = {};
+  const schemaToUse = actualSchema;
+
+  Object.entries(values).forEach(([key, value]) => {
+    const prop = schemaToUse.properties?.[key];
+    if (prop && value !== null && value !== undefined && value !== "") {
+      switch (prop.type) {
+        case "boolean":
+          convertedValues[key] = value === "true" || value === true;
+          break;
+        case "number":
+        case "integer": {
+          const numericValue = Number(value);
+          convertedValues[key] = Number.isNaN(numericValue)
+            ? value
+            : prop.type === "integer"
+              ? Math.trunc(numericValue)
+              : numericValue;
+          break;
+        }
+        case "object":
+        case "array": {
+          try {
+            const parsed = typeof value === "string" ? JSON.parse(value) : value;
+            const isValidObject =
+              prop.type === "object" &&
+              parsed !== null &&
+              typeof parsed === "object" &&
+              !Array.isArray(parsed);
+            const isValidArray = prop.type === "array" && Array.isArray(parsed);
+            if ((prop.type === "object" && isValidObject) || (prop.type === "array" && isValidArray)) {
+              convertedValues[key] = parsed;
+            } else {
+              convertedValues[key] = value;
+            }
+          } catch {
+            convertedValues[key] = value;
+          }
+          break;
+        }
+        case "string":
+          convertedValues[key] = String(value);
+          break;
+        default:
+          convertedValues[key] = value;
+      }
+    } else if (value !== null && value !== undefined && value !== "") {
+      convertedValues[key] = value;
+    }
+  });
+
+  const isNestedParams =
+    schema.properties?.params?.type === "object" && schema.properties.params.properties;
+
+  return isNestedParams ? { params: convertedValues } : convertedValues;
+}
+
+export interface MCPToolArgumentsFormRef {
+  getSubmitValues: () => Promise<Record<string, any>>;
+}
+
+interface MCPToolArgumentsFormProps {
+  tool: MCPTool;
+  className?: string;
+}
+
+const MCPToolArgumentsForm = forwardRef<MCPToolArgumentsFormRef, MCPToolArgumentsFormProps>(
+  ({ tool, className }, ref) => {
+    const [form] = Form.useForm();
+
+    const schema: InputSchema = useMemo(() => {
+      if (typeof tool.inputSchema === "string") {
+        return {
+          type: "object",
+          properties: {
+            input: {
+              type: "string",
+              description: "Input for this tool",
+            },
+          },
+          required: ["input"],
+        };
+      }
+      return tool.inputSchema as InputSchema;
+    }, [tool.inputSchema]);
+
+    const actualSchema: InputSchema = useMemo(() => {
+      if (
+        schema.properties?.params?.type === "object" &&
+        schema.properties.params.properties
+      ) {
+        return {
+          type: "object",
+          properties: schema.properties.params.properties,
+          required: schema.properties.params.required || [],
+        };
+      }
+      return schema;
+    }, [schema]);
+
+    useImperativeHandle(ref, () => ({
+      getSubmitValues: async () => {
+        const values = await form.validateFields();
+        return convertFormValues(values, actualSchema, schema);
+      },
+    }));
+
+    React.useEffect(() => {
+      form.resetFields();
+      if (!actualSchema.properties) return;
+
+      const initialValues: Record<string, any> = {};
+      Object.entries(actualSchema.properties).forEach(([key, prop]) => {
+        initialValues[key] = getInitialValueForField(prop);
+      });
+      form.setFieldsValue(initialValues);
+    }, [form, actualSchema, tool]);
+
+    if (typeof tool.inputSchema === "string") {
+      return (
+        <Form form={form} layout="vertical" className={className}>
+          <Form.Item
+            label={
+              <span className="text-sm font-medium text-gray-700">
+                Input <span className="text-red-500">*</span>
+              </span>
+            }
+            name="input"
+            rules={[{ required: true, message: "Please enter input for this tool" }]}
+          >
+            <Input placeholder="Enter input for this tool" />
+          </Form.Item>
+        </Form>
+      );
+    }
+
+    if (!actualSchema.properties) {
+      return (
+        <Form form={form} layout="vertical" className={className}>
+          <div className="py-4 text-center text-sm text-gray-500">
+            No parameters required for this tool.
+          </div>
+        </Form>
+      );
+    }
+
+    return (
+      <Form form={form} layout="vertical" className={className}>
+        {Object.entries(actualSchema.properties).map(([key, prop]) => {
+          const initialValue = getInitialValueForField(prop);
+          const fieldKey = `${tool.name}-${key}`;
+          return (
+            <Form.Item
+              key={fieldKey}
+              label={
+                <span className="text-sm font-medium text-gray-700 flex items-center">
+                  {key} {actualSchema.required?.includes(key) && <span className="text-red-500">*</span>}
+                  {prop.description && (
+                    <Tooltip title={prop.description}>
+                      <InfoCircleOutlined className="ml-2 text-gray-400 hover:text-gray-600" />
+                    </Tooltip>
+                  )}
+                </span>
+              }
+              name={key}
+              initialValue={initialValue}
+              rules={[
+                {
+                  required: actualSchema.required?.includes(key),
+                  message: `Please enter ${key}`,
+                },
+                ...(prop.type === "object" || prop.type === "array"
+                  ? [
+                      {
+                        validator: (_rule: any, value: any) => {
+                          if (
+                            (value === undefined || value === null || value === "") &&
+                            !actualSchema.required?.includes(key)
+                          ) {
+                            return Promise.resolve();
+                          }
+                          try {
+                            const parsed = typeof value === "string" ? JSON.parse(value) : value;
+                            const isValidObject =
+                              prop.type === "object" &&
+                              parsed !== null &&
+                              typeof parsed === "object" &&
+                              !Array.isArray(parsed);
+                            const isValidArray = prop.type === "array" && Array.isArray(parsed);
+                            if (
+                              (prop.type === "object" && isValidObject) ||
+                              (prop.type === "array" && isValidArray)
+                            ) {
+                              return Promise.resolve();
+                            }
+                            return Promise.reject(
+                              new Error(
+                                prop.type === "object" ? "Please enter a JSON object" : "Please enter a JSON array",
+                              ),
+                            );
+                          } catch {
+                            return Promise.reject(new Error("Invalid JSON"));
+                          }
+                        },
+                      },
+                    ]
+                  : []),
+              ]}
+            >
+              {prop.type === "string" && prop.enum ? (
+                <Select
+                  placeholder={`Select ${key}`}
+                  allowClear={!actualSchema.required?.includes(key)}
+                  options={prop.enum.map((v) => ({ value: v, label: v }))}
+                />
+              ) : prop.type === "string" && !prop.enum ? (
+                <Input
+                  placeholder={prop.description || `Enter ${key}`}
+                  allowClear
+                />
+              ) : prop.type === "number" || prop.type === "integer" ? (
+                <InputNumber
+                  step={prop.type === "integer" ? 1 : undefined}
+                  placeholder={prop.description || `Enter ${key}`}
+                  className="w-full"
+                  style={{ width: "100%" }}
+                />
+              ) : prop.type === "boolean" ? (
+                <Select
+                  placeholder={`Select ${key}`}
+                  allowClear={!actualSchema.required?.includes(key)}
+                  options={[
+                    { value: true, label: "True" },
+                    { value: false, label: "False" },
+                  ]}
+                />
+              ) : (prop.type === "object" || prop.type === "array") ? (
+                <Input.TextArea
+                  rows={prop.type === "object" ? 4 : 3}
+                  placeholder={
+                    prop.description ||
+                    (prop.type === "object"
+                      ? `Enter JSON object for ${key}`
+                      : `Enter JSON array for ${key}`)
+                  }
+                  spellCheck={false}
+                  className="font-mono"
+                />
+              ) : (
+                <Input
+                  placeholder={prop.description || `Enter ${key}`}
+                  allowClear
+                />
+              )}
+            </Form.Item>
+          );
+        })}
+      </Form>
+    );
+  },
+);
+
+MCPToolArgumentsForm.displayName = "MCPToolArgumentsForm";
+
+export default MCPToolArgumentsForm;

ui/litellm-dashboard/src/components/molecules/notifications_manager.tsx

@@ -1,8 +1,18 @@
 import React from "react";
-import { notification } from "antd";
+import { notification as staticNotification } from "antd";
+import type { NotificationInstance } from "antd/es/notification/interface";
 import { parseErrorMessage } from "../shared/errorUtils";
 import { ArgsProps } from "antd/es/notification";
 
+let notificationInstance: NotificationInstance | null = null;
+
+export const setNotificationInstance = (instance: NotificationInstance) => {
+  notificationInstance = instance;
+};
+
+// Helper to get the best available notification instance
+const getNotification = () => notificationInstance || staticNotification;
+
 type Placement = "top" | "topLeft" | "topRight" | "bottom" | "bottomLeft" | "bottomRight";
 
 type NotificationConfig = {
@@ -251,7 +261,7 @@ function looksErrorPayload(input: any, status?: number): boolean {
 const NotificationManager = {
   error(input: string | NotificationConfig) {
     const cfg = normalize(input, "Error");
-    notification.error({
+    getNotification().error({
       ...COMMON_NOTIFICATION_PROPS,
       ...cfg,
       placement: cfg.placement ?? defaultPlacement(),
@@ -261,7 +271,7 @@ const NotificationManager = {
 
   warning(input: string | NotificationConfig) {
     const cfg = normalize(input, "Warning");
-    notification.warning({
+    getNotification().warning({
       ...COMMON_NOTIFICATION_PROPS,
       ...cfg,
       placement: cfg.placement ?? defaultPlacement(),
@@ -271,7 +281,7 @@ const NotificationManager = {
 
   info(input: string | NotificationConfig) {
     const cfg = normalize(input, "Info");
-    notification.info({
+    getNotification().info({
       ...COMMON_NOTIFICATION_PROPS,
       ...cfg,
       placement: cfg.placement ?? defaultPlacement(),
@@ -281,7 +291,7 @@ const NotificationManager = {
 
   success(input: string | React.ReactNode | NotificationConfig) {
     if (React.isValidElement(input)) {
-      notification.success({
+      getNotification().success({
         ...COMMON_NOTIFICATION_PROPS,
         message: "Success",
         description: input,
@@ -291,7 +301,7 @@ const NotificationManager = {
       return;
     }
     const cfg = normalize(input as string | NotificationConfig, "Success");
-    notification.success({
+    getNotification().success({
       ...COMMON_NOTIFICATION_PROPS,
       ...cfg,
       placement: cfg.placement ?? defaultPlacement(),
@@ -316,11 +326,11 @@ const NotificationManager = {
         title === "Content Blocked" ||
         title === "Integration Error"
       ) {
-        notification.warning({ ...COMMON_NOTIFICATION_PROPS, ...payload, duration: extra?.duration ?? 7 });
+        getNotification().warning({ ...COMMON_NOTIFICATION_PROPS, ...payload, duration: extra?.duration ?? 7 });
         return;
       }
       if (title === "Server Error") {
-        notification.error({ ...COMMON_NOTIFICATION_PROPS, ...payload, duration: extra?.duration ?? 8 });
+        getNotification().error({ ...COMMON_NOTIFICATION_PROPS, ...payload, duration: extra?.duration ?? 8 });
         return;
       }
       if (
@@ -331,10 +341,10 @@ const NotificationManager = {
         title === "Error" ||
         title === "Already Exists"
       ) {
-        notification.error({ ...COMMON_NOTIFICATION_PROPS, ...payload, duration: extra?.duration ?? 6 });
+        getNotification().error({ ...COMMON_NOTIFICATION_PROPS, ...payload, duration: extra?.duration ?? 6 });
         return;
       }
-      notification.info({ ...COMMON_NOTIFICATION_PROPS, ...payload, duration: extra?.duration ?? 4 });
+      getNotification().info({ ...COMMON_NOTIFICATION_PROPS, ...payload, duration: extra?.duration ?? 4 });
       return;
     }
 
@@ -343,18 +353,18 @@ const NotificationManager = {
     const payload = { ...base, message: cls?.title ?? "Info" };
 
     if (cls?.kind === "success") {
-      notification.success({ ...COMMON_NOTIFICATION_PROPS, ...payload, duration: extra?.duration ?? 3.5 });
+      getNotification().success({ ...COMMON_NOTIFICATION_PROPS, ...payload, duration: extra?.duration ?? 3.5 });
       return;
     }
     if (cls?.kind === "warning") {
-      notification.warning({ ...COMMON_NOTIFICATION_PROPS, ...payload, duration: extra?.duration ?? 6 });
+      getNotification().warning({ ...COMMON_NOTIFICATION_PROPS, ...payload, duration: extra?.duration ?? 6 });
       return;
     }
-    notification.info({ ...COMMON_NOTIFICATION_PROPS, ...payload, duration: extra?.duration ?? 4 });
+    getNotification().info({ ...COMMON_NOTIFICATION_PROPS, ...payload, duration: extra?.duration ?? 4 });
   },
 
   clear() {
-    notification.destroy();
+    getNotification().destroy();
   },
 };
 

ui/litellm-dashboard/src/components/networking.tsx

@@ -6685,11 +6685,16 @@ export const listMCPTools = async (accessToken: string, serverId: string) => {
   }
 };
 
+export interface CallMCPToolOptions {
+  guardrails?: string[];
+}
+
 export const callMCPTool = async (
   accessToken: string,
   serverId: string,
   toolName: string,
   toolArguments: Record<string, any>,
+  options?: CallMCPToolOptions,
 ) => {
   try {
     // Construct base URL
@@ -6702,14 +6707,19 @@ export const callMCPTool = async (
       "Content-Type": "application/json",
     };
 
+    const body: Record<string, any> = {
+      server_id: serverId,
+      name: toolName,
+      arguments: toolArguments,
+    };
+    if (options?.guardrails && options.guardrails.length > 0) {
+      body.litellm_metadata = { guardrails: options.guardrails };
+    }
+
     const response = await fetch(url, {
       method: "POST",
       headers,
-      body: JSON.stringify({
-        server_id: serverId,
-        name: toolName,
-        arguments: toolArguments,
-      }),
+      body: JSON.stringify(body),
     });
 
     if (!response.ok) {

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

@@ -31,9 +31,10 @@ import { v4 as uuidv4 } from "uuid";
 import { truncateString } from "../../../utils/textUtils";
 import GuardrailSelector from "../../guardrails/GuardrailSelector";
 import PolicySelector from "../../policies/PolicySelector";
+import MCPToolArgumentsForm, { MCPToolArgumentsFormRef } from "../../mcp_tools/MCPToolArgumentsForm";
 import { MCPServer } from "../../mcp_tools/types";
 import NotificationsManager from "../../molecules/notifications_manager";
-import { fetchMCPServers, listMCPTools } from "../../networking";
+import { callMCPTool, fetchMCPServers, listMCPTools } from "../../networking";
 import TagSelector from "../../tag_management/TagSelector";
 import VectorStoreSelector from "../../vector_store_management/VectorStoreSelector";
 import { makeA2ASendMessageRequest } from "../llm_calls/a2a_send_message";
@@ -85,7 +86,11 @@ interface ChatUIProps {
   };
 }
 
-const MCP_SUPPORTED_ENDPOINTS = new Set<EndpointType>([EndpointType.CHAT, EndpointType.RESPONSES]);
+const MCP_SUPPORTED_ENDPOINTS = new Set<EndpointType>([
+  EndpointType.CHAT,
+  EndpointType.RESPONSES,
+  EndpointType.MCP,
+]);
 
 const ChatUI: React.FC<ChatUIProps> = ({
   accessToken,
@@ -107,6 +112,8 @@ const ChatUI: React.FC<ChatUIProps> = ({
   });
   const [isLoadingMCPServers, setIsLoadingMCPServers] = useState(false);
   const [serverToolsMap, setServerToolsMap] = useState<Record<string, any[]>>({});
+  const [selectedMCPDirectTool, setSelectedMCPDirectTool] = useState<string | undefined>(undefined);
+  const mcpToolArgsFormRef = useRef<MCPToolArgumentsFormRef>(null);
   const [mcpServerToolRestrictions, setMCPServerToolRestrictions] = useState<Record<string, string[]>>(() => {
     const saved = sessionStorage.getItem("mcpServerToolRestrictions");
     try {
@@ -396,6 +403,18 @@ const ChatUI: React.FC<ChatUIProps> = ({
     loadMCPServers();
   }, [accessToken, userID, userRole, apiKeySource, apiKey, token]);
 
+  // Load tools when MCP direct mode has a server selected
+  useEffect(() => {
+    if (
+      endpointType === EndpointType.MCP &&
+      selectedMCPServers.length === 1 &&
+      selectedMCPServers[0] !== "__all__" &&
+      !serverToolsMap[selectedMCPServers[0]]
+    ) {
+      loadServerTools(selectedMCPServers[0]);
+    }
+  }, [endpointType, selectedMCPServers, serverToolsMap]);
+
   // Fetch agents when A2A endpoint is selected
   useEffect(() => {
     const userApiKey = apiKeySource === "session" ? accessToken : apiKey;
@@ -769,8 +788,13 @@ const ChatUI: React.FC<ChatUIProps> = ({
     setUploadedAudio(null);
   };
 
-  const handleSendMessage = async () => {
-    if (inputMessage.trim() === "" && endpointType !== EndpointType.TRANSCRIPTION) return;
+    const handleSendMessage = async () => {
+    if (
+      inputMessage.trim() === "" &&
+      endpointType !== EndpointType.TRANSCRIPTION &&
+      endpointType !== EndpointType.MCP
+    )
+      return;
 
     // For image edits, require both image and prompt
     if (endpointType === EndpointType.IMAGE_EDITS && uploadedImages.length === 0) {
@@ -790,7 +814,39 @@ const ChatUI: React.FC<ChatUIProps> = ({
       return;
     }
 
-    // Require model selection for all model-based endpoints
+    // For MCP direct mode, require server and tool selection, and get form values early
+    let mcpToolArguments: Record<string, any> = {};
+    if (endpointType === EndpointType.MCP) {
+      const mcpServerId =
+        selectedMCPServers.length === 1 && selectedMCPServers[0] !== "__all__"
+          ? selectedMCPServers[0]
+          : null;
+      if (!mcpServerId) {
+        NotificationsManager.fromBackend("Please select an MCP server to test");
+        return;
+      }
+      if (!selectedMCPDirectTool) {
+        NotificationsManager.fromBackend("Please select an MCP tool to call");
+        return;
+      }
+      const mcpTool = (serverToolsMap[selectedMCPServers[0]] || []).find(
+        (t: any) => t.name === selectedMCPDirectTool,
+      );
+      if (!mcpTool) {
+        NotificationsManager.fromBackend("Please wait for tool schema to load");
+        return;
+      }
+      try {
+        mcpToolArguments = (await mcpToolArgsFormRef.current?.getSubmitValues()) ?? {};
+      } catch (err) {
+        NotificationsManager.fromBackend(
+          err instanceof Error ? err.message : "Please fill in all required parameters",
+        );
+        return;
+      }
+    }
+
+    // Require model selection for all model-based endpoints (MCP direct mode does not need a model)
     const modelRequiredEndpoints = [
       EndpointType.CHAT,
       EndpointType.IMAGE,
@@ -874,6 +930,10 @@ const ChatUI: React.FC<ChatUIProps> = ({
         ? `🎵 Audio file: ${uploadedAudio.name}\nPrompt: ${inputMessage}`
         : `🎵 Audio file: ${uploadedAudio.name}`;
       displayMessage = createDisplayMessage(audioMessage, false);
+    } else if (endpointType === EndpointType.MCP && selectedMCPDirectTool) {
+      // For MCP direct mode, show tool name and arguments from form
+      const mcpMessage = `🔧 MCP Tool: ${selectedMCPDirectTool}\nArguments: ${JSON.stringify(mcpToolArguments, null, 2)}`;
+      displayMessage = createDisplayMessage(mcpMessage, false);
     } else {
       displayMessage = createDisplayMessage(inputMessage, false);
     }
@@ -1057,6 +1117,32 @@ const ChatUI: React.FC<ChatUIProps> = ({
         }
       }
 
+      // Handle MCP direct tool calls (no chat completions)
+      if (endpointType === EndpointType.MCP) {
+        const mcpServerId =
+          selectedMCPServers.length === 1 && selectedMCPServers[0] !== "__all__"
+            ? selectedMCPServers[0]
+            : null;
+        if (mcpServerId && selectedMCPDirectTool) {
+          const result = await callMCPTool(
+            effectiveApiKey,
+            mcpServerId,
+            selectedMCPDirectTool,
+            mcpToolArguments,
+            selectedGuardrails.length > 0 ? { guardrails: selectedGuardrails } : undefined,
+          );
+          const resultText =
+            result?.content?.length > 0
+              ? JSON.stringify(
+                  result.content.map((c: any) => (c.type === "text" ? c.text : c)).filter(Boolean),
+                  null,
+                  2,
+                )
+              : JSON.stringify(result, null, 2);
+          updateTextUI("assistant", resultText || "Tool executed successfully.");
+        }
+      }
+
       // Handle A2A agent calls (separate from model-based calls) - use streaming
       if (endpointType === EndpointType.A2A_AGENTS && selectedAgent) {
         await makeA2ASendMessageRequest(
@@ -1069,6 +1155,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
           updateTotalLatency,
           updateA2AMetadata,
           customProxyBaseUrl || undefined,
+          selectedGuardrails.length > 0 ? selectedGuardrails : undefined,
         );
       }
     } catch (error) {
@@ -1253,10 +1340,17 @@ const ChatUI: React.FC<ChatUIProps> = ({
                     setSelectedModel(undefined);
                     setSelectedAgent(undefined);
                     setShowCustomModelInput(false);
+                    setSelectedMCPDirectTool(undefined);
+                    // For MCP direct mode, require single server (clear __all__ or multiple)
+                    if (value === EndpointType.MCP) {
+                      setSelectedMCPServers((prev) =>
+                        prev.length === 1 && prev[0] !== "__all__" ? prev : [],
+                      );
+                    }
                     try {
                       sessionStorage.removeItem("selectedModel");
                       sessionStorage.removeItem("selectedAgent");
-                    } catch { }
+                    } catch {}
                   }}
                   className="mb-4"
                 />
@@ -1290,8 +1384,8 @@ const ChatUI: React.FC<ChatUIProps> = ({
                 />
               </div>
 
-              {/* Model Selector - shown when NOT using A2A Agents */}
-              {endpointType !== EndpointType.A2A_AGENTS && (
+              {/* Model Selector - shown when NOT using A2A Agents or MCP direct mode */}
+              {endpointType !== EndpointType.A2A_AGENTS && endpointType !== EndpointType.MCP && (
                 <div>
                   <Text className="font-medium block mb-2 text-gray-700 flex items-center justify-between">
                     <span className="flex items-center">
@@ -1449,36 +1543,59 @@ const ChatUI: React.FC<ChatUIProps> = ({
               {/* MCP Server Selection */}
               <div>
                 <Text className="font-medium block mb-2 text-gray-700 flex items-center">
-                  <ToolOutlined className="mr-2" /> MCP Servers
-                  <Tooltip className="ml-1" title="Select MCP servers to use in your conversation.">
+                  <ToolOutlined className="mr-2" />
+                  {endpointType === EndpointType.MCP ? "MCP Server" : "MCP Servers"}
+                  <Tooltip
+                    className="ml-1"
+                    title={
+                      endpointType === EndpointType.MCP
+                        ? "Select an MCP server to test tools directly."
+                        : "Select MCP servers to use in your conversation."
+                    }
+                  >
                     <InfoCircleOutlined />
                   </Tooltip>
                 </Text>
                 <Select
-                  mode="multiple"
+                  mode={endpointType === EndpointType.MCP ? undefined : "multiple"}
                   style={{ width: "100%" }}
-                  placeholder="Select MCP servers"
-                  value={selectedMCPServers}
+                  placeholder={
+                    endpointType === EndpointType.MCP ? "Select MCP server" : "Select MCP servers"
+                  }
+                  value={
+                    endpointType === EndpointType.MCP
+                      ? selectedMCPServers[0] !== "__all__" && selectedMCPServers.length === 1
+                        ? selectedMCPServers[0]
+                        : undefined
+                      : selectedMCPServers
+                  }
                   onChange={(value) => {
-                    if (value.includes("__all__")) {
-                      setSelectedMCPServers(["__all__"]);
-                      setMCPServerToolRestrictions({});
+                    if (endpointType === EndpointType.MCP) {
+                      const serverId = value as string | undefined;
+                      setSelectedMCPServers(serverId ? [serverId] : []);
+                      setSelectedMCPDirectTool(undefined);
+                      if (serverId && !serverToolsMap[serverId]) {
+                        loadServerTools(serverId);
+                      }
                     } else {
-                      setSelectedMCPServers(value);
-                      // Clean up tool restrictions for removed servers
-                      setMCPServerToolRestrictions((prev) => {
-                        const updated = { ...prev };
-                        Object.keys(updated).forEach((serverId) => {
-                          if (!value.includes(serverId)) delete updated[serverId];
+                      if ((value as string[]).includes("__all__")) {
+                        setSelectedMCPServers(["__all__"]);
+                        setMCPServerToolRestrictions({});
+                      } else {
+                        setSelectedMCPServers(value as string[]);
+                        setMCPServerToolRestrictions((prev) => {
+                          const updated = { ...prev };
+                          Object.keys(updated).forEach((serverId) => {
+                            if (!(value as string[]).includes(serverId)) delete updated[serverId];
+                          });
+                          return updated;
                         });
-                        return updated;
-                      });
-                      // Load tools for newly selected servers
-                      value.forEach((serverId) => {
-                        if (!serverToolsMap[serverId]) {
-                          loadServerTools(serverId);
-                        }
-                      });
+                        (value as string[]).forEach((serverId) => {
+                          if (!serverToolsMap[serverId]) {
+                            loadServerTools(serverId);
+                          }
+                        });
+                      }
                     }
                   }}
                   loading={isLoadingMCPServers}
@@ -1486,15 +1603,17 @@ const ChatUI: React.FC<ChatUIProps> = ({
                   allowClear
                   optionLabelProp="label"
                   disabled={!MCP_SUPPORTED_ENDPOINTS.has(endpointType as EndpointType)}
-                  maxTagCount="responsive"
+                  maxTagCount={endpointType === EndpointType.MCP ? 1 : "responsive"}
                 >
-                  {/* All MCP Servers option */}
-                  <Select.Option key="__all__" value="__all__" label="All MCP Servers">
-                    <div className="flex flex-col py-1">
-                      <span className="font-medium">All MCP Servers</span>
-                      <span className="text-xs text-gray-500 mt-1">Use all available MCP servers</span>
-                    </div>
-                  </Select.Option>
+                  {/* All MCP Servers option - hidden for MCP direct mode */}
+                  {endpointType !== EndpointType.MCP && (
+                    <Select.Option key="__all__" value="__all__" label="All MCP Servers">
+                      <div className="flex flex-col py-1">
+                        <span className="font-medium">All MCP Servers</span>
+                        <span className="text-xs text-gray-500 mt-1">Use all available MCP servers</span>
+                      </div>
+                    </Select.Option>
+                  )}
 
                   {/* Individual servers */}
                   {mcpServers.map((server) => (
@@ -1502,7 +1621,9 @@ const ChatUI: React.FC<ChatUIProps> = ({
                       key={server.server_id}
                       value={server.server_id}
                       label={server.alias || server.server_name || server.server_id}
-                      disabled={selectedMCPServers.includes("__all__")}
+                      disabled={
+                        endpointType === EndpointType.MCP ? false : selectedMCPServers.includes("__all__")
+                      }
                     >
                       <div className="flex flex-col py-1">
                         <span className="font-medium">{server.alias || server.server_name || server.server_id}</span>
@@ -1512,9 +1633,31 @@ const ChatUI: React.FC<ChatUIProps> = ({
                   ))}
                 </Select>
 
-                {/* Tool restrictions UI (optional) */}
+                {/* MCP Tool selector - only for MCP direct mode */}
+                {endpointType === EndpointType.MCP &&
+                  selectedMCPServers.length === 1 &&
+                  selectedMCPServers[0] !== "__all__" && (
+                    <div className="mt-3">
+                      <Text className="text-xs text-gray-600 mb-1 block">Select Tool</Text>
+                      <Select
+                        style={{ width: "100%" }}
+                        placeholder="Select a tool to call"
+                        value={selectedMCPDirectTool}
+                        onChange={(value) => setSelectedMCPDirectTool(value)}
+                        options={(serverToolsMap[selectedMCPServers[0]] || []).map((tool: any) => ({
+                          value: tool.name,
+                          label: tool.name,
+                        }))}
+                        allowClear
+                        className="rounded-md"
+                      />
+                    </div>
+                  )}
+
+                {/* Tool restrictions UI (optional) - hidden for MCP direct mode */}
                 {selectedMCPServers.length > 0 &&
                   !selectedMCPServers.includes("__all__") &&
+                  endpointType !== EndpointType.MCP &&
                   MCP_SUPPORTED_ENDPOINTS.has(endpointType as EndpointType) && (
                     <div className="mt-3 space-y-2">
                       {selectedMCPServers.map((serverId) => {
@@ -2085,8 +2228,8 @@ const ChatUI: React.FC<ChatUIProps> = ({
                 </div>
               )}
 
-              {/* Suggested prompts - show when chat is empty and not loading */}
-              {chatHistory.length === 0 && !isLoading && (
+              {/* Suggested prompts - show when chat is empty and not loading (skip for MCP - uses structured form) */}
+              {chatHistory.length === 0 && !isLoading && endpointType !== EndpointType.MCP && (
                 <div className="flex items-center gap-2 mb-3 overflow-x-auto">
                   {(endpointType === EndpointType.A2A_AGENTS
                     ? ["What can you help me with?", "Tell me about yourself", "What tasks can you perform?"]
@@ -2151,46 +2294,79 @@ const ChatUI: React.FC<ChatUIProps> = ({
                     )}
                   </div>
 
-                  {/* Middle: input field */}
-                  <TextArea
-                    value={inputMessage}
-                    onChange={(e) => setInputMessage(e.target.value)}
-                    onKeyDown={handleKeyDown}
-                    placeholder={
-                      endpointType === EndpointType.CHAT ||
+                  {/* Middle: input field or MCP structured form */}
+                  {endpointType === EndpointType.MCP &&
+                  selectedMCPServers.length === 1 &&
+                  selectedMCPServers[0] !== "__all__" &&
+                  selectedMCPDirectTool ? (
+                    <div className="flex-1 overflow-y-auto max-h-48 min-h-[44px] p-2 border border-gray-200 rounded-lg bg-gray-50/50">
+                      {(() => {
+                        const mcpTool = (serverToolsMap[selectedMCPServers[0]] || []).find(
+                          (t: any) => t.name === selectedMCPDirectTool,
+                        );
+                        return mcpTool ? (
+                          <MCPToolArgumentsForm
+                            ref={mcpToolArgsFormRef}
+                            tool={mcpTool}
+                            className="space-y-2"
+                          />
+                        ) : (
+                          <div className="flex items-center justify-center h-10 text-sm text-gray-500">
+                            Loading tool schema...
+                          </div>
+                        );
+                      })()}
+                    </div>
+                  ) : (
+                    <TextArea
+                      value={inputMessage}
+                      onChange={(e) => setInputMessage(e.target.value)}
+                      onKeyDown={handleKeyDown}
+                      placeholder={
+                        endpointType === EndpointType.CHAT ||
                         endpointType === EndpointType.EMBEDDINGS ||
                         endpointType === EndpointType.RESPONSES ||
                         endpointType === EndpointType.ANTHROPIC_MESSAGES
-                        ? "Type your message... (Shift+Enter for new line)"
-                        : endpointType === EndpointType.A2A_AGENTS
-                          ? "Send a message to the A2A agent..."
-                          : endpointType === EndpointType.IMAGE_EDITS
-                            ? "Describe how you want to edit the image..."
-                            : endpointType === EndpointType.SPEECH
-                              ? "Enter text to convert to speech..."
-                              : endpointType === EndpointType.TRANSCRIPTION
-                                ? "Optional: Add context or prompt for transcription..."
-                                : "Describe the image you want to generate..."
-                    }
-                    disabled={isLoading}
-                    className="flex-1"
-                    autoSize={{ minRows: 1, maxRows: 4 }}
-                    style={{
-                      resize: "none",
-                      border: "none",
-                      boxShadow: "none",
-                      background: "transparent",
-                      padding: "4px 0",
-                      fontSize: "14px",
-                      lineHeight: "20px",
-                    }}
-                  />
+                          ? "Type your message... (Shift+Enter for new line)"
+                          : endpointType === EndpointType.A2A_AGENTS
+                            ? "Send a message to the A2A agent..."
+                            : endpointType === EndpointType.IMAGE_EDITS
+                              ? "Describe how you want to edit the image..."
+                              : endpointType === EndpointType.SPEECH
+                                ? "Enter text to convert to speech..."
+                                : endpointType === EndpointType.TRANSCRIPTION
+                                  ? "Optional: Add context or prompt for transcription..."
+                                  : "Describe the image you want to generate..."
+                      }
+                      disabled={isLoading}
+                      className="flex-1"
+                      autoSize={{ minRows: 1, maxRows: 4 }}
+                      style={{
+                        resize: "none",
+                        border: "none",
+                        boxShadow: "none",
+                        background: "transparent",
+                        padding: "4px 0",
+                        fontSize: "14px",
+                        lineHeight: "20px",
+                      }}
+                    />
+                  )}
 
                   {/* Right: send button - matching blue theme */}
                   <TremorButton
                     onClick={handleSendMessage}
                     disabled={
-                      isLoading || (endpointType === EndpointType.TRANSCRIPTION ? !uploadedAudio : !inputMessage.trim())
+                      isLoading ||
+                      (endpointType === EndpointType.MCP
+                        ? !(
+                            selectedMCPServers.length === 1 &&
+                            selectedMCPServers[0] !== "__all__" &&
+                            selectedMCPDirectTool
+                          )
+                        : endpointType === EndpointType.TRANSCRIPTION
+                          ? !uploadedAudio
+                          : !inputMessage.trim())
                     }
                     className="flex-shrink-0 ml-2 !w-8 !h-8 !min-w-8 !p-0 !rounded-full !bg-blue-600 hover:!bg-blue-700 disabled:!bg-gray-300 !border-none !text-white disabled:!text-gray-500 !flex !items-center !justify-center"
                   >

ui/litellm-dashboard/src/components/playground/chat_ui/chatConstants.ts

@@ -43,4 +43,5 @@ export const ENDPOINT_OPTIONS = [
   { value: EndpointType.SPEECH, label: "/v1/audio/speech" },
   { value: EndpointType.TRANSCRIPTION, label: "/v1/audio/transcriptions" },
   { value: EndpointType.A2A_AGENTS, label: "/v1/a2a/message/send" },
+  { value: EndpointType.MCP, label: "/mcp-rest/tools/call" },
 ];

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

@@ -26,6 +26,7 @@ export enum EndpointType {
   SPEECH = "speech",
   TRANSCRIPTION = "transcription",
   A2A_AGENTS = "a2a_agents",
+  MCP = "mcp",
   // add additional endpoint types if required
 }
 

ui/litellm-dashboard/src/components/playground/llm_calls/a2a_send_message.tsx

@@ -23,6 +23,7 @@ interface A2AJsonRpcRequest {
   method: string;
   params: {
     message: A2AMessage;
+    metadata?: { guardrails?: string[] };
   };
 }
 
@@ -114,6 +115,7 @@ export const makeA2ASendMessageRequest = async (
   onTotalLatency?: (totalLatency: number) => void,
   onA2AMetadata?: (metadata: A2ATaskMetadata) => void,
   customBaseUrl?: string,
+  guardrails?: string[],
 ): Promise<void> => {
   const proxyBaseUrl = customBaseUrl || getProxyBaseUrl();
   const url = proxyBaseUrl
@@ -137,6 +139,10 @@ export const makeA2ASendMessageRequest = async (
     },
   };
 
+  if (guardrails && guardrails.length > 0) {
+    jsonRpcRequest.params.metadata = { guardrails };
+  }
+
   const startTime = performance.now();
 
   try {

ui/litellm-dashboard/src/contexts/AntdGlobalProvider.tsx

@@ -0,0 +1,24 @@
+"use client";
+
+import React, { useEffect, useRef } from "react";
+import { notification } from "antd";
+import { setNotificationInstance } from "@/components/molecules/notifications_manager";
+
+export default function AntdGlobalProvider({ children }: { children: React.ReactNode }) {
+  const [api, contextHolder] = notification.useNotification();
+  const initialized = useRef(false);
+
+  useEffect(() => {
+    if (!initialized.current) {
+      setNotificationInstance(api);
+      initialized.current = true;
+    }
+  }, [api]);
+
+  return (
+    <>
+      {contextHolder}
+      {children}
+    </>
+  );
+}