Merge pull request #40461 from BerriAI/litellm_lit_7373_responses_custom_tool_call_guardrail

fix(guardrails): scan and rewrite Responses custom_tool_call output items

litellm/completion_extras/litellm_responses_transformation/transformation.py

@@ -227,7 +227,7 @@ class _ChatToolCallDict(ChatCompletionToolCallChunk, total=False):
     provider_specific_fields: Mapping[str, object]
 
 
-def _tool_call_dict_from_output_item(item: Mapping[str, Any], index: int) -> _ChatToolCallDict:
+def tool_call_dict_from_output_item(item: Mapping[str, Any], index: int) -> _ChatToolCallDict:
     """Convert a ``function_call`` or ``custom_tool_call`` output item dict to a chat
     completions tool_call dict. Custom (grammar/freeform) tool calls carry their raw
     string payload in ``input`` rather than ``arguments``; both map to
@@ -755,7 +755,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
                     # Tool calls accumulate into the single trailing tool_calls choice
                     # like the typed branches above; a choice per call would hide every
                     # call after choices[0] from chat clients
-                    accumulated_tool_calls.append(_tool_call_dict_from_output_item(raw_item, tool_call_index))
+                    accumulated_tool_calls.append(tool_call_dict_from_output_item(raw_item, tool_call_index))
                     tool_call_index += 1
                 elif handle_raw_dict_callback is not None:
                     choice, index = handle_raw_dict_callback(item=raw_item, index=index)
@@ -1409,7 +1409,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
             # New output item added
             output_item = parsed_chunk.get("item", {})
             if output_item.get("type") in ("function_call", "custom_tool_call"):
-                converted: Final = _tool_call_dict_from_output_item(output_item, parsed_chunk.get("output_index", 0))
+                converted: Final = tool_call_dict_from_output_item(output_item, parsed_chunk.get("output_index", 0))
                 provider_specific_fields: Final = converted.get("provider_specific_fields")
 
                 function_chunk: Final = ChatCompletionToolCallFunctionChunk(
@@ -1484,7 +1484,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
                                 index=0,
                                 delta=Delta(
                                     tool_calls=(
-                                        _tool_call_dict_from_output_item(
+                                        tool_call_dict_from_output_item(
                                             output_item, parsed_chunk.get("output_index", 0)
                                         ),
                                     )

litellm/llms/openai/responses/guardrail_translation/handler.py

@@ -37,14 +37,14 @@ from itertools import accumulate, chain, repeat
 from types import MappingProxyType
 from typing import TYPE_CHECKING, Any, Final, NamedTuple, Union, cast
 
-from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall
-from pydantic import BaseModel, TypeAdapter
+from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
 from typing_extensions import ReadOnly, TypedDict
 
 from litellm._logging import verbose_proxy_logger
 from litellm.completion_extras.litellm_responses_transformation.transformation import (
     LiteLLMResponsesTransformationHandler,
     OpenAiResponsesToChatCompletionStreamIterator,
+    tool_call_dict_from_output_item,
 )
 from litellm.llms.base_llm.guardrail_translation.base_translation import (
     BaseTranslation,
@@ -84,7 +84,6 @@ from litellm.types.llms.openai import (
 )
 from litellm.types.responses.main import (
     GenericResponseOutputItem,
-    OutputFunctionToolCall,
     OutputText,
 )
 from litellm.types.utils import GenericGuardrailAPIInputs
@@ -106,6 +105,19 @@ class _ToolCallShape(NamedTuple):
     arguments: str
 
 
+class _ToolCallFunctionFields(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    name: str | None = None
+    arguments: str = ""
+
+
+class _ToolCallFields(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    function: _ToolCallFunctionFields
+
+
 def _tool_call_shapes(tool_calls: Sequence[ChatCompletionToolCallChunk]) -> tuple[_ToolCallShape, ...]:
     return tuple(
         _ToolCallShape(name=tool_call["function"].get("name"), arguments=tool_call["function"].get("arguments", ""))
@@ -113,6 +125,47 @@ def _tool_call_shapes(tool_calls: Sequence[ChatCompletionToolCallChunk]) -> tupl
     )
 
 
+def _returned_tool_call_shape(tool_call: object) -> _ToolCallShape | None:
+    payload: Final = tool_call.model_dump() if isinstance(tool_call, BaseModel) else tool_call
+    try:
+        fields: Final = _ToolCallFields.model_validate(payload)
+    except ValidationError:
+        return None
+    return _ToolCallShape(name=fields.function.name, arguments=fields.function.arguments)
+
+
+def _post_guardrail_tool_call_shapes(
+    returned_tool_calls: Sequence[object] | None,
+    pre_guardrail_tool_calls: tuple[_ToolCallShape, ...],
+    guardrail_name: str | None,
+) -> tuple[_ToolCallShape, ...]:
+    if not pre_guardrail_tool_calls:
+        return pre_guardrail_tool_calls
+    if returned_tool_calls is None or len(returned_tool_calls) != len(pre_guardrail_tool_calls):
+        verbose_proxy_logger.warning(
+            "OpenAI Responses API: guardrail %s returned %s tool calls for the %d scanned, "
+            "leaving the tool call output items unchanged",
+            guardrail_name,
+            "no" if returned_tool_calls is None else len(returned_tool_calls),
+            len(pre_guardrail_tool_calls),
+        )
+        return pre_guardrail_tool_calls
+    returned_shapes: Final = tuple(_returned_tool_call_shape(tool_call) for tool_call in returned_tool_calls)
+    validated_shapes: Final = tuple(shape for shape in returned_shapes if shape is not None)
+    if len(validated_shapes) != len(returned_shapes):
+        verbose_proxy_logger.warning(
+            "OpenAI Responses API: guardrail %s returned tool calls without a function name and arguments, "
+            "leaving the tool call output items unchanged",
+            guardrail_name,
+        )
+        return pre_guardrail_tool_calls
+    return validated_shapes
+
+
+def _tool_call_rewrite(before: _ToolCallShape, after: _ToolCallShape) -> _ToolCallShape:
+    return _ToolCallShape(name=after.name if after.name != before.name else None, arguments=after.arguments)
+
+
 class ResponseOutputEnvelope(TypedDict, total=False):
     """Dict form of a Responses API response, as far as guardrail write-back reads it."""
 
@@ -140,8 +193,18 @@ _TERMINAL_ENVELOPE_EVENT_TYPES: Final = frozenset(
 )
 
 
-_FUNCTION_CALL_ARGUMENT_EVENT_TYPES: Final = frozenset(
-    {"response.function_call_arguments.delta", "response.function_call_arguments.done"}
+_TOOL_CALL_ITEM_TYPES: Final = frozenset({"function_call", "custom_tool_call"})
+_TOOL_CALL_PAYLOAD_FIELDS: Final[Mapping[str, str]] = MappingProxyType(
+    {"function_call": "arguments", "custom_tool_call": "input"}
+)
+_TOOL_CALL_PAYLOAD_DELTA_EVENT_TYPES: Final = frozenset(
+    {"response.function_call_arguments.delta", "response.custom_tool_call_input.delta"}
+)
+_TOOL_CALL_PAYLOAD_DONE_EVENT_FIELDS: Final[Mapping[str, str]] = MappingProxyType(
+    {"response.function_call_arguments.done": "arguments", "response.custom_tool_call_input.done": "input"}
+)
+_TOOL_CALL_PAYLOAD_EVENT_TYPES: Final = _TOOL_CALL_PAYLOAD_DELTA_EVENT_TYPES | frozenset(
+    _TOOL_CALL_PAYLOAD_DONE_EVENT_FIELDS
 )
 _OUTPUT_ITEM_EVENT_TYPES: Final = frozenset({"response.output_item.added", "response.output_item.done"})
 _PATCHABLE_ITEM_FIELDS: Final[Mapping[str, str]] = MappingProxyType(
@@ -180,8 +243,20 @@ def _rewritten_input_item(item: Mapping[str, object], rewritten: object) -> Mapp
     return {**item, field: converted_value}  # mutable-ok: request input items must stay JSON-plain dicts
 
 
-def _is_function_call_item(item: object) -> bool:
-    return isinstance(item, Mapping) and item.get("type") in ("function_call", "custom_tool_call")
+def _is_tool_call_item(item: object) -> bool:
+    return isinstance(item, Mapping) and item.get("type") in _TOOL_CALL_ITEM_TYPES
+
+
+def _tool_call_output_item_mapping(item: object) -> Mapping[str, object] | None:
+    if stream_item_field(item, "type") not in _TOOL_CALL_ITEM_TYPES:
+        return None
+    if isinstance(item, Mapping):
+        return cast("Mapping[str, object]", item)  # cast-ok: output items are str-keyed JSON objects
+    return item.model_dump() if isinstance(item, BaseModel) else None
+
+
+def _is_tool_call_output_item(item: object) -> bool:
+    return _tool_call_output_item_mapping(item) is not None
 
 
 def _last_message_role(messages: Sequence[object]) -> str | None:
@@ -205,7 +280,7 @@ def _provenance_unit_bounds(
     start_indexes: Final = tuple(
         index
         for index in range(len(raw_input))
-        if index == 0 or not (_is_function_call_item(raw_input[index]) and trailing_roles[index - 1] == "assistant")
+        if index == 0 or not (_is_tool_call_item(raw_input[index]) and trailing_roles[index - 1] == "assistant")
     )
     return tuple(zip(start_indexes, (*start_indexes[1:], len(raw_input))))
 
@@ -603,7 +678,7 @@ class OpenAIResponsesHandler(BaseTranslation):
             - response.output is a list of output items
             - Each output item can be:
               * GenericResponseOutputItem with a content list of OutputText objects
-              * ResponseFunctionToolCall with tool call data
+              * ResponseFunctionToolCall or CustomToolCallOutputItem with tool call data
             - Each OutputText object has a text field
         """
 
