Agent Guardrails - working on output + streaming output (#20636)

* Add http support to custom code guardrails + Unified guardrails for MCP + Agent guardrail support (#20619)

* fix: fix styling

* fix(custom_code_guardrail.py): add http support for custom code guardrails

allows users to call external guardrails on litellm with minimal code changes (no custom handlers)

Test guardrail integrations more easily

* feat(a2a/): add guardrails for agent interactions

allows the same guardrails for llm's to be applied to agents as well

* fix(a2a/): support passing guardrails to a2a from the UI

* style(code-editor): allow editing custom code guardrails on ui + add examples of pre/post calls for custom code guardrails

* feat(mcp/): support custom code guardrails for mcp calls

allows custom code guardrails to work on mcp input

* feat(chatui.tsx): support guardrails on mcp tool calls on playground

* fix(ui/): add mcp input as an example for custom code guardrails

* feat(a2a/): ensure a2a guardrails works on response output

* feat(a2a/): support streaming guardrails

* test: address greptile comments

* test: address greptile comments

litellm/cost_calculator.py

@@ -73,6 +73,7 @@ from litellm.llms.vertex_ai.cost_calculator import (
 from litellm.llms.vertex_ai.cost_calculator import cost_router as google_cost_router
 from litellm.llms.xai.cost_calculator import cost_per_token as xai_cost_per_token
 from litellm.responses.utils import ResponseAPILoggingUtils
+from litellm.types.agents import LiteLLMSendMessageResponse
 from litellm.types.llms.openai import (
     HttpxBinaryResponseContent,
     ImageGenerationRequestQuality,
@@ -149,32 +150,33 @@ def _get_additional_costs(
 ) -> Optional[dict]:
     """
     Calculate additional costs beyond standard token costs.
-    
+
     This function delegates to provider-specific config classes to calculate
     any additional costs like routing fees, infrastructure costs, etc.
-    
+
     Args:
         model: The model name
         custom_llm_provider: The provider name (optional)
         prompt_tokens: Number of prompt tokens
         completion_tokens: Number of completion tokens
-        
+
     Returns:
         Optional dictionary with cost names and amounts, or None if no additional costs
     """
     if not custom_llm_provider:
         return None
-        
+
     try:
         config_class = None
         if custom_llm_provider == "azure_ai":
             from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo
+
             config_class = AzureFoundryModelInfo.get_azure_ai_config_for_model(model)
         # Add more providers here as needed
         # elif custom_llm_provider == "other_provider":
         #     config_class = get_other_provider_config(model)
-        
-        if config_class and hasattr(config_class, 'calculate_additional_costs'):
+
+        if config_class and hasattr(config_class, "calculate_additional_costs"):
             return config_class.calculate_additional_costs(
                 model=model,
                 prompt_tokens=prompt_tokens,
@@ -182,7 +184,7 @@ def _get_additional_costs(
             )
     except Exception as e:
         verbose_logger.debug(f"Error calculating additional costs: {e}")
-    
+
     return None
 
 
@@ -747,6 +749,8 @@ def _infer_call_type(
         return "image_generation"
     elif isinstance(completion_response, TextCompletionResponse):
         return "text_completion"
+    elif isinstance(completion_response, LiteLLMSendMessageResponse):
+        return "send_message"
 
     return call_type
 
@@ -1030,9 +1034,9 @@ def completion_cost(  # noqa: PLR0915
                     or isinstance(completion_response, dict)
                 ):  # tts returns a custom class
                     if isinstance(completion_response, dict):
-                        usage_obj: Optional[
-                            Union[dict, Usage]
-                        ] = completion_response.get("usage", {})
+                        usage_obj: Optional[Union[dict, Usage]] = (
+                            completion_response.get("usage", {})
+                        )
                     else:
                         usage_obj = getattr(completion_response, "usage", {})
                     if isinstance(usage_obj, BaseModel) and not _is_known_usage_objects(
@@ -1386,7 +1390,7 @@ def completion_cost(  # noqa: PLR0915
                     service_tier=service_tier,
                     response=completion_response,
                 )
-                
+
                 # Get additional costs from provider (e.g., routing fees, infrastructure costs)
                 additional_costs = _get_additional_costs(
                     model=model,
@@ -1394,7 +1398,7 @@ def completion_cost(  # noqa: PLR0915
                     prompt_tokens=prompt_tokens,
                     completion_tokens=completion_tokens,
                 )
-                
+
                 _final_cost = (
                     prompt_tokens_cost_usd_dollar + completion_tokens_cost_usd_dollar
                 )

litellm/litellm_core_utils/api_route_to_call_types.py

@@ -3,6 +3,9 @@ Dictionary mapping API routes to their corresponding CallTypes in LiteLLM.
 
 This dictionary maps each API endpoint to the CallTypes that can be used for that route.
 Each route can have both async (prefixed with 'a') and sync call types.
+
+Route patterns may contain placeholders like {agent_id}, {model}, {batch_id}; these
+match a single path segment when resolving call types for a concrete path.
 """
 
 from typing import List, Optional
@@ -10,17 +13,43 @@ from typing import List, Optional
 from litellm.types.utils import API_ROUTE_TO_CALL_TYPES, CallTypes
 
 
+def _route_matches_pattern(route: str, pattern: str) -> bool:
+    """
+    Return True if the concrete route matches the pattern.
+    Pattern segments like {param} match any single path segment.
+    """
+    route_parts = route.strip("/").split("/")
+    pattern_parts = pattern.strip("/").split("/")
+    if len(route_parts) != len(pattern_parts):
+        return False
+    for r, p in zip(route_parts, pattern_parts):
+        if p.startswith("{") and p.endswith("}"):
+            continue
+        if r != p:
+            return False
+    return True
+
+
 def get_call_types_for_route(route: str) -> Optional[List[CallTypes]]:
     """
     Get the list of CallTypes for a given API route.
 
+    Supports both exact keys and dynamic patterns (e.g. /a2a/my-agent/message/send
+    matches /a2a/{agent_id}/message/send).
+
     Args:
-        route: API route path (e.g., "/chat/completions")
+        route: API route path (e.g., "/chat/completions" or "/a2a/my-pydantic-agent/message/send")
 
     Returns:
         List of CallTypes for that route, or None if route not found
     """
-    return API_ROUTE_TO_CALL_TYPES.get(route, None)
+    exact = API_ROUTE_TO_CALL_TYPES.get(route, None)
+    if exact is not None:
+        return exact
+    for pattern, call_types in API_ROUTE_TO_CALL_TYPES.items():
+        if _route_matches_pattern(route, pattern):
+            return call_types
+    return None
 
 
 def get_routes_for_call_type(call_type: CallTypes) -> list:

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

@@ -10,6 +10,7 @@ A2A Protocol Format:
 - Output: JSON-RPC 2.0 with result containing message/artifact parts
 """
 
+import json
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
 
 from litellm._logging import verbose_proxy_logger
@@ -206,6 +207,132 @@ class A2AGuardrailHandler(BaseTranslation):
             response["result"] = result
             return response
 
+    async def process_output_streaming_response(
+        self,
+        responses_so_far: List[Any],
+        guardrail_to_apply: "CustomGuardrail",
+        litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None,
+        user_api_key_dict: Optional["UserAPIKeyAuth"] = None,
+    ) -> List[Any]:
+        """
+        Process A2A streaming output by applying guardrails to accumulated text.
+
+        IN-PLACE MODIFICATION: This method mutates responses_so_far in place. Callers
+        must deep-copy the current chunk before calling if they intend to yield that
+        chunk, because we put the full guardrailed text in the first chunk and clear
+        all subsequent text parts to "".
+
+        Algorithm:
+        1. Parse each item (dict or NDJSON str) and collect text from result.artifact.parts,
+           result.message.parts, result.parts, etc.
+        2. Concatenate all texts in order, apply guardrail once to the combined string.
+        3. Write the full guardrailed text into the FIRST chunk's first text part.
+        4. Clear all other text parts in all chunks to "" (in-place).
+
+        responses_so_far: List of JSON-RPC 2.0 objects (dict or NDJSON str).
+        Returns: The same list (modified in place).
+        """
+        from litellm.llms.a2a.common_utils import extract_text_from_a2a_response
+
+        # Parse each item; keep alignment with responses_so_far (None where unparseable)
+        parsed: List[Optional[Dict[str, Any]]] = [None] * len(responses_so_far)
+        for i, item in enumerate(responses_so_far):
+            if isinstance(item, dict):
+                obj = item
+            elif isinstance(item, str):
+                try:
+                    obj = json.loads(item.strip())
+                except (json.JSONDecodeError, TypeError):
+                    continue
+            else:
+                continue
+            if isinstance(obj.get("result"), dict):
+                parsed[i] = obj
+
+        valid_parsed = [(i, obj) for i, obj in enumerate(parsed) if obj is not None]
+        if not valid_parsed:
+            return responses_so_far
+
+        # Collect text from each chunk in order (by original index in responses_so_far)
+        text_parts: List[str] = []
+        chunk_indices_with_text: List[int] = []  # indices into valid_parsed
+        for idx, (orig_i, obj) in enumerate(valid_parsed):
+            t = extract_text_from_a2a_response(obj)
+            if t:
+                text_parts.append(t)
+                chunk_indices_with_text.append(orig_i)
+
+        combined_text = "".join(text_parts)
+        if not combined_text:
+            return responses_so_far
+
+        request_data: dict = {"responses_so_far": responses_so_far}
+        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=[combined_text])
+        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", [])
+        if not guardrailed_texts or len(guardrailed_texts) != 1:
+            # Guardrail should return exactly one text for combined input
+            verbose_proxy_logger.warning(
+                "A2A streaming guardrail returned unexpected texts count: %s",
+                len(guardrailed_texts or []),
+            )
+            return responses_so_far
+        guardrailed_text = guardrailed_texts[0]
+
+        # Find first chunk (by original index) that has text; put full guardrailed text there and clear rest
+        first_chunk_with_text: Optional[int] = (
+            chunk_indices_with_text[0] if chunk_indices_with_text else None
+        )
+
+        for orig_i, obj in valid_parsed:
+            result = obj.get("result", {})
+            if not isinstance(result, dict):
+                continue
+            texts_in_chunk: List[str] = []
+            mappings: List[Tuple[Tuple[str, ...], int]] = []
+            self._extract_texts_from_result(
+                result=result,
+                texts_to_check=texts_in_chunk,
+                task_mappings=mappings,
+            )
+            if not mappings:
+                continue
+            if orig_i == first_chunk_with_text:
+                # Put full guardrailed text in first text part; clear others in this chunk
+                for task_idx, (path, part_idx) in enumerate(mappings):
+                    text = guardrailed_text if task_idx == 0 else ""
+                    self._apply_text_to_path(
+                        result=result,
+                        path=path,
+                        part_idx=part_idx,
+                        text=text,
+                    )
+            else:
+                # Clear all text parts in non-first chunks (in-place)
+                for path, part_idx in mappings:
+                    self._apply_text_to_path(
+                        result=result,
+                        path=path,
+                        part_idx=part_idx,
+                        text="",
+                    )
+
+        # Write back to responses_so_far where we had NDJSON strings
+        for i, item in enumerate(responses_so_far):
+            if isinstance(item, str) and parsed[i] is not None:
+                responses_so_far[i] = json.dumps(parsed[i]) + "\n"
+
+        return responses_so_far
+
     def _extract_texts_from_result(
         self,
         result: Dict[str, Any],
@@ -301,15 +428,21 @@ class A2AGuardrailHandler(BaseTranslation):
         part_idx: int,
         text: str,
     ) -> None:
-        """Apply guardrailed text back to the specified path in the result."""
-        # Navigate to the parts list
-        current = result
+        """Apply guardrailed text back to the specified path in the result (in-place)."""
+        current: Any = 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
+        if not isinstance(current, list) or part_idx >= len(current):
+            verbose_proxy_logger.warning(
+                "A2A _apply_text_to_path: invalid path or index path=%s part_idx=%s",
+                path,
+                part_idx,
+            )
+            return
+        part = current[part_idx]
+        if isinstance(part, dict) and "text" in part:
+            part["text"] = text

litellm/proxy/_experimental/mcp_server/mcp_server_manager.py

@@ -72,6 +72,7 @@ try:
     from mcp.shared.tool_name_validation import (
         SEP_986_URL,
     )
+    from mcp.shared.tool_name_validation import SEP_986_URL
 except ImportError:
     from pydantic import BaseModel
 

litellm/proxy/_experimental/mcp_server/rest_endpoints.py

@@ -381,43 +381,10 @@ if MCP_AVAILABLE:
             if "metadata" in data and "user_api_key_auth" in data["metadata"]:
                 data["user_api_key_auth"] = data["metadata"]["user_api_key_auth"]
 
-            # Get all auth contexts
-            auth_contexts = await build_effective_auth_contexts(user_api_key_dict)
-
-            # Collect allowed server IDs from all contexts, then apply IP filtering
-            _rest_client_ip = IPAddressUtils.get_mcp_client_ip(request)
-            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)
-
-            allowed_server_ids_set = set(
-                global_mcp_server_manager.filter_server_ids_by_ip(
-                    list(allowed_server_ids_set), _rest_client_ip
-                )
+            allowed_mcp_servers = await _resolve_allowed_mcp_servers_for_tool_call(
+                user_api_key_dict, server_id
             )
 
-            # Check if the specified server_id is allowed
-            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}",
-                    },
-                )
-
-            # Build allowed_mcp_servers list (only include allowed servers)
-            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)
-
             # Call execute_mcp_tool directly (permission checks already done)
             result = await execute_mcp_tool(
                 name=tool_name,

litellm/proxy/agent_endpoints/a2a_endpoints.py

@@ -6,7 +6,7 @@ The A2A SDK can point to LiteLLM's URL and invoke agents registered with LiteLLM
 """
 
 import json
-from typing import Optional
+from typing import Any, Optional
 
 from fastapi import APIRouter, Depends, HTTPException, Request, Response
 from fastapi.responses import JSONResponse, StreamingResponse
@@ -54,14 +54,21 @@ async def _handle_stream_message(
     agent_id: Optional[str] = None,
     metadata: Optional[dict] = None,
     proxy_server_request: Optional[dict] = None,
+    *,
+    user_api_key_dict: Optional[UserAPIKeyAuth] = None,
+    request_data: Optional[dict] = None,
+    proxy_logging_obj: Optional[Any] = None,
 ) -> StreamingResponse:
-    """Handle message/stream method via SDK functions."""
+    """Handle message/stream method via SDK functions.
+
+    When user_api_key_dict, request_data, and proxy_logging_obj are provided,
+    uses common_request_processing.async_streaming_data_generator with NDJSON
+    serializers so proxy hooks and cost injection apply.
+    """
     from litellm.a2a_protocol import asend_message_streaming
     from litellm.a2a_protocol.main import A2A_SDK_AVAILABLE
 
-    # Check is handled in invoke_agent_a2a, but if called directly:
     if not A2A_SDK_AVAILABLE:
-        # Return a streaming response that yields an error
         async def _error_stream():
             yield json.dumps(
                 {
@@ -78,29 +85,82 @@ async def _handle_stream_message(
 
     from a2a.types import MessageSendParams, SendStreamingMessageRequest
 
+    use_proxy_hooks = (
+        user_api_key_dict is not None
+        and request_data is not None
+        and proxy_logging_obj is not None
+    )
+
     async def stream_response():
         try:
             a2a_request = SendStreamingMessageRequest(
                 id=request_id,
                 params=MessageSendParams(**params),
             )
-            async for chunk in asend_message_streaming(
+            a2a_stream = asend_message_streaming(
                 request=a2a_request,
                 api_base=api_base,
                 litellm_params=litellm_params,
                 agent_id=agent_id,
                 metadata=metadata,
                 proxy_server_request=proxy_server_request,
-            ):
-                # Chunk may be dict or object depending on bridge vs standard path
-                if hasattr(chunk, "model_dump"):
-                    yield json.dumps(
-                        chunk.model_dump(mode="json", exclude_none=True)
+            )
+
+            if use_proxy_hooks and user_api_key_dict is not None and request_data is not None and proxy_logging_obj is not None:
+                from litellm.proxy.common_request_processing import (
+                    ProxyBaseLLMRequestProcessing,
+                )
+
+                def _ndjson_chunk(chunk: Any) -> str:
+                    if hasattr(chunk, "model_dump"):
+                        obj = chunk.model_dump(mode="json", exclude_none=True)
+                    else:
+                        obj = chunk
+                    return json.dumps(obj) + "\n"
+
+                def _ndjson_error(proxy_exc: Any) -> str:
+                    return json.dumps(
+                        {
+                            "jsonrpc": "2.0",
+                            "id": request_id,
+                            "error": {
+                                "code": -32603,
+                                "message": getattr(
+                                    proxy_exc, "message", f"Streaming error: {proxy_exc!s}"
+                                ),
+                            },
+                        }
                     ) + "\n"
-                else:
-                    yield json.dumps(chunk) + "\n"
+
+                async for line in ProxyBaseLLMRequestProcessing.async_streaming_data_generator(
+                    response=a2a_stream,
+                    user_api_key_dict=user_api_key_dict,
+                    request_data=request_data,
+                    proxy_logging_obj=proxy_logging_obj,
+                    serialize_chunk=_ndjson_chunk,
+                    serialize_error=_ndjson_error,
+                ):
+                    yield line
+            else:
+                async for chunk in a2a_stream:
+                    if hasattr(chunk, "model_dump"):
+                        yield json.dumps(
+                            chunk.model_dump(mode="json", exclude_none=True)
+                        ) + "\n"
+                    else:
+                        yield json.dumps(chunk) + "\n"
         except Exception as e:
             verbose_proxy_logger.exception(f"Error streaming A2A response: {e}")
+            if use_proxy_hooks and proxy_logging_obj is not None and user_api_key_dict is not None and request_data is not None:
+                transformed_exception = await proxy_logging_obj.post_call_failure_hook(
+                    user_api_key_dict=user_api_key_dict,
+                    original_exception=e,
+                    request_data=request_data,
+                )
+                if transformed_exception is not None:
+                    e = transformed_exception
+            if isinstance(e, HTTPException):
+                raise
             yield json.dumps(
                 {
                     "jsonrpc": "2.0",
@@ -323,8 +383,18 @@ async def invoke_agent_a2a(
                 metadata=data.get("metadata", {}),
                 proxy_server_request=data.get("proxy_server_request"),
             )
+
+            response = await proxy_logging_obj.post_call_success_hook(
+                user_api_key_dict=user_api_key_dict,
+                data=data,
+                response=response,
+            )
             return JSONResponse(
-                content=response.model_dump(mode="json", exclude_none=True)
+                content=(
+                    response.model_dump(mode="json", exclude_none=True)  # type: ignore
+                    if hasattr(response, "model_dump")
+                    else response
+                )
             )
 
         elif method == "message/stream":
@@ -336,6 +406,9 @@ async def invoke_agent_a2a(
                 agent_id=agent.agent_id,
                 metadata=data.get("metadata", {}),
                 proxy_server_request=data.get("proxy_server_request"),
+                user_api_key_dict=user_api_key_dict,
+                request_data=data,
+                proxy_logging_obj=proxy_logging_obj,
             )
         else:
             return _jsonrpc_error(request_id, -32601, f"Method '{method}' not found")

litellm/proxy/common_request_processing.py

@@ -43,6 +43,11 @@ from litellm.proxy.utils import ProxyLogging
 from litellm.router import Router
 from litellm.types.utils import ServerToolUse
 
+# Type alias for streaming chunk serializer (chunk after hooks + cost injection -> wire format)
+StreamChunkSerializer = Callable[[Any], str]
+# Type alias for streaming error serializer (ProxyException -> wire format)
+StreamErrorSerializer = Callable[[ProxyException], str]
+
 if TYPE_CHECKING:
     from litellm.proxy.proxy_server import ProxyConfig as _ProxyConfig
 
@@ -1193,22 +1198,24 @@ class ProxyBaseLLMRequestProcessing:
             return chunk
 
     @staticmethod
-    async def async_sse_data_generator(
-        response,
+    async def async_streaming_data_generator(
+        response: Any,
         user_api_key_dict: UserAPIKeyAuth,
         request_data: dict,
         proxy_logging_obj: ProxyLogging,
-    ):
+        *,
+        serialize_chunk: StreamChunkSerializer,
+        serialize_error: StreamErrorSerializer,
+    ) -> AsyncGenerator[str, None]:
         """
-        Anthropic /messages and Google /generateContent streaming data generator require SSE events
+        Shared streaming data generator: runs proxy iterator hook, per-chunk hook,
+        cost injection, then yields chunks via serialize_chunk; on exception runs
+        failure hook and yields via serialize_error. Use for SSE or NDJSON.
         """
-
         verbose_proxy_logger.debug("inside generator")
         try:
             str_so_far = ""
-            async for (
-                chunk
-            ) in proxy_logging_obj.async_post_call_streaming_iterator_hook(
+            async for chunk in proxy_logging_obj.async_post_call_streaming_iterator_hook(
                 user_api_key_dict=user_api_key_dict,
                 response=response,
                 request_data=request_data,
@@ -1216,7 +1223,6 @@ class ProxyBaseLLMRequestProcessing:
                 verbose_proxy_logger.debug(
                     "async_data_generator: received streaming chunk - {}".format(chunk)
                 )
-                ### CALL HOOKS ### - modify outgoing data
                 chunk = await proxy_logging_obj.async_post_call_streaming_hook(
                     user_api_key_dict=user_api_key_dict,
                     response=chunk,
@@ -1227,30 +1233,34 @@ class ProxyBaseLLMRequestProcessing:
                 if isinstance(chunk, (ModelResponse, ModelResponseStream)):
                     response_str = litellm.get_response_string(response_obj=chunk)
                     str_so_far += response_str
+                elif hasattr(chunk, "model_dump"):
+                    try:
+                        d = chunk.model_dump(mode="json", exclude_none=True)
+                        if isinstance(d, dict):
+                            str_so_far += str(d.get("content", ""))
+                    except Exception:
+                        pass
+                elif isinstance(chunk, dict):
+                    str_so_far += str(chunk.get("content", ""))
 
-                # Inject cost into Anthropic-style SSE usage for /v1/messages for any provider
                 model_name = request_data.get("model", "")
                 chunk = (
                     ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection(
                         chunk, model_name
                     )
                 )
-
-                # Format chunk using helper function
-                yield ProxyBaseLLMRequestProcessing.return_sse_chunk(chunk)
+                yield serialize_chunk(chunk)
         except Exception as e:
             verbose_proxy_logger.exception(
                 "litellm.proxy.proxy_server.async_data_generator(): Exception occured - {}".format(
                     str(e)
                 )
             )
-            # Allow callbacks to transform the error response
             transformed_exception = await proxy_logging_obj.post_call_failure_hook(
                 user_api_key_dict=user_api_key_dict,
                 original_exception=e,
                 request_data=request_data,
             )
-            # Use transformed exception if callback returned one, otherwise use original
             if transformed_exception is not None:
                 e = transformed_exception
             verbose_proxy_logger.debug(
@@ -1259,18 +1269,36 @@ class ProxyBaseLLMRequestProcessing:
 
             if isinstance(e, HTTPException):
                 raise e
-            else:
-                error_traceback = traceback.format_exc()
-                error_msg = f"{str(e)}\n\n{error_traceback}"
-
+            error_traceback = traceback.format_exc()
+            error_msg = f"{str(e)}\n\n{error_traceback}"
             proxy_exception = ProxyException(
                 message=getattr(e, "message", error_msg),
                 type=getattr(e, "type", "None"),
                 param=getattr(e, "param", "None"),
                 code=getattr(e, "status_code", 500),
             )
-            error_returned = json.dumps({"error": proxy_exception.to_dict()})
-            yield f"{STREAM_SSE_DATA_PREFIX}{error_returned}\n\n"
+            yield serialize_error(proxy_exception)
+
+    @staticmethod
+    async def async_sse_data_generator(
+        response: Any,
+        user_api_key_dict: UserAPIKeyAuth,
+        request_data: dict,
+        proxy_logging_obj: ProxyLogging,
+    ) -> AsyncGenerator[str, None]:
+        """
+        Anthropic /messages and Google /generateContent streaming data generator require SSE events.
+        Delegates to async_streaming_data_generator with SSE serializers.
+        """
+        async for chunk in ProxyBaseLLMRequestProcessing.async_streaming_data_generator(
+            response=response,
+            user_api_key_dict=user_api_key_dict,
+            request_data=request_data,
+            proxy_logging_obj=proxy_logging_obj,
+            serialize_chunk=ProxyBaseLLMRequestProcessing.return_sse_chunk,
+            serialize_error=lambda proxy_exc: f"{STREAM_SSE_DATA_PREFIX}{json.dumps({'error': proxy_exc.to_dict()})}\n\n",
+        ):
+            yield chunk
 
     @staticmethod
     def _process_chunk_with_cost_injection(chunk: Any, model_name: str) -> Any:

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

@@ -7,8 +7,11 @@ Unified Guardrail, leveraging LiteLLM's /applyGuardrail endpoint
 """
 
 import copy
+import json
 from typing import Any, AsyncGenerator, List, Optional, Union
 
+from fastapi import HTTPException
+
 from litellm._logging import verbose_proxy_logger
 from litellm.caching.caching import DualCache
 from litellm.cost_calculator import _infer_call_type
@@ -18,9 +21,80 @@ from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_fo
 from litellm.llms import load_guardrail_translation_mappings
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.types.guardrails import GuardrailEventHooks
-from litellm.types.utils import CallTypes, CallTypesLiteral, ModelResponseStream
+from litellm.types.utils import CallTypes, CallTypesLiteral
+
+# Call types that use NDJSON streaming (A2A); guardrail HTTPException is emitted as in-stream error
+A2A_CALL_TYPES = (CallTypes.asend_message, CallTypes.send_message)
 
 GUARDRAIL_NAME = "unified_llm_guardrails"
+
+
+def _get_a2a_request_id(
+    responses_so_far: List[Any], request_data: dict
+) -> Optional[str]:
+    """Get JSON-RPC request id from first A2A chunk or request body for in-stream error reporting."""
+    for item in responses_so_far:
+        if isinstance(item, dict) and "id" in item:
+            return item.get("id")
+        if isinstance(item, str):
+            try:
+                obj = json.loads(item.strip())
+                if isinstance(obj, dict) and "id" in obj:
+                    return obj.get("id")
+            except (json.JSONDecodeError, TypeError):
+                continue
+    body = request_data.get("body") or request_data.get("data") or {}
+    if isinstance(body, dict):
+        return body.get("id")
+    return None
+
+
+def _format_a2a_guardrail_error_chunk(
+    e: HTTPException,
+    call_type: Optional[str],
+    responses_so_far: List[Any],
+    request_data: dict,
+) -> Optional[str]:
+    """
+    Format HTTPException from guardrail as JSON-RPC 2.0 error chunk for A2A streaming.
+
+    When the response has already started, we cannot send an HTTP 4xx. For A2A (NDJSON)
+    streams, we yield an in-stream JSON-RPC error so the client receives the rejection.
+
+    Returns:
+        JSON-RPC error chunk string with trailing newline if call_type is A2A, else None.
+        Caller should yield the result when not None; otherwise re-raise the exception.
+    """
+    if call_type is None or CallTypes(call_type) not in A2A_CALL_TYPES:
+        return None
+    request_id = _get_a2a_request_id(responses_so_far, request_data)
+    detail = (
+        e.detail
+        if isinstance(e.detail, dict)
+        else {"message": str(e.detail)}
+    )
+    return (
+        json.dumps(
+            {
+                "jsonrpc": "2.0",
+                "id": request_id,
+                "error": {
+                    "code": -32603,
+                    "message": detail.get(
+                        "error", detail.get("message", str(e.detail))
+                    ),
+                    "data": {
+                        k: v
+                        for k, v in detail.items()
+                        if k not in ("error", "message")
+                    },
+                },
+            }
+        )
+        + "\n"
+    )
+
+
 endpoint_guardrail_translation_mappings = None
 
 
@@ -233,7 +307,7 @@ class UnifiedLLMGuardrails(CustomLogger):
         user_api_key_dict: UserAPIKeyAuth,
         response: Any,
         request_data: dict,
-    ) -> AsyncGenerator[ModelResponseStream, None]:
+    ) -> AsyncGenerator[Any, None]:
         """
         Passes the entire stream to the guardrail
 
@@ -364,13 +438,21 @@ class UnifiedLLMGuardrails(CustomLogger):
                     CallTypes(call_type)
                 ]()
 
-                await endpoint_translation.process_output_streaming_response(
-                    responses_so_far=responses_so_far,
-                    guardrail_to_apply=guardrail_to_apply,
-                    litellm_logging_obj=request_data.get("litellm_logging_obj"),
-                    user_api_key_dict=user_api_key_dict,
-                )
-
+                try:
+                    await endpoint_translation.process_output_streaming_response(
+                        responses_so_far=responses_so_far,
+                        guardrail_to_apply=guardrail_to_apply,
+                        litellm_logging_obj=request_data.get("litellm_logging_obj"),
+                        user_api_key_dict=user_api_key_dict,
+                    )
+                except HTTPException as e:
+                    error_chunk = _format_a2a_guardrail_error_chunk(
+                        e, call_type, responses_so_far, request_data
+                    )
+                    if error_chunk is not None:
+                        yield error_chunk
+                        return
+                    raise
                 yield original_item
             else:
                 yield item
@@ -390,9 +472,18 @@ class UnifiedLLMGuardrails(CustomLogger):
                 CallTypes(call_type)
             ]()
 
-            await endpoint_translation.process_output_streaming_response(
-                responses_so_far=responses_so_far,
-                guardrail_to_apply=guardrail_to_apply,
-                litellm_logging_obj=request_data.get("litellm_logging_obj"),
-                user_api_key_dict=user_api_key_dict,
-            )
+            try:
+                await endpoint_translation.process_output_streaming_response(
+                    responses_so_far=responses_so_far,
+                    guardrail_to_apply=guardrail_to_apply,
+                    litellm_logging_obj=request_data.get("litellm_logging_obj"),
+                    user_api_key_dict=user_api_key_dict,
+                )
+            except HTTPException as e:
+                error_chunk = _format_a2a_guardrail_error_chunk(
+                    e, call_type, responses_so_far, request_data
+                )
+                if error_chunk is not None:
+                    yield error_chunk
+                else:
+                    raise

litellm/proxy/utils.py

@@ -945,6 +945,7 @@ class ProxyLogging:
         data: dict,
         user_api_key_dict: Optional[UserAPIKeyAuth],
         call_type: CallTypesLiteral,
+        event_type: GuardrailEventHooks,
     ) -> Optional[dict]:
         """
         Process a guardrail callback during pre-call hook.
@@ -964,8 +965,10 @@ class ProxyLogging:
         from litellm.types.guardrails import GuardrailEventHooks
 
         # Determine the event type based on call type
-        event_type = GuardrailEventHooks.pre_call
-        if call_type == CallTypes.call_mcp_tool.value:
+        if (
+            event_type is GuardrailEventHooks.pre_call
+            and call_type == CallTypes.call_mcp_tool.value
+        ):
             event_type = GuardrailEventHooks.pre_mcp_call
 
         # Check if the guardrail should run for this request
@@ -1222,6 +1225,7 @@ class ProxyLogging:
                         data=data,  # type: ignore
                         user_api_key_dict=user_api_key_dict,
                         call_type=call_type,
+                        event_type=GuardrailEventHooks.pre_call,
                     )
                     if result is None:
                         continue
@@ -1375,11 +1379,11 @@ class ProxyLogging:
         # Note: user_info is a CallInfo that can represent user/team/org level info. For team budgets,
         # alert_emails is populated from team_object.metadata.soft_budget_alerting_emails (see auth_checks.py)
         is_soft_budget_with_alert_emails = (
-            type == "soft_budget" 
-            and user_info.alert_emails is not None 
+            type == "soft_budget"
+            and user_info.alert_emails is not None
             and len(user_info.alert_emails) > 0
         )
-        
+
         if self.alerting is None and not is_soft_budget_with_alert_emails:
             # do nothing if alerting is not switched on (unless it's a soft_budget alert with team-specific emails)
             return
@@ -1395,10 +1399,9 @@ class ProxyLogging:
         # 1. "email" is in alerting config, OR
         # 2. It's a soft_budget alert with team-specific alert_emails (bypasses global alerting config)
         should_send_email = (
-            (self.alerting is not None and "email" in self.alerting) 
-            or is_soft_budget_with_alert_emails
-        )
-        
+            self.alerting is not None and "email" in self.alerting
+        ) or is_soft_budget_with_alert_emails
+
         if should_send_email and self.email_logging_instance is not None:
             await self.email_logging_instance.budget_alerts(
                 type=type,
@@ -1762,6 +1765,7 @@ class ProxyLogging:
 
         from litellm.types.guardrails import GuardrailEventHooks
 
+
         guardrail_callbacks: List[CustomGuardrail] = []
         other_callbacks: List[CustomLogger] = []
         try:
@@ -1867,6 +1871,10 @@ class ProxyLogging:
             )
         return merged_headers
 
