diff --git a/.circleci/config.yml b/.circleci/config.yml index 8672561f654..e171759f1c4 100644 --- a/.circleci/config.yml +++ b/.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 diff --git a/docs/my-website/docs/proxy/guardrails/custom_code_guardrail.md b/docs/my-website/docs/proxy/guardrails/custom_code_guardrail.md index cb246144497..8cbc247ae5e 100644 --- a/docs/my-website/docs/proxy/guardrails/custom_code_guardrail.md +++ b/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 diff --git a/litellm/a2a_protocol/__init__.py b/litellm/a2a_protocol/__init__.py index d8d349bb98a..85c03687e25 100644 --- a/litellm/a2a_protocol/__init__.py +++ b/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", ] diff --git a/litellm/a2a_protocol/card_resolver.py b/litellm/a2a_protocol/card_resolver.py index 7c4c5af149d..4c5dd3e3ba6 100644 --- a/litellm/a2a_protocol/card_resolver.py +++ b/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}. " diff --git a/litellm/a2a_protocol/exception_mapping_utils.py b/litellm/a2a_protocol/exception_mapping_utils.py new file mode 100644 index 00000000000..8463080b358 --- /dev/null +++ b/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, + ) diff --git a/litellm/a2a_protocol/exceptions.py b/litellm/a2a_protocol/exceptions.py new file mode 100644 index 00000000000..546b23105be --- /dev/null +++ b/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, + ) diff --git a/litellm/a2a_protocol/main.py b/litellm/a2a_protocol/main.py index b326f9e7ed5..642dfaf023c 100644 --- a/litellm/a2a_protocol/main.py +++ b/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( diff --git a/litellm/constants.py b/litellm/constants.py index 25decd363a6..3c618723d64 100644 --- a/litellm/constants.py +++ b/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 ### diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index 3e8f9bc337b..5e21ff9754f 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/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) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 1652ec2aa0c..bbd55a59bce 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/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", {}) diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index a99bd1cd0f3..6b9e51034c0 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/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: diff --git a/litellm/llms/a2a/chat/guardrail_translation/README.md b/litellm/llms/a2a/chat/guardrail_translation/README.md new file mode 100644 index 00000000000..1e18f5cda3a --- /dev/null +++ b/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 diff --git a/litellm/llms/a2a/chat/guardrail_translation/__init__.py b/litellm/llms/a2a/chat/guardrail_translation/__init__.py new file mode 100644 index 00000000000..13c20677485 --- /dev/null +++ b/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"] diff --git a/litellm/llms/a2a/chat/guardrail_translation/handler.py b/litellm/llms/a2a/chat/guardrail_translation/handler.py new file mode 100644 index 00000000000..770453f2def --- /dev/null +++ b/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 diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 0da47634a94..4704549e716 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.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, diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 7e70b5baae4..786cfbfb008 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -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 diff --git a/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py b/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py index 4d53ae7059d..14bbb82808d 100644 --- a/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py +++ b/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", ) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 187f0e3a3c5..e7174f943b2 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -74,6 +74,7 @@ try: ) 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 +476,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 +1956,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 +1996,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 @@ -2301,28 +2303,24 @@ class MCPServerManager: non-public servers are hidden from external IPs. """ registry = self.get_registry() - # Pass 1: Match by alias (highest priority) for server in registry.values(): if server.alias == server_name: if not self._is_server_accessible_from_ip(server, client_ip): return None return server - # Pass 2: Match by server_name for server in registry.values(): if server.server_name == server_name: if not self._is_server_accessible_from_ip(server, client_ip): return None return server - # Pass 3: Match by name (lowest priority) for server in registry.values(): if server.name == server_name: if not self._is_server_accessible_from_ip(server, client_ip): return None return server - return None def get_filtered_registry( diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 940389b1935..2cd56e0ff3f 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/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 diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 5c51254e9b0..890c4ae8fb2 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -725,7 +725,6 @@ if MCP_AVAILABLE: def _get_client_ip_from_context() -> Optional[str]: """ Extract client_ip from auth context. - Returns None if context not set (caller should handle this as "no IP filtering"). """ try: @@ -748,7 +747,6 @@ if MCP_AVAILABLE: mcp_servers: Optional list of server names to filter to. client_ip: Client IP for IP-based access control. If None, falls back to auth context. Pass explicitly from request handlers for safety. - Note: If client_ip is None and auth context is not set, IP filtering is skipped. This is intentional for internal callers but may indicate a bug if called from a request handler without proper context setup. diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 6fc2ad3f166..6e4131c3ced 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2173,6 +2173,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=()) diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 24727aacd75..12b2d5c4dfe 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/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, ) diff --git a/litellm/proxy/auth/ip_address_utils.py b/litellm/proxy/auth/ip_address_utils.py index 4dbc9d49cc3..651d6785333 100644 --- a/litellm/proxy/auth/ip_address_utils.py +++ b/litellm/proxy/auth/ip_address_utils.py @@ -49,7 +49,6 @@ class IPAddressUtils: ) -> List[Union[ipaddress.IPv4Network, ipaddress.IPv6Network]]: """ Parse trusted proxy CIDR ranges for XFF validation. - Returns empty list if not configured (XFF will not be trusted). """ if not configured_ranges: @@ -153,5 +152,4 @@ class IPAddressUtils: "XFF header from untrusted IP %s, ignoring", direct_ip ) return direct_ip - return _get_request_ip_address(request, use_x_forwarded_for=use_xff) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 769c250d9fc..f33b2412260 100644 --- a/litellm/proxy/common_request_processing.py +++ b/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 diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py index a0ca324411c..68f9dfd7abc 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py +++ b/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, diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py index 695e59977c8..de7690635d8 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py +++ b/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, diff --git a/litellm/proxy/guardrails/guardrail_hooks/pangea/pangea.py b/litellm/proxy/guardrails/guardrail_hooks/pangea/pangea.py index ca3e6a90801..55c7e72c36f 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/pangea/pangea.py +++ b/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): diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index f07f65d10f5..cc05358baf7 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/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 diff --git a/litellm/proxy/management_endpoints/types.py b/litellm/proxy/management_endpoints/types.py index ad2ad0a5fe5..a35fc4a5f3f 100644 --- a/litellm/proxy/management_endpoints/types.py +++ b/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) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 2d248dc81f3..278f3bdaafd 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/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( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7c5e25cf324..cf9e00512b1 100644 --- a/litellm/proxy/proxy_server.py +++ b/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, diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index 6baa2047d73..e941964644e 100644 --- a/litellm/proxy/route_llm_request.py +++ b/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 [ diff --git a/litellm/responses/litellm_completion_transformation/handler.py b/litellm/responses/litellm_completion_transformation/handler.py index 7e0b4cfa243..74ad6675e1d 100644 --- a/litellm/responses/litellm_completion_transformation/handler.py +++ b/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)}") diff --git a/litellm/router.py b/litellm/router.py index 6dc3278225e..374f80db361 100644 --- a/litellm/router.py +++ b/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 diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 0da47634a94..4704549e716 100644 --- a/model_prices_and_context_window.json +++ b/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, diff --git a/tests/agent_tests/local_vertex_agent.py b/tests/agent_tests/local_only_agent_tests/local_vertex_agent.py similarity index 100% rename from tests/agent_tests/local_vertex_agent.py rename to tests/agent_tests/local_only_agent_tests/local_vertex_agent.py diff --git a/tests/agent_tests/test_a2a.py b/tests/agent_tests/local_only_agent_tests/test_a2a.py similarity index 100% rename from tests/agent_tests/test_a2a.py rename to tests/agent_tests/local_only_agent_tests/test_a2a.py diff --git a/tests/agent_tests/test_a2a_completion_bridge.py b/tests/agent_tests/local_only_agent_tests/test_a2a_completion_bridge.py similarity index 100% rename from tests/agent_tests/test_a2a_completion_bridge.py rename to tests/agent_tests/local_only_agent_tests/test_a2a_completion_bridge.py diff --git a/tests/agent_tests/test_a2a_agent.py b/tests/agent_tests/test_a2a_agent.py new file mode 100644 index 00000000000..acaf615ec88 --- /dev/null +++ b/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)}") diff --git a/tests/mcp_tests/test_mcp_client_unit.py b/tests/mcp_tests/test_mcp_client_unit.py index 9533e56bcc8..c70d0c42cd8 100644 --- a/tests/mcp_tests/test_mcp_client_unit.py +++ b/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() diff --git a/tests/test_litellm/a2a_protocol/test_card_resolver.py b/tests/test_litellm/a2a_protocol/test_card_resolver.py index a1bd33107b6..1bdab50860c 100644 --- a/tests/test_litellm/a2a_protocol/test_card_resolver.py +++ b/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/" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index e1e4b3a8b6d..68afe784988 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -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""" diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index 41096503a2e..16f80826798 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/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 diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 69cf8240c63..7bebe00d61e 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/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): """ diff --git a/tests/test_litellm/proxy/test_route_llm_request.py b/tests/test_litellm/proxy/test_route_llm_request.py index 90eace63714..1283d2ccbe7 100644 --- a/tests/test_litellm/proxy/test_route_llm_request.py +++ b/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 diff --git a/ui/litellm-dashboard/src/app/layout.tsx b/ui/litellm-dashboard/src/app/layout.tsx index 95c485fe2f0..1233da9046f 100644 --- a/ui/litellm-dashboard/src/app/layout.tsx +++ b/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 ( - {children} + + {children} + ); } diff --git a/ui/litellm-dashboard/src/components/guardrails/custom_code/CustomCodeModal.tsx b/ui/litellm-dashboard/src/components/guardrails/custom_code/CustomCodeModal.tsx index 73d96c590ec..1b8a8c034bd 100644 --- a/ui/litellm-dashboard/src/components/guardrails/custom_code/CustomCodeModal.tsx +++ b/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 = ({ @@ -125,16 +168,72 @@ const CustomCodeModal: React.FC = ({ onClose, onSuccess, accessToken, + editData, }) => { + const isEditMode = !!editData; const [guardrailName, setGuardrailName] = useState(""); - const [mode, setMode] = useState("pre_call"); + const [mode, setMode] = useState(["pre_call"]); const [defaultOn, setDefaultOn] = useState(false); const [selectedTemplate, setSelectedTemplate] = useState("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(null); const [copiedPrimitive, setCopiedPrimitive] = useState(null); const textareaRef = useRef(null); @@ -145,18 +244,35 @@ const CustomCodeModal: React.FC = ({ 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 = ({ } }; - // 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 = ({ 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 = ({ 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 = ({ 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 = ({
{/* Header */}
-

Create Custom Guardrail

+

+ {isEditMode ? "Edit Custom Guardrail" : "Create Custom Guardrail"} +

Define custom logic using Python-like syntax

@@ -311,14 +467,16 @@ const CustomCodeModal: React.FC = ({ placeholder="e.g., block-pii-custom" />
-
- +
+