feat: add SDK-only custom data for tool outputs (#3486)

Author: seratch

src/agents/__init__.py

@@ -133,12 +133,20 @@
 )
 from .tool import (
     ApplyPatchTool,
+    ApplyPatchToolCustomDataContext,
+    ApplyPatchToolCustomDataExtractor,
     CodeInterpreterTool,
     ComputerProvider,
     ComputerTool,
+    ComputerToolCustomDataContext,
+    ComputerToolCustomDataExtractor,
     CustomTool,
+    CustomToolCustomDataContext,
+    CustomToolCustomDataExtractor,
     FileSearchTool,
     FunctionTool,
+    FunctionToolCustomDataContext,
+    FunctionToolCustomDataExtractor,
     FunctionToolResult,
     HostedMCPTool,
     ImageGenerationTool,
@@ -458,10 +466,16 @@ def enable_verbose_stdout_logging():
     "AgentUpdatedStreamEvent",
     "StreamEvent",
     "FunctionTool",
+    "FunctionToolCustomDataContext",
+    "FunctionToolCustomDataExtractor",
     "FunctionToolResult",
     "ComputerTool",
+    "ComputerToolCustomDataContext",
+    "ComputerToolCustomDataExtractor",
     "ComputerProvider",
     "CustomTool",
+    "CustomToolCustomDataContext",
+    "CustomToolCustomDataExtractor",
     "FileSearchTool",
     "CodeInterpreterTool",
     "ImageGenerationTool",
@@ -494,6 +508,8 @@ def enable_verbose_stdout_logging():
     "ApplyPatchOperation",
     "ApplyPatchResult",
     "ApplyPatchTool",
+    "ApplyPatchToolCustomDataContext",
+    "ApplyPatchToolCustomDataExtractor",
     "Tool",
     "WebSearchTool",
     "HostedMCPTool",

src/agents/items.py

@@ -403,6 +403,13 @@ class ToolCallOutputItem(RunItemBase[Any]):
     tool_origin: ToolOrigin | None = None
     """Optional metadata describing the source of a function-tool-backed item."""
 
+    custom_data: dict[str, Any] | None = None
+    """SDK-only custom data attached to this tool output.
+
+    This data is not part of ``raw_item`` and is not sent back to the model when the output item is
+    replayed as input.
+    """
+
     @property
     def call_id(self) -> str | None:
         """Return the call identifier from the raw item, if available."""

src/agents/mcp/__init__.py

@@ -17,6 +17,8 @@
     )
 
 from .util import (
+    MCPToolCustomDataContext,
+    MCPToolCustomDataExtractor,
     MCPToolMetaContext,
     MCPToolMetaResolver,
     MCPUtil,
@@ -50,6 +52,8 @@
     "MCPServerManager",
     "LocalMCPApprovalCallable",
     "MCPUtil",
+    "MCPToolCustomDataContext",
+    "MCPToolCustomDataExtractor",
     "MCPToolMetaContext",
     "MCPToolMetaResolver",
     "ToolFilter",

src/agents/mcp/server.py

@@ -45,6 +45,7 @@
 from ..util._types import MaybeAwaitable
 from .util import (
     HttpClientFactory,
+    MCPToolCustomDataExtractor,
     MCPToolMetaResolver,
     ToolFilter,
     ToolFilterContext,
@@ -229,6 +230,7 @@ def __init__(
         require_approval: RequireApprovalSetting = None,
         failure_error_function: ToolErrorFunction | None | _UnsetType = _UNSET,
         tool_meta_resolver: MCPToolMetaResolver | None = None,
+        custom_data_extractor: MCPToolCustomDataExtractor | None = None,
     ):
         """
         Args:
@@ -248,13 +250,16 @@ def __init__(
                 SDK default) will be used.
             tool_meta_resolver: Optional callable that produces MCP request metadata (`_meta`) for
                 tool calls. It is invoked by the Agents SDK before calling `call_tool`.
+            custom_data_extractor: Optional callable that produces SDK-only custom data for
+                emitted MCP tool output items.
         """
         self.use_structured_content = use_structured_content
         self._needs_approval_policy = self._normalize_needs_approval(
             require_approval=require_approval
         )
         self._failure_error_function = failure_error_function
         self.tool_meta_resolver = tool_meta_resolver
+        self.custom_data_extractor = custom_data_extractor
 
     @abc.abstractmethod
     async def connect(self):
@@ -544,6 +549,7 @@ def __init__(
         require_approval: RequireApprovalSetting = None,
         failure_error_function: ToolErrorFunction | None | _UnsetType = _UNSET,
         tool_meta_resolver: MCPToolMetaResolver | None = None,
+        custom_data_extractor: MCPToolCustomDataExtractor | None = None,
     ):
         """
         Args:
@@ -576,12 +582,15 @@ def __init__(
                 SDK default) will be used.
             tool_meta_resolver: Optional callable that produces MCP request metadata (`_meta`) for
                 tool calls. It is invoked by the Agents SDK before calling `call_tool`.
+            custom_data_extractor: Optional callable that produces SDK-only custom data for
+                emitted MCP tool output items.
         """
         super().__init__(
             use_structured_content=use_structured_content,
             require_approval=require_approval,
             failure_error_function=failure_error_function,
             tool_meta_resolver=tool_meta_resolver,
+            custom_data_extractor=custom_data_extractor,
         )
         self.session: ClientSession | None = None
         self.exit_stack: AsyncExitStack = AsyncExitStack()
@@ -1108,6 +1117,7 @@ def __init__(
         require_approval: RequireApprovalSetting = None,
         failure_error_function: ToolErrorFunction | None | _UnsetType = _UNSET,
         tool_meta_resolver: MCPToolMetaResolver | None = None,
+        custom_data_extractor: MCPToolCustomDataExtractor | None = None,
     ):
         """Create a new MCP server based on the stdio transport.
 
@@ -1145,6 +1155,8 @@ def __init__(
                 SDK default) will be used.
             tool_meta_resolver: Optional callable that produces MCP request metadata (`_meta`) for
                 tool calls. It is invoked by the Agents SDK before calling `call_tool`.
+            custom_data_extractor: Optional callable that produces SDK-only custom data for
+                emitted MCP tool output items.
         """
         super().__init__(
             cache_tools_list=cache_tools_list,
@@ -1157,6 +1169,7 @@ def __init__(
             require_approval=require_approval,
             failure_error_function=failure_error_function,
             tool_meta_resolver=tool_meta_resolver,
+            custom_data_extractor=custom_data_extractor,
         )
 
         self.params = StdioServerParameters(
@@ -1229,6 +1242,7 @@ def __init__(
         require_approval: RequireApprovalSetting = None,
         failure_error_function: ToolErrorFunction | None | _UnsetType = _UNSET,
         tool_meta_resolver: MCPToolMetaResolver | None = None,
+        custom_data_extractor: MCPToolCustomDataExtractor | None = None,
     ):
         """Create a new MCP server based on the HTTP with SSE transport.
 
@@ -1268,6 +1282,8 @@ def __init__(
                 SDK default) will be used.
             tool_meta_resolver: Optional callable that produces MCP request metadata (`_meta`) for
                 tool calls. It is invoked by the Agents SDK before calling `call_tool`.
+            custom_data_extractor: Optional callable that produces SDK-only custom data for
+                emitted MCP tool output items.
         """
         super().__init__(
             cache_tools_list=cache_tools_list,
@@ -1280,6 +1296,7 @@ def __init__(
             require_approval=require_approval,
             failure_error_function=failure_error_function,
             tool_meta_resolver=tool_meta_resolver,
+            custom_data_extractor=custom_data_extractor,
         )
 
         self.params = params
@@ -1365,6 +1382,7 @@ def __init__(
         require_approval: RequireApprovalSetting = None,
         failure_error_function: ToolErrorFunction | None | _UnsetType = _UNSET,
         tool_meta_resolver: MCPToolMetaResolver | None = None,
+        custom_data_extractor: MCPToolCustomDataExtractor | None = None,
     ):
         """Create a new MCP server based on the Streamable HTTP transport.
 
@@ -1405,6 +1423,8 @@ def __init__(
                 SDK default) will be used.
             tool_meta_resolver: Optional callable that produces MCP request metadata (`_meta`) for
                 tool calls. It is invoked by the Agents SDK before calling `call_tool`.
+            custom_data_extractor: Optional callable that produces SDK-only custom data for
+                emitted MCP tool output items.
         """
         super().__init__(
             cache_tools_list=cache_tools_list,
@@ -1417,6 +1437,7 @@ def __init__(
             require_approval=require_approval,
             failure_error_function=failure_error_function,
             tool_meta_resolver=tool_meta_resolver,
+            custom_data_extractor=custom_data_extractor,
         )
 
         self.params = params

src/agents/mcp/util.py

@@ -7,8 +7,9 @@
 import inspect
 import json
 from collections import Counter
-from collections.abc import Awaitable, Callable
+from collections.abc import Awaitable, Callable, Mapping
 from dataclasses import dataclass
+from types import MappingProxyType
 from typing import TYPE_CHECKING, Any, Protocol, Union
 
 import httpx
@@ -39,6 +40,7 @@
 )
 from ..tool_context import ToolContext
 from ..tracing import FunctionSpanData, get_current_span, mcp_tools_span
+from ..util._custom_data import maybe_extract_custom_data
 from ..util._types import MaybeAwaitable
 
 if TYPE_CHECKING:
@@ -149,13 +151,50 @@ class MCPToolMetaContext:
     """The parsed tool arguments."""
 
 
+@dataclass(frozen=True)
+class MCPToolCustomDataContext:
+    """Context passed to MCP tool custom data extractors."""
+
+    run_context: RunContextWrapper[Any]
+    """The current run context."""
+
+    server_name: str
+    """The name of the MCP server."""
+
+    tool_name: str
+    """The original MCP tool name invoked on the server."""
+
+    tool_display_name: str
+    """The public tool name exposed through the Agents SDK."""
+
+    arguments: Mapping[str, Any]
+    """The parsed tool arguments."""
+
+    result_meta: Mapping[str, Any] | None
+    """The MCP tool result ``_meta`` payload, if present."""
+
+    structured_content: Mapping[str, Any] | None
+    """The MCP tool result ``structuredContent`` payload, if present."""
+
+    is_error: bool | None
+    """The MCP tool result ``isError`` flag, if present."""
+
+    tool_output: ToolOutput
+    """The model-visible tool output produced by the Agents SDK."""
+
+
 if TYPE_CHECKING:
     MCPToolMetaResolver = Callable[
         [MCPToolMetaContext],
         MaybeAwaitable[dict[str, Any] | None],
     ]
+    MCPToolCustomDataExtractor = Callable[
+        [MCPToolCustomDataContext],
+        MaybeAwaitable[Mapping[str, Any] | None],
+    ]
 else:
     MCPToolMetaResolver = Callable[..., Any]
+    MCPToolCustomDataExtractor = Callable[..., Any]
 """A function that produces MCP request metadata for tool calls.
 
 Args:
@@ -164,6 +203,7 @@ class MCPToolMetaContext:
 Returns:
     A dict to send as MCP `_meta`, or None to omit metadata.
 """
+"""A function that produces SDK-only custom data for MCP tool output items."""
 
 
 def create_static_tool_filter(
@@ -541,6 +581,41 @@ def _merge_mcp_meta(
             merged.update(copy.deepcopy(explicit_meta))
         return merged
 
+    @staticmethod
+    def _copy_mapping_proxy(value: Any) -> Mapping[str, Any] | None:
+        if not isinstance(value, dict):
+            return None
+        return MappingProxyType(copy.deepcopy(value))
+
+    @classmethod
+    async def _extract_custom_data(
+        cls,
+        *,
+        server: MCPServer,
+        context: RunContextWrapper[Any],
+        tool_name: str,
+        tool_display_name: str,
+        arguments: dict[str, Any],
+        result: Any,
+        tool_output: ToolOutput,
+    ) -> dict[str, Any] | None:
+        extractor = getattr(server, "custom_data_extractor", None)
+        if extractor is None:
+            return None
+
+        extractor_context = MCPToolCustomDataContext(
+            run_context=context,
+            server_name=server.name,
+            tool_name=tool_name,
+            tool_display_name=tool_display_name,
+            arguments=MappingProxyType(copy.deepcopy(arguments)),
+            result_meta=cls._copy_mapping_proxy(getattr(result, "meta", None)),
+            structured_content=cls._copy_mapping_proxy(getattr(result, "structuredContent", None)),
+            is_error=getattr(result, "isError", None),
+            tool_output=copy.deepcopy(tool_output),
+        )
+        return await maybe_extract_custom_data(extractor, extractor_context)
+
     @classmethod
     async def _resolve_meta(
         cls,
@@ -688,6 +763,18 @@ async def invoke_mcp_tool(
             else:
                 tool_output = tool_output_list
 
+        custom_data = await cls._extract_custom_data(
+            server=server,
+            context=context,
+            tool_name=tool.name,
+            tool_display_name=tool_name_for_display,
+            arguments=json_data,
+            result=result,
+            tool_output=tool_output,
+        )
+        if custom_data and isinstance(context, ToolContext):
+            context._custom_data = custom_data
+
         current_span = get_current_span()
         if current_span:
             if isinstance(current_span.span_data, FunctionSpanData):