+    def is_a2a_streaming_response(self, response: dict) -> bool:
+        expected_keys = ["jsonrpc", "id", "result"]
+        return all(key in response for key in expected_keys)
+
     async def async_post_call_streaming_hook(
         self,
         data: dict,
@@ -1887,6 +1895,10 @@ class ProxyLogging:
         response_str: Optional[str] = None
         if isinstance(response, (ModelResponse, ModelResponseStream)):
             response_str = litellm.get_response_string(response_obj=response)
+        elif isinstance(response, dict) and self.is_a2a_streaming_response(response):
+            from litellm.llms.a2a.common_utils import extract_text_from_a2a_response
+
+            response_str = extract_text_from_a2a_response(response)
         if response_str is not None:
             for callback in litellm.callbacks:
                 try:

litellm/types/utils.py

@@ -33,6 +33,7 @@ from litellm.types.llms.base import (
 from litellm.types.mcp import MCPServerCostInfo
 
 from ..litellm_core_utils.core_helpers import map_finish_reason
+from .agents import LiteLLMSendMessageResponse
 from .guardrails import GuardrailEventHooks
 from .llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse
 from .llms.base import HiddenParams
@@ -777,6 +778,7 @@ API_ROUTE_TO_CALL_TYPES = {
     "/mcp/call_tool": [CallTypes.call_mcp_tool],
     # A2A (Agent-to-Agent)
     "/a2a/{agent_id}": [CallTypes.asend_message, CallTypes.send_message],
+    "/a2a/{agent_id}/message/send": [CallTypes.asend_message, CallTypes.send_message],
     # Passthrough endpoints
     "/llm_passthrough": [
         CallTypes.llm_passthrough_route,
@@ -2139,7 +2141,14 @@ class ImageObject(OpenAIImage):
     revised_prompt: Optional[str] = None
     provider_specific_fields: Optional[Dict[str, Any]] = None
 
-    def __init__(self, b64_json=None, url=None, revised_prompt=None, provider_specific_fields=None, **kwargs):
+    def __init__(
+        self,
+        b64_json=None,
+        url=None,
+        revised_prompt=None,
+        provider_specific_fields=None,
+        **kwargs,
+    ):
         super().__init__(b64_json=b64_json, url=url, revised_prompt=revised_prompt)  # type: ignore
         if provider_specific_fields:
             self.provider_specific_fields = provider_specific_fields
@@ -2641,7 +2650,9 @@ class CostBreakdown(TypedDict, total=False):
     )
     total_cost: float  # Total cost (input + output + tool usage)
     tool_usage_cost: float  # Cost of usage of built-in tools
-    additional_costs: Dict[str, float]  # Free-form additional costs (e.g., {"azure_model_router_flat_cost": 0.00014})
+    additional_costs: Dict[
+        str, float
+    ]  # Free-form additional costs (e.g., {"azure_model_router_flat_cost": 0.00014})
     original_cost: float  # Cost before discount (optional)
     discount_percent: float  # Discount percentage applied (e.g., 0.05 = 5%) (optional)
     discount_amount: float  # Discount amount in USD (optional)
@@ -3355,6 +3366,7 @@ LLMResponseTypes = Union[
     LiteLLMFineTuningJob,
     AnthropicMessagesResponse,
     ResponsesAPIResponse,
+    LiteLLMSendMessageResponse,
 ]
 
 

tests/test_litellm/llms/a2a/chat/guardrail_translation/__init__.py

@@ -0,0 +1 @@
+# Tests for A2A guardrail translation

tests/test_litellm/llms/a2a/chat/guardrail_translation/test_a2a_guardrail_handler.py

@@ -0,0 +1,314 @@
+"""
+Test A2A Guardrail Translation Handler
+
+Unit tests for the A2A protocol guardrail handler, covering:
+- Text extraction from A2A message parts (input and output formats)
+- In-place modification logic for streaming responses
+- Defensive handling of malformed or empty inputs
+"""
+
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from litellm.llms.a2a.chat.guardrail_translation.handler import A2AGuardrailHandler
+from litellm.types.utils import CallTypes
+
+
+@pytest.fixture
+def mock_guardrail():
+    """Guardrail mock that echoes input texts."""
+    guardrail = MagicMock()
+    guardrail.apply_guardrail = AsyncMock(
+        side_effect=lambda inputs, **kwargs: {"texts": inputs.get("texts", [])}
+    )
+    return guardrail
+
+
+class TestA2AGuardrailHandlerProcessInputMessages:
+    """Tests for process_input_messages (pre-call hook)."""
+
+    @pytest.mark.asyncio
+    async def test_extracts_and_applies_guardrail_to_text_parts(self, mock_guardrail):
+        """Should extract text from kind=text parts and apply guardrail."""
+        handler = A2AGuardrailHandler()
+        mock_guardrail.apply_guardrail = AsyncMock(
+            return_value={"texts": ["guardrailed hello", "guardrailed world"]}
+        )
+
+        data = {
+            "params": {
+                "message": {
+                    "parts": [
+                        {"kind": "text", "text": "hello"},
+                        {"kind": "text", "text": "world"},
+                    ]
+                }
+            }
+        }
+
+        result = await handler.process_input_messages(
+            data=data,
+            guardrail_to_apply=mock_guardrail,
+        )
+
+        mock_guardrail.apply_guardrail.assert_called_once()
+        call_inputs = mock_guardrail.apply_guardrail.call_args.kwargs["inputs"]
+        assert call_inputs["texts"] == ["hello", "world"]
+
+        assert result["params"]["message"]["parts"][0]["text"] == "guardrailed hello"
+        assert result["params"]["message"]["parts"][1]["text"] == "guardrailed world"
+
+    @pytest.mark.asyncio
+    async def test_skips_empty_parts(self, mock_guardrail):
+        """Should skip parts with no text content."""
+        handler = A2AGuardrailHandler()
+
+        data = {
+            "params": {
+                "message": {
+                    "parts": [
+                        {"kind": "text", "text": ""},
+                        {"kind": "model", "model": "gpt-4"},
+                    ]
+                }
+            }
+        }
+
+        result = await handler.process_input_messages(
+            data=data,
+            guardrail_to_apply=mock_guardrail,
+        )
+
+        mock_guardrail.apply_guardrail.assert_not_called()
+        assert result == data
+
+    @pytest.mark.asyncio
+    async def test_returns_unchanged_when_no_parts(self, mock_guardrail):
+        """Should return data unchanged when message has no parts."""
+        handler = A2AGuardrailHandler()
+        data = {"params": {"message": {}}}
+
+        result = await handler.process_input_messages(
+            data=data,
+            guardrail_to_apply=mock_guardrail,
+        )
+
+        mock_guardrail.apply_guardrail.assert_not_called()
+        assert result == data
+
+
+class TestA2AGuardrailHandlerProcessOutputResponse:
+    """Tests for process_output_response (post-call, non-streaming)."""
+
+    @pytest.mark.asyncio
+    async def test_applies_guardrail_to_direct_message_parts(self, mock_guardrail):
+        """Should process result.parts format."""
+        mock_guardrail.apply_guardrail = AsyncMock(
+            return_value={"texts": ["guardrailed output"]}
+        )
+        handler = A2AGuardrailHandler()
+
+        response = {
+            "result": {
+                "kind": "message",
+                "parts": [{"kind": "text", "text": "original output"}],
+            }
+        }
+
+        result = await handler.process_output_response(
+            response=response,
+            guardrail_to_apply=mock_guardrail,
+        )
+
+        assert result["result"]["parts"][0]["text"] == "guardrailed output"
+
+    @pytest.mark.asyncio
+    async def test_applies_guardrail_to_nested_message_parts(self, mock_guardrail):
+        """Should process result.message.parts format."""
+        mock_guardrail.apply_guardrail = AsyncMock(
+            return_value={"texts": ["guardrailed nested"]}
+        )
+        handler = A2AGuardrailHandler()
+
+        response = {
+            "result": {
+                "message": {
+                    "parts": [{"kind": "text", "text": "nested text"}],
+                }
+            }
+        }
+
+        result = await handler.process_output_response(
+            response=response,
+            guardrail_to_apply=mock_guardrail,
+        )
+
+        assert result["result"]["message"]["parts"][0]["text"] == "guardrailed nested"
+
+    @pytest.mark.asyncio
+    async def test_applies_guardrail_to_artifact_parts(self, mock_guardrail):
+        """Should process result.artifact.parts (streaming artifact-update format)."""
+        mock_guardrail.apply_guardrail = AsyncMock(
+            return_value={"texts": ["guardrailed artifact"]}
+        )
+        handler = A2AGuardrailHandler()
+
+        response = {
+            "result": {
+                "kind": "artifact-update",
+                "artifact": {
+                    "parts": [{"kind": "text", "text": "artifact text"}],
+                },
+            }
+        }
+
+        result = await handler.process_output_response(
+            response=response,
+            guardrail_to_apply=mock_guardrail,
+        )
+
+        assert result["result"]["artifact"]["parts"][0]["text"] == "guardrailed artifact"
+
+
+class TestA2AGuardrailHandlerProcessOutputStreamingResponse:
+    """
+    Tests for process_output_streaming_response.
+
+    IMPORTANT: This method modifies responses_so_far IN-PLACE. It:
+    1. Concatenates all text from chunks in order
+    2. Applies guardrail once to the combined text
+    3. Writes the full guardrailed text into the FIRST chunk that had text
+    4. CLEARS all other text parts in subsequent chunks to "" (in-place)
+    """
+
+    @pytest.mark.asyncio
+    async def test_streaming_combines_text_and_puts_in_first_chunk(self, mock_guardrail):
+        """Combined guardrailed text should be placed in first chunk; others cleared."""
+        mock_guardrail.apply_guardrail = AsyncMock(
+            return_value={"texts": ["COMBINED_GUARDRAILED"]}
+        )
+        handler = A2AGuardrailHandler()
+
+        chunk1 = {
+            "result": {
+                "artifact": {"parts": [{"kind": "text", "text": "chunk1 "}]},
+            }
+        }
+        chunk2 = {
+            "result": {
+                "artifact": {"parts": [{"kind": "text", "text": "chunk2"}]},
+            }
+        }
+        responses_so_far = [chunk1, chunk2]
+
+        result = await handler.process_output_streaming_response(
+            responses_so_far=responses_so_far,
+            guardrail_to_apply=mock_guardrail,
+        )
+
+        # In-place: first chunk gets full guardrailed text
+        assert chunk1["result"]["artifact"]["parts"][0]["text"] == "COMBINED_GUARDRAILED"
+        # Second chunk's text is cleared
+        assert chunk2["result"]["artifact"]["parts"][0]["text"] == ""
+
+        assert result is responses_so_far  # Same list, modified in place
+
+    @pytest.mark.asyncio
+    async def test_streaming_handles_ndjson_strings(self, mock_guardrail):
+        """Should parse NDJSON strings and write back as NDJSON."""
+        mock_guardrail.apply_guardrail = AsyncMock(
+            return_value={"texts": ["GUARDRAILED"]}
+        )
+        handler = A2AGuardrailHandler()
+
+        responses_so_far = [
+            '{"result":{"artifact":{"parts":[{"kind":"text","text":"hello"}]}}}\n',
+        ]
+
+        result = await handler.process_output_streaming_response(
+            responses_so_far=responses_so_far,
+            guardrail_to_apply=mock_guardrail,
+        )
+
+        # responses_so_far is modified in place; NDJSON string is updated
+        parsed = __import__("json").loads(responses_so_far[0].strip())
+        assert parsed["result"]["artifact"]["parts"][0]["text"] == "GUARDRAILED"
+
+    @pytest.mark.asyncio
+    async def test_streaming_returns_early_when_no_text(self, mock_guardrail):
+        """Should return responses_so_far unchanged when no text content."""
+        handler = A2AGuardrailHandler()
+
+        responses_so_far = [
+            {"result": {"artifact": {"parts": [{"kind": "model", "model": "gpt-4"}]}}},
+        ]
+
+        result = await handler.process_output_streaming_response(
+            responses_so_far=responses_so_far,
+            guardrail_to_apply=mock_guardrail,
+        )
+
+        mock_guardrail.apply_guardrail.assert_not_called()
+        assert result == responses_so_far
+
+
+class TestA2AGuardrailHandlerExtractTextsFromResult:
+    """Tests for _extract_texts_from_result helper."""
+
+    def test_extracts_from_multiple_formats(self):
+        """Should extract text from parts, message.parts, artifact.parts, etc."""
+        handler = A2AGuardrailHandler()
+        texts: list = []
+        mappings: list = []
+
+        result = {
+            "parts": [{"kind": "text", "text": "direct"}],
+            "message": {"parts": [{"kind": "text", "text": "nested"}]},
+            "artifact": {"parts": [{"kind": "text", "text": "artifact"}]},
+        }
+
+        handler._extract_texts_from_result(result, texts, mappings)
+
+        assert texts == ["direct", "nested", "artifact"]
+        assert len(mappings) == 3
+
+
+class TestA2AGuardrailHandlerApplyTextToPath:
+    """Tests for _apply_text_to_path helper."""
+
+    def test_applies_text_to_nested_path(self):
+        """Should navigate path and update part text."""
+        handler = A2AGuardrailHandler()
+        result = {
+            "message": {
+                "parts": [
+                    {"kind": "text", "text": "old"},
+                ]
+            }
+        }
+
+        handler._apply_text_to_path(
+            result=result,
+            path=("message", "parts"),
+            part_idx=0,
+            text="new",
+        )
+
+        assert result["message"]["parts"][0]["text"] == "new"
+
+
+def test_a2a_guardrail_translation_mappings():
+    """A2A handler should be registered for send_message and asend_message."""
+    from litellm.llms.a2a.chat.guardrail_translation import (
+        guardrail_translation_mappings,
+    )
+
+    assert CallTypes.send_message in guardrail_translation_mappings
+    assert CallTypes.asend_message in guardrail_translation_mappings
+    assert (
+        guardrail_translation_mappings[CallTypes.send_message] == A2AGuardrailHandler
+    )
+    assert (
+        guardrail_translation_mappings[CallTypes.asend_message] == A2AGuardrailHandler
+    )

tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py

@@ -1,6 +1,9 @@
 """Tests for unified guardrail."""
 
+import json
+
 import pytest
+from fastapi import HTTPException
 
 from litellm.caching import DualCache
 from litellm.integrations.custom_guardrail import CustomGuardrail
@@ -9,7 +12,9 @@ from litellm.proxy._experimental.mcp_server.guardrail_translation.handler import
     MCPGuardrailTranslationHandler,
 )
 from litellm.proxy._types import UserAPIKeyAuth
-from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail import unified_guardrail as unified_module
+from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail import (
+    unified_guardrail as unified_module,
+)
 from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
     UnifiedLLMGuardrails,
 )
@@ -229,3 +234,129 @@ class TestUnifiedLLMGuardrails:
                     f"Chunk {i} lost its content (got {content!r}). "
                     f"Expected non-empty content for every streamed chunk."
                 )