@@ -668,6 +743,7 @@ class OpenAIResponsesHandler(BaseTranslation):
             if response_model:
                 inputs["model"] = response_model
 
+            pre_guardrail_tool_calls: Final = _tool_call_shapes(tool_calls_to_check)
             guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
                 inputs=inputs,
                 request_data=request_data,
@@ -676,6 +752,11 @@ class OpenAIResponsesHandler(BaseTranslation):
             )
 
             guardrailed_texts: Final = guardrailed_inputs.get("texts", [])
+            post_guardrail_tool_calls: Final = _post_guardrail_tool_call_shapes(
+                returned_tool_calls=guardrailed_inputs.get("tool_calls"),
+                pre_guardrail_tool_calls=pre_guardrail_tool_calls,
+                guardrail_name=guardrail_to_apply.guardrail_name,
+            )
 
             # Step 3: Map guardrail responses back to original response structure
             await self._apply_guardrail_responses_to_output(
@@ -683,6 +764,11 @@ class OpenAIResponsesHandler(BaseTranslation):
                 responses=guardrailed_texts,
                 task_mappings=task_mappings,
             )
+            self._write_tool_call_rewrites_to_output(
+                tool_call_items=tuple(item for item in response_output if _is_tool_call_output_item(item)),
+                pre_guardrail_tool_calls=pre_guardrail_tool_calls,
+                post_guardrail_tool_calls=post_guardrail_tool_calls,
+            )
 
         verbose_proxy_logger.debug("OpenAI Responses API: Processed output response: %s", response)
 
