Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 16 additions & 0 deletions src/agents/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -130,12 +130,20 @@
)
from .tool import (
ApplyPatchTool,
ApplyPatchToolCustomDataContext,
ApplyPatchToolCustomDataExtractor,
CodeInterpreterTool,
ComputerProvider,
ComputerTool,
ComputerToolCustomDataContext,
ComputerToolCustomDataExtractor,
CustomTool,
CustomToolCustomDataContext,
CustomToolCustomDataExtractor,
FileSearchTool,
FunctionTool,
FunctionToolCustomDataContext,
FunctionToolCustomDataExtractor,
FunctionToolResult,
HostedMCPTool,
ImageGenerationTool,
Expand Down Expand Up @@ -446,10 +454,16 @@ def enable_verbose_stdout_logging():
"AgentUpdatedStreamEvent",
"StreamEvent",
"FunctionTool",
"FunctionToolCustomDataContext",
"FunctionToolCustomDataExtractor",
"FunctionToolResult",
"ComputerTool",
"ComputerToolCustomDataContext",
"ComputerToolCustomDataExtractor",
"ComputerProvider",
"CustomTool",
"CustomToolCustomDataContext",
"CustomToolCustomDataExtractor",
"FileSearchTool",
"CodeInterpreterTool",
"ImageGenerationTool",
Expand Down Expand Up @@ -482,6 +496,8 @@ def enable_verbose_stdout_logging():
"ApplyPatchOperation",
"ApplyPatchResult",
"ApplyPatchTool",
"ApplyPatchToolCustomDataContext",
"ApplyPatchToolCustomDataExtractor",
"Tool",
"WebSearchTool",
"HostedMCPTool",
Expand Down
7 changes: 7 additions & 0 deletions src/agents/items.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand Down
4 changes: 4 additions & 0 deletions src/agents/mcp/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,8 @@
)

from .util import (
MCPToolCustomDataContext,
MCPToolCustomDataExtractor,
MCPToolMetaContext,
MCPToolMetaResolver,
MCPUtil,
Expand Down Expand Up @@ -50,6 +52,8 @@
"MCPServerManager",
"LocalMCPApprovalCallable",
"MCPUtil",
"MCPToolCustomDataContext",
"MCPToolCustomDataExtractor",
"MCPToolMetaContext",
"MCPToolMetaResolver",
"ToolFilter",
Expand Down
21 changes: 21 additions & 0 deletions src/agents/mcp/server.py
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,7 @@
from ..util._types import MaybeAwaitable
from .util import (
HttpClientFactory,
MCPToolCustomDataExtractor,
MCPToolMetaResolver,
ToolFilter,
ToolFilterContext,
Expand Down Expand Up @@ -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:
Expand All @@ -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):
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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.

Expand Down Expand Up @@ -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,
Expand All @@ -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(
Expand Down Expand Up @@ -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.

Expand Down Expand Up @@ -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,
Expand All @@ -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
Expand Down Expand Up @@ -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.

Expand Down Expand Up @@ -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,
Expand All @@ -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
Expand Down
89 changes: 88 additions & 1 deletion src/agents/mcp/util.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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:
Expand All @@ -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(
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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):
Expand Down
Loading
Loading