+
+        @pytest.mark.asyncio
+        async def test_a2a_streaming_httpexception_yields_jsonrpc_error_chunk(self):
+            """
+            When A2A streaming guardrail raises HTTPException, the hook should yield
+            a JSON-RPC 2.0 error chunk so the client sees the rejection in-stream,
+            since the HTTP response has already started.
+            """
+
+            class _HTTPExceptionRaisingTranslation(BaseTranslation):
+                """Raises HTTPException to simulate guardrail rejection."""
+
+                async def process_input_messages(self, data, guardrail_to_apply, litellm_logging_obj=None):  # type: ignore[override]
+                    return data
+
+                async def process_output_response(self, response, guardrail_to_apply, litellm_logging_obj=None, user_api_key_dict=None):  # type: ignore[override]
+                    return response
+
+                async def process_output_streaming_response(
+                    self,
+                    responses_so_far,
+                    guardrail_to_apply,
+                    litellm_logging_obj=None,
+                    user_api_key_dict=None,
+                ):
+                    raise HTTPException(status_code=400, detail={"error": "Content blocked"})
+
+            unified_module.endpoint_guardrail_translation_mappings = {
+                CallTypes.asend_message: _HTTPExceptionRaisingTranslation,
+            }
+
+            handler = UnifiedLLMGuardrails()
+            guardrail = RecordingGuardrail()
+
+            # A2A NDJSON chunk format
+            chunks = [
+                '{"jsonrpc":"2.0","id":"req-1","result":{"artifact":{"parts":[{"kind":"text","text":"blocked"}]}}}\n',
+            ]
+
+            async def mock_stream():
+                for chunk in chunks:
+                    yield chunk
+
+            user_api_key_dict = UserAPIKeyAuth(
+                api_key="test-key",
+                request_route="/a2a/test-agent/message/send",
+            )
+
+            request_data = {
+                "guardrail_to_apply": guardrail,
+                "body": {"id": "req-1"},
+            }
+
+            yielded = []
+            async for item in handler.async_post_call_streaming_iterator_hook(
+                user_api_key_dict=user_api_key_dict,
+                response=mock_stream(),
+                request_data=request_data,
+            ):
+                yielded.append(item)
+
+            # Should have yielded the first chunk(s) plus a JSON-RPC error
+            assert len(yielded) >= 1
+            last = yielded[-1]
+            if isinstance(last, str):
+                parsed = json.loads(last.strip())
+                assert "error" in parsed
+                assert parsed["error"]["code"] == -32603
+                assert "Content blocked" in str(parsed["error"]["message"])
+
+        @pytest.mark.asyncio
+        async def test_non_a2a_streaming_httpexception_re_raises(self):
+            """When non-A2A streaming guardrail raises HTTPException, it should re-raise."""
+
+            class _HTTPExceptionRaisingTranslation(BaseTranslation):
+                async def process_input_messages(self, data, guardrail_to_apply, litellm_logging_obj=None):  # type: ignore[override]
+                    return data
+
+                async def process_output_response(self, response, guardrail_to_apply, litellm_logging_obj=None, user_api_key_dict=None):  # type: ignore[override]
+                    return response
+
+                async def process_output_streaming_response(
+                    self,
+                    responses_so_far,
+                    guardrail_to_apply,
+                    litellm_logging_obj=None,
+                    user_api_key_dict=None,
+                ):
+                    raise HTTPException(status_code=400, detail="Blocked")
+
+            unified_module.endpoint_guardrail_translation_mappings = {
+                CallTypes.acompletion: _HTTPExceptionRaisingTranslation,
+            }
+
+            handler = UnifiedLLMGuardrails()
+            guardrail = RecordingGuardrail()
+
+            chunks = [
+                ModelResponseStream(
+                    choices=[StreamingChoices(delta=Delta(content="hi", role="assistant"), finish_reason=None)],
+                ),
+            ]
+
+            async def mock_stream():
+                for chunk in chunks:
+                    yield chunk
+
+            user_api_key_dict = UserAPIKeyAuth(
+                api_key="test-key",
+                request_route="/v1/chat/completions",
+            )
+
+            request_data = {
+                "guardrail_to_apply": guardrail,
+                "model": "gpt-4",
+            }
+
+            with pytest.raises(HTTPException) as exc_info:
+                async for _ in handler.async_post_call_streaming_iterator_hook(
+                    user_api_key_dict=user_api_key_dict,
+                    response=mock_stream(),
+                    request_data=request_data,
+                ):
+                    pass
+
+            assert exc_info.value.status_code == 400

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