@@ -779,11 +865,10 @@ class OpenAIResponsesHandler(BaseTranslation):
                 )
 
                 guardrailed_texts: Final = guardrailed_inputs.get("texts", [])
-                returned_tool_calls: Final = guardrailed_inputs.get("tool_calls")
-                post_guardrail_tool_calls: Final = _tool_call_shapes(
-                    returned_tool_calls
-                    if isinstance(returned_tool_calls, list) and len(returned_tool_calls) == len(tool_calls_to_check)
-                    else tool_calls_to_check
+                post_guardrail_tool_calls: Final = _post_guardrail_tool_call_shapes(
+                    returned_tool_calls=guardrailed_inputs.get("tool_calls"),
+                    pre_guardrail_tool_calls=pre_guardrail_tool_calls,
+                    guardrail_name=guardrail_to_apply.guardrail_name,
                 )
 
                 # Write guardrailed texts back into the output items in-place.
@@ -933,11 +1018,12 @@ class OpenAIResponsesHandler(BaseTranslation):
         guardrail_name: str,
     ) -> None:
         """Write ended-stream guardrail tool-call rewrites into the completed
-        envelope's ``function_call`` items and sync the earlier stream events,
-        keyed by ``call_id``. The guardrail sees the envelope's function calls
-        in output order, which is how a rewritten call finds its ``call_id``;
-        the stream events find their call through the ``call_id`` on
-        ``output_item`` events and the ``item_id`` on argument events, since an
+        envelope's ``function_call`` and ``custom_tool_call`` items and sync the
+        earlier stream events, keyed by ``call_id``. The guardrail sees the
+        envelope's tool calls in output order, which is how a rewritten call
+        finds its ``call_id``; the stream events find their call through the
+        ``call_id`` on ``output_item`` events and the ``item_id`` on argument
+        and custom-input events, since an
         event's ``output_index`` need not match the envelope's (the chat bridge
         numbers tool calls from 1 while the envelope lists them after the
         message). A rewrite whose calls do not line up with the envelope, or
@@ -945,32 +1031,30 @@ class OpenAIResponsesHandler(BaseTranslation):
         pipeline executor discards it and releases the original events."""
         if post_guardrail_tool_calls == pre_guardrail_tool_calls:
             return
-        function_call_items: Final = tuple(
-            output_item for output_item in outputs if stream_item_field(output_item, "type") == "function_call"
-        )
+        tool_call_items: Final = tuple(output_item for output_item in outputs if _is_tool_call_output_item(output_item))
         call_ids: Final = tuple(
             call_id
-            for output_item in function_call_items
+            for output_item in tool_call_items
             if isinstance(call_id := stream_item_field(output_item, "call_id"), str) and call_id
         )
         stream_events: Final = responses_so_far[:-1]