@@ -230,6 +230,45 @@ const CustomCodeModal: React.FC<CustomCodeModalProps> = ({
         structured_messages: [],
         model: "gpt-4"
       }
+    },
+    pre_mcp_call: {
+      name: "Pre MCP (MCP tool as OpenAI tool)",
+      data: {
+        texts: [
+          "Tool: read_wiki_structure\nArguments: {\"repoName\": \"BerriAI/litellm\"}"
+        ],
+        images: [],
+        tools: [
+          {
+            type: "function",
+            function: {
+              name: "read_wiki_structure",
+              description: "Read the structure of a GitHub repository (MCP tool passed as OpenAI tool)",
+              parameters: {
+                type: "object",
+                properties: {
+                  repoName: { type: "string", description: "Repository name, e.g. BerriAI/litellm" }
+                },
+                required: ["repoName"]
+              }
+            }
+          }
+        ],
+        tool_calls: [
+          {
+            id: "call_mcp_001",
+            type: "function",
+            function: {
+              name: "read_wiki_structure",
+              arguments: "{\"repoName\": \"BerriAI/litellm\"}"
+            }
+          }
+        ],
+        structured_messages: [
+          { role: "user", content: "Tool: read_wiki_structure\nArguments: {\"repoName\": \"BerriAI/litellm\"}" }
+        ],
+        model: "mcp-tool-call"
+      }
     }
   };
   
@@ -559,6 +598,13 @@ const CustomCodeModal: React.FC<CustomCodeModalProps> = ({
                         >
                           Pre-call
                         </button>
+                        <button
+                          type="button"
+                          onClick={() => setTestInput(JSON.stringify(TEST_INPUT_EXAMPLES.pre_mcp_call.data, null, 2))}
+                          className="px-2 py-1 text-xs rounded border border-purple-200 bg-purple-50 text-purple-700 hover:bg-purple-100 transition-colors"
+                        >
+                          Pre MCP
+                        </button>
                         <button
                           type="button"
                           onClick={() => setTestInput(JSON.stringify(TEST_INPUT_EXAMPLES.post_call.data, null, 2))}
@@ -572,7 +618,7 @@ const CustomCodeModal: React.FC<CustomCodeModalProps> = ({
                       <div className="grid grid-cols-2 gap-x-4 gap-y-1">
                         <div><strong>texts</strong>: Message content (always)</div>
                         <div><strong>images</strong>: Base64 images (vision)</div>
-                        <div><strong>tools</strong>: Tool definitions <span className="text-orange-600">(pre_call)</span></div>
+                        <div><strong>tools</strong>: Tool definitions <span className="text-orange-600">(pre_call)</span>, MCP as OpenAI tool <span className="text-purple-600">(pre_mcp_call)</span></div>
                         <div><strong>tool_calls</strong>: LLM tool calls <span className="text-green-600">(post_call)</span></div>
                         <div><strong>structured_messages</strong>: Full messages <span className="text-orange-600">(pre_call)</span></div>
                         <div><strong>model</strong>: Model name (always)</div>