-        call_id_by_item_id: Final = self._function_call_ids_by_item_id(stream_events)
+        call_id_by_item_id: Final = self._tool_call_ids_by_item_id(stream_events)
         event_call_ids: Final = tuple(
-            self._function_call_event_call_id(event, call_id_by_item_id) for event in stream_events
+            self._tool_call_event_call_id(event, call_id_by_item_id) for event in stream_events
         )
         rewrites_by_call_id: Final = MappingProxyType(
             {
-                call_id: after
+                call_id: _tool_call_rewrite(before, after)
                 for call_id, before, after in zip(call_ids, pre_guardrail_tool_calls, post_guardrail_tool_calls)
                 if after != before
             }
         )
         unresolved_argument_event: Final = any(
-            call_id is None and stream_item_field(event, "type") in _FUNCTION_CALL_ARGUMENT_EVENT_TYPES
+            call_id is None and stream_item_field(event, "type") in _TOOL_CALL_PAYLOAD_EVENT_TYPES
             for event, call_id in zip(stream_events, event_call_ids)
         )
         if (
-            len(call_ids) != len(function_call_items)
+            len(call_ids) != len(tool_call_items)
             or len(frozenset(call_ids)) != len(call_ids)
             or len(call_ids) != len(post_guardrail_tool_calls)
             or unresolved_argument_event
@@ -981,10 +1065,10 @@ class OpenAIResponsesHandler(BaseTranslation):
             raise UndeliverableStreamRewrite(guardrail_name)
         for output_item, rewrite in (
             (output_item, rewrites_by_call_id[call_id])
-            for output_item, call_id in zip(function_call_items, call_ids)
+            for output_item, call_id in zip(tool_call_items, call_ids)
             if call_id in rewrites_by_call_id
         ):
-            self._write_function_call_item(output_item, rewrite.name, rewrite.arguments)
+            self._write_tool_call_item(output_item, rewrite.name, rewrite.arguments)
         delta_replacements: Final = MappingProxyType(
             {call_id: chain((rewrite.arguments,), repeat("")) for call_id, rewrite in rewrites_by_call_id.items()}
         )
@@ -992,16 +1076,18 @@ class OpenAIResponsesHandler(BaseTranslation):
             if call_id not in rewrites_by_call_id:
                 continue
             match stream_item_field(event, "type"):
-                case "response.function_call_arguments.delta":
+                case str() as event_type if event_type in _TOOL_CALL_PAYLOAD_DELTA_EVENT_TYPES:
                     self._write_event_field(event, "delta", next(delta_replacements[call_id]))
-                case "response.function_call_arguments.done":
-                    self._write_event_field(event, "arguments", rewrites_by_call_id[call_id].arguments)
+                case str() as event_type if event_type in _TOOL_CALL_PAYLOAD_DONE_EVENT_FIELDS:
+                    self._write_event_field(
+                        event, _TOOL_CALL_PAYLOAD_DONE_EVENT_FIELDS[event_type], rewrites_by_call_id[call_id].arguments
+                    )
                 case "response.output_item.added":
-                    self._write_function_call_item(
+                    self._write_tool_call_item(
                         stream_item_field(event, "item"), rewrites_by_call_id[call_id].name, None
                     )
                 case "response.output_item.done":
-                    self._write_function_call_item(
+                    self._write_tool_call_item(
                         stream_item_field(event, "item"),
                         rewrites_by_call_id[call_id].name,
                         rewrites_by_call_id[call_id].arguments,
@@ -1009,8 +1095,23 @@ class OpenAIResponsesHandler(BaseTranslation):
                 case _:
                     pass
 
+    def _write_tool_call_rewrites_to_output(
+        self,
+        tool_call_items: Sequence[object],
+        pre_guardrail_tool_calls: tuple[_ToolCallShape, ...],
+        post_guardrail_tool_calls: tuple[_ToolCallShape, ...],
+    ) -> None:
+        if len(tool_call_items) != len(post_guardrail_tool_calls):
+            return
+        for output_item, rewrite in (
+            (output_item, _tool_call_rewrite(before, after))
+            for output_item, before, after in zip(tool_call_items, pre_guardrail_tool_calls, post_guardrail_tool_calls)
+            if after != before
+        ):
+            self._write_tool_call_item(output_item, rewrite.name, rewrite.arguments)
+
     @staticmethod
-    def _function_call_ids_by_item_id(stream_events: Sequence[object]) -> Mapping[str, str]:
+    def _tool_call_ids_by_item_id(stream_events: Sequence[object]) -> Mapping[str, str]:
         items: Final = tuple(
             stream_item_field(event, "item")
             for event in stream_events
@@ -1020,32 +1121,35 @@ class OpenAIResponsesHandler(BaseTranslation):
             {
                 item_id: call_id
                 for item in items
-                if stream_item_field(item, "type") == "function_call"
+                if stream_item_field(item, "type") in _TOOL_CALL_ITEM_TYPES
                 and isinstance(item_id := stream_item_field(item, "id"), str)
                 and isinstance(call_id := stream_item_field(item, "call_id"), str)
             }
         )
 
     @staticmethod
-    def _function_call_event_call_id(event: object, call_id_by_item_id: Mapping[str, str]) -> str | None:
+    def _tool_call_event_call_id(event: object, call_id_by_item_id: Mapping[str, str]) -> str | None:
         event_type: Final = stream_item_field(event, "type")
-        if event_type in _FUNCTION_CALL_ARGUMENT_EVENT_TYPES:
+        if event_type in _TOOL_CALL_PAYLOAD_EVENT_TYPES:
             item_id: Final = stream_item_field(event, "item_id")
             return call_id_by_item_id.get(item_id) if isinstance(item_id, str) else None
         if event_type not in _OUTPUT_ITEM_EVENT_TYPES:
             return None
         item: Final = stream_item_field(event, "item")
         call_id: Final = stream_item_field(item, "call_id")
-        return call_id if stream_item_field(item, "type") == "function_call" and isinstance(call_id, str) else None
+        return (
+            call_id if stream_item_field(item, "type") in _TOOL_CALL_ITEM_TYPES and isinstance(call_id, str) else None
+        )
 
     @staticmethod
-    def _write_function_call_item(item: object, name: str | None, arguments: str | None) -> None:
+    def _write_tool_call_item(item: object, name: str | None, payload: str | None) -> None:
         if item is None:
             return
         if name is not None:
             OpenAIResponsesHandler._write_event_field(item, "name", name)
-        if arguments is not None:
-            OpenAIResponsesHandler._write_event_field(item, "arguments", arguments)
+        item_type: Final = stream_item_field(item, "type")
+        if payload is not None and isinstance(item_type, str) and item_type in _TOOL_CALL_PAYLOAD_FIELDS:
+            OpenAIResponsesHandler._write_event_field(item, _TOOL_CALL_PAYLOAD_FIELDS[item_type], payload)
 
     def _check_streaming_has_ended(self, responses_so_far: Sequence[object]) -> bool:
         """
@@ -1073,7 +1177,7 @@ class OpenAIResponsesHandler(BaseTranslation):
     def _completed_response_scan_key(response: object) -> StreamingScanKey:
         output_items: Final = stream_item_items(response, "output")
         message_items: Final = tuple(
-            item for item in output_items if stream_item_field(item, "type") != "function_call"
+            item for item in output_items if stream_item_field(item, "type") not in _TOOL_CALL_ITEM_TYPES
         )
         return StreamingScanKey(
             texts=tuple(
@@ -1085,7 +1189,7 @@ class OpenAIResponsesHandler(BaseTranslation):
             tool_calls=tuple(
                 stream_item_fingerprint(item)
                 for item in output_items
-                if stream_item_field(item, "type") == "function_call"
+                if stream_item_field(item, "type") in _TOOL_CALL_ITEM_TYPES
             ),
             stream_ended=True,
         )
@@ -1196,34 +1300,10 @@ class OpenAIResponsesHandler(BaseTranslation):
         Override this method to customize text/image/tool extraction logic.
         """
 
-        # Check if this is a tool call (OutputFunctionToolCall)
-        if isinstance(output_item, OutputFunctionToolCall) or (
-            isinstance(output_item, BaseModel)
-            and hasattr(output_item, "type")
-            and getattr(output_item, "type") == "function_call"
-        ):
+        tool_call_item: Final = _tool_call_output_item_mapping(output_item)
+        if tool_call_item is not None:
             if tool_calls_to_check is not None:
-                tool_call_dict = (
-                    LiteLLMCompletionResponsesConfig.convert_response_function_tool_call_to_chat_completion_tool_call(
-                        tool_call_item=output_item,
-                        index=output_idx,
-                    )
-                )
-                tool_calls_to_check.append(cast(ChatCompletionToolCallChunk, tool_call_dict))
-            return
-        elif isinstance(output_item, dict) and output_item.get("type") == "function_call":
-            # Handle dict representation of tool call
-            if tool_calls_to_check is not None:
-                # Convert dict to ResponseFunctionToolCall for processing
-                try:
-                    tool_call_obj: Final = ResponseFunctionToolCall(**output_item)
-                    tool_call_dict = LiteLLMCompletionResponsesConfig.convert_response_function_tool_call_to_chat_completion_tool_call(
-                        tool_call_item=tool_call_obj,
-                        index=output_idx,
-                    )
-                    tool_calls_to_check.append(cast(ChatCompletionToolCallChunk, tool_call_dict))
-                except Exception:
-                    pass
+                tool_calls_to_check.append(tool_call_dict_from_output_item(tool_call_item, output_idx))
             return
 
         # Handle both GenericResponseOutputItem and dict

tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py

@@ -10,11 +10,19 @@ from collections.abc import Callable
 from typing import Any, List, Literal, Optional, Tuple
 from unittest.mock import AsyncMock, MagicMock
 
+import logging
+
 import pytest
 
 
 from fastapi import HTTPException
-from openai.types.responses import ResponseFunctionToolCall
+from pydantic import BaseModel
+from openai.types.responses import (
+    ResponseCustomToolCall,
+    ResponseCustomToolCallInputDeltaEvent,
+    ResponseCustomToolCallInputDoneEvent,
+    ResponseFunctionToolCall,
+)
 
 from litellm.integrations.custom_guardrail import CustomGuardrail
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
@@ -23,11 +31,12 @@ from litellm.llms.openai.responses.guardrail_translation.handler import (
     OpenAIResponsesHandler,
 )
 from litellm.llms.openai.responses.guardrail_translation.tool_merge import merge_guardrailed_tools
+from litellm.types.llms.openai import ChatCompletionToolCallChunk
 from litellm.responses.litellm_completion_transformation.transformation import (
     LiteLLMCompletionResponsesConfig,
 )
 from litellm.types.llms.openai import ResponsesAPIResponse
-from litellm.types.responses.main import GenericResponseOutputItem, OutputText
+from litellm.types.responses.main import CustomToolCallOutputItem, GenericResponseOutputItem, OutputText
 from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs
 
 
@@ -57,6 +66,60 @@ class MockGuardrail(CustomGuardrail):
         return inputs
 
 
+class PersimmonMaskingGuardrail(CustomGuardrail):
+    async def apply_guardrail(
+        self,
+        inputs: GenericGuardrailAPIInputs,
+        request_data: dict,
+        input_type: Literal["request", "response"],
+        logging_obj: Optional[LiteLLMLoggingObj] = None,
+    ) -> GenericGuardrailAPIInputs:
+        tool_calls = [
+            {
+                **tool_call,
+                "function": {
+                    **tool_call["function"],
+                    "arguments": tool_call["function"]["arguments"].replace("persimmon", "[MASKED]"),
+                },
+            }
+            for tool_call in inputs.get("tool_calls", [])
+        ]
+        return {**inputs, "tool_calls": tool_calls}
+
+
+class FlatShapeGuardrail(CustomGuardrail):
+    async def apply_guardrail(
+        self,
+        inputs: GenericGuardrailAPIInputs,
+        request_data: dict,
+        input_type: Literal["request", "response"],
+        logging_obj: Optional[LiteLLMLoggingObj] = None,
+    ) -> GenericGuardrailAPIInputs:
+        flat_tool_calls = [{"name": "exec", "input": "rm -rf /"} for _ in inputs.get("tool_calls", [])]
+        return {**inputs, "tool_calls": flat_tool_calls}
+
+
+class DroppingGuardrail(CustomGuardrail):
+    async def apply_guardrail(
+        self,
+        inputs: GenericGuardrailAPIInputs,
+        request_data: dict,
+        input_type: Literal["request", "response"],
+        logging_obj: Optional[LiteLLMLoggingObj] = None,
+    ) -> GenericGuardrailAPIInputs:
+        return {**inputs, "tool_calls": []}
+
+
+CUSTOM_TOOL_CALL_ITEM = {
+    "type": "custom_tool_call",
+    "id": "ctc_1",
+    "call_id": "call_exec_1",
+    "name": "exec",
+    "input": "echo persimmon",
+    "status": "completed",
+}
+
+
 class TestOpenAIResponsesHandlerDiscovery:
     """Test that the handler is properly discovered by the guardrail system"""
 
@@ -557,7 +620,7 @@ class TestOpenAIResponsesHandlerToolCallExtraction:
 
         texts_to_check: List[str] = []
         images_to_check: List[str] = []
-        tool_calls_to_check: List[Any] = []
+        tool_calls_to_check: List[ChatCompletionToolCallChunk] = []
         task_mappings: List[Tuple[int, int]] = []
 
         # Extract tool calls
@@ -628,6 +691,123 @@ class TestOpenAIResponsesHandlerToolCallExtraction:
             == '{"location":"Boston, MA","unit":"celsius"}'
         )
 
+    @pytest.mark.parametrize(
+        "output_item",
+        [
+            dict(CUSTOM_TOOL_CALL_ITEM),
+            CustomToolCallOutputItem(**CUSTOM_TOOL_CALL_ITEM),
+            ResponseCustomToolCall(**{key: value for key, value in CUSTOM_TOOL_CALL_ITEM.items() if key != "status"}),
+        ],
+        ids=["dict", "litellm_typed", "openai_typed"],
+    )
+    def test_extract_custom_tool_call_input_as_arguments(self, output_item):
+        handler = OpenAIResponsesHandler()
+        texts_to_check: List[str] = []
+        tool_calls_to_check: List[Any] = []
+
+        handler._extract_output_text_and_images(
+            output_item=output_item,
+            output_idx=2,
+            texts_to_check=texts_to_check,
+            images_to_check=[],
+            task_mappings=[],
+            tool_calls_to_check=tool_calls_to_check,
+        )
+
+        assert texts_to_check == []
+        assert tool_calls_to_check == [
+            {
+                "id": "call_exec_1",
+                "type": "function",
+                "function": {"name": "exec", "arguments": "echo persimmon"},
+                "index": 2,
+            }
+        ]
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("typed", [False, True], ids=["dict", "typed"])
+    async def test_process_output_response_writes_tool_call_rewrites_back(self, typed):
+        handler = OpenAIResponsesHandler()
+        function_call = {
+            "type": "function_call",
+            "id": "fc_1",
+            "call_id": "call_fn_1",
+            "name": "lookup_fruit",
+            "arguments": '{"fruit": "persimmon"}',
+            "status": "completed",
+        }
+        message = {
+            "type": "message",
+            "id": "msg_1",
+            "role": "assistant",
+            "status": "completed",
+            "content": [{"type": "output_text", "text": "running persimmon", "annotations": []}],
+        }
+        payload = {
+            "id": "resp_1",
+            "created_at": 1,
+            "model": "gpt-5.6",
+            "object": "response",
+            "status": "completed",
+            "output": [message, function_call, dict(CUSTOM_TOOL_CALL_ITEM)],
+        }
+        response = ResponsesAPIResponse.model_validate(payload) if typed else payload
+
+        result = await handler.process_output_response(response, PersimmonMaskingGuardrail(guardrail_name="mask"))
+
+        output = result.output if typed else result["output"]
+        function_item, custom_item = output[1], output[2]
+        assert (function_item.arguments if typed else function_item["arguments"]) == '{"fruit": "[MASKED]"}'
+        assert (custom_item.input if typed else custom_item["input"]) == "echo [MASKED]"
+        assert (custom_item.name if typed else custom_item["name"]) == "exec"
+        assert (output[0].content[0].text if typed else output[0]["content"][0]["text"]) == "running persimmon"
+
+    @staticmethod
+    def _custom_tool_call_response(item: dict) -> dict:
+        return {
+            "id": "resp_1",
+            "created_at": 1,
+            "model": "gpt-5.6",
+            "object": "response",
+            "status": "completed",
+            "output": [item],
+        }
+
+    @pytest.mark.asyncio
+    async def test_process_output_response_ignores_tool_call_rewrites_in_another_shape(self):
+        handler = OpenAIResponsesHandler()
+        response = self._custom_tool_call_response(dict(CUSTOM_TOOL_CALL_ITEM))
+
+        result = await handler.process_output_response(response, FlatShapeGuardrail(guardrail_name="flat"))
+
+        assert result["output"][0]["input"] == "echo persimmon"
+        assert result["output"][0]["name"] == "exec"
+
+    @pytest.mark.asyncio
+    async def test_process_output_response_warns_when_guardrail_drops_tool_calls(self, caplog):
+        handler = OpenAIResponsesHandler()
+        response = self._custom_tool_call_response(dict(CUSTOM_TOOL_CALL_ITEM))
+
+        with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
+            result = await handler.process_output_response(response, DroppingGuardrail(guardrail_name="dropper"))
+
+        assert result["output"][0]["input"] == "echo persimmon"
+        assert any(
+            "dropper" in record.getMessage() and "0 tool calls for the 1 scanned" in record.getMessage()
+            for record in caplog.records
+        )
+
+    @pytest.mark.asyncio
+    async def test_process_output_response_keeps_a_nameless_custom_tool_call_nameless(self):
+        handler = OpenAIResponsesHandler()
+        nameless_item = {key: value for key, value in CUSTOM_TOOL_CALL_ITEM.items() if key != "name"}
+        response = self._custom_tool_call_response(nameless_item)
+
+        result = await handler.process_output_response(response, PersimmonMaskingGuardrail(guardrail_name="mask"))
+
+        assert result["output"][0]["input"] == "echo [MASKED]"
+        assert "name" not in result["output"][0]
+
     @pytest.mark.asyncio
     async def test_process_output_response_with_tool_calls(self):
         """Test processing output response containing function tool calls"""
@@ -1315,6 +1495,128 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing:
         assert completed_event.response.output[0].arguments == '{"fruit": "[MASKED]"}'
         assert completed_event.response.output[0].name == "lookup_fruit"
 
+    @staticmethod
+    def _ended_custom_tool_call_stream_events() -> List[dict]:
+        def item(input_text: str, status: str) -> dict:
+            return {**CUSTOM_TOOL_CALL_ITEM, "input": input_text, "status": status}
+
+        return [
+            {"type": "response.output_item.added", "output_index": 0, "item": item("", "in_progress")},
+            {"type": "response.custom_tool_call_input.delta", "item_id": "ctc_1", "output_index": 0, "delta": "echo "},
+            {"type": "response.custom_tool_call_input.delta", "item_id": "ctc_1", "output_index": 0, "delta": "persimmon"},
+            {"type": "response.custom_tool_call_input.done", "item_id": "ctc_1", "output_index": 0, "input": "echo persimmon"},
+            {"type": "response.output_item.done", "output_index": 0, "item": item("echo persimmon", "completed")},
+            {
+                "type": "response.completed",
+                "response": {
+                    "id": "resp_123",
+                    "created_at": 1,
+                    "model": "gpt-5.6",
+                    "output": [item("echo persimmon", "completed")],
+                    "status": "completed",
+                },
+            },
+        ]
+
+    @pytest.mark.asyncio
+    async def test_deliver_ended_stream_rewrites_syncs_custom_tool_call_events(self):
+        handler = OpenAIResponsesHandler()
+        events = self._ended_custom_tool_call_stream_events()
+
+        result = await handler.process_output_streaming_response(
+            responses_so_far=events,
+            guardrail_to_apply=PersimmonMaskingGuardrail(guardrail_name="mask"),
+            litellm_logging_obj=None,
+            deliver_ended_stream_rewrites=True,
+        )
+
+        assert result is events
+        assert events[0]["item"]["input"] == ""
+        assert events[1]["delta"] == "echo [MASKED]"
+        assert events[2]["delta"] == ""
+        assert events[3]["input"] == "echo [MASKED]"
+        assert events[4]["item"]["input"] == "echo [MASKED]"
+        assert events[5]["response"]["output"][0]["input"] == "echo [MASKED]"
+        assert events[5]["response"]["output"][0]["name"] == "exec"
+        assert "arguments" not in events[5]["response"]["output"][0]
+
+    @pytest.mark.asyncio
+    async def test_deliver_ended_stream_rewrites_keep_a_nameless_custom_tool_call_nameless(self):
+        handler = OpenAIResponsesHandler()
+        events = self._ended_custom_tool_call_stream_events()
+        items = [events[0]["item"], events[4]["item"], events[5]["response"]["output"][0]]
+        for item in items:
+            del item["name"]
+
+        await handler.process_output_streaming_response(
+            responses_so_far=events,
+            guardrail_to_apply=PersimmonMaskingGuardrail(guardrail_name="mask"),
+            litellm_logging_obj=None,
+            deliver_ended_stream_rewrites=True,
+        )
+
+        assert events[3]["input"] == "echo [MASKED]"
+        assert events[5]["response"]["output"][0]["input"] == "echo [MASKED]"
+        assert all("name" not in item for item in items)
+
+    @pytest.mark.asyncio
+    async def test_deliver_ended_stream_rewrites_syncs_typed_custom_tool_call_events(self):
+        from litellm.types.llms.openai import (
+            OutputItemAddedEvent,
+            OutputItemDoneEvent,
+            ResponseCompletedEvent,
+        )
+
+        handler = OpenAIResponsesHandler()
+        typed_events: List[BaseModel] = [
+            model.model_validate({**event, "sequence_number": sequence_number})
+            for sequence_number, (model, event) in enumerate(
+                zip(
+                    (
+                        OutputItemAddedEvent,
+                        ResponseCustomToolCallInputDeltaEvent,
+                        ResponseCustomToolCallInputDeltaEvent,
+                        ResponseCustomToolCallInputDoneEvent,
+                        OutputItemDoneEvent,
+                        ResponseCompletedEvent,
+                    ),
+                    self._ended_custom_tool_call_stream_events(),
+                )
+            )
+        ]
+        completed_event = typed_events[5]
+        assert isinstance(completed_event.response.output[0], CustomToolCallOutputItem)
+
+        await handler.process_output_streaming_response(
+            responses_so_far=typed_events,
+            guardrail_to_apply=PersimmonMaskingGuardrail(guardrail_name="mask"),
+            litellm_logging_obj=None,
+            deliver_ended_stream_rewrites=True,
+        )
+
+        assert typed_events[1].delta == "echo [MASKED]"
+        assert typed_events[2].delta == ""
+        assert typed_events[3].input == "echo [MASKED]"
+        assert typed_events[4].item.input == "echo [MASKED]"
+        assert completed_event.response.output[0].input == "echo [MASKED]"
+        assert completed_event.response.output[0].name == "exec"
+
+    @pytest.mark.asyncio
+    async def test_deliver_ended_stream_custom_tool_call_rewrite_without_matching_events_fails_closed(self):
+        from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
+
+        handler = OpenAIResponsesHandler()
+        events = self._ended_custom_tool_call_stream_events()
+        events[5]["response"]["output"] = [{**events[5]["response"]["output"][0], "call_id": "call_999"}]
+
+        with pytest.raises(UndeliverableStreamRewrite):
+            await handler.process_output_streaming_response(
+                responses_so_far=events,
+                guardrail_to_apply=PersimmonMaskingGuardrail(guardrail_name="mask"),
+                litellm_logging_obj=None,
+                deliver_ended_stream_rewrites=True,
+            )
+
     @staticmethod
     def _bridged_function_call_stream_events() -> List[dict]:
         reasoning = {"type": "reasoning", "id": "rs_1", "summary": []}
@@ -2747,8 +3049,21 @@ class TestOpenAIResponsesHandlerStreamingScanKey:
         assert len(ended_key.tool_calls) == 1 and "get_weather" in ended_key.tool_calls[0]
         assert ended_key != open_key
 
+    def test_completed_event_with_a_custom_tool_call_changes_the_key(self):
+        handler = OpenAIResponsesHandler()
+        message = {"type": "message", "content": [{"type": "output_text", "text": "hi"}]}
+        ended_key = handler.get_streaming_scan_key(
+            [self._delta(0, "hi"), self._completed(1, [message, dict(CUSTOM_TOOL_CALL_ITEM)])]
+        )
+        rewritten_key = handler.get_streaming_scan_key(
+            [self._delta(0, "hi"), self._completed(1, [message, {**CUSTOM_TOOL_CALL_ITEM, "input": "echo kumquat"}])]
+        )
+        assert ended_key.texts == ("hi",)
+        assert len(ended_key.tool_calls) == 1 and "echo persimmon" in ended_key.tool_calls[0]
+        assert rewritten_key != ended_key
+
     def test_completed_event_reads_every_output_text_part(self):
-        from litellm.types.responses.main import GenericResponseOutputItem, OutputText
+        from litellm.types.responses.main import CustomToolCallOutputItem, GenericResponseOutputItem, OutputText
 
         item = GenericResponseOutputItem(
             type="message",