66"""
77
88import logging
9+ from datetime import timedelta
910from typing import TYPE_CHECKING , Any
1011
1112from mcp .types import Tool as MCPTool
@@ -28,20 +29,28 @@ class MCPAgentTool(AgentTool):
2829 seamlessly within the agent framework.
2930 """
3031
31- def __init__ (self , mcp_tool : MCPTool , mcp_client : "MCPClient" , name_override : str | None = None ) -> None :
32+ def __init__ (
33+ self ,
34+ mcp_tool : MCPTool ,
35+ mcp_client : "MCPClient" ,
36+ name_override : str | None = None ,
37+ timeout : timedelta | None = None ,
38+ ) -> None :
3239 """Initialize a new MCPAgentTool instance.
3340
3441 Args:
3542 mcp_tool: The MCP tool to adapt
3643 mcp_client: The MCP server connection to use for tool invocation
3744 name_override: Optional name to use for the agent tool (for disambiguation)
3845 If None, uses the original MCP tool name
46+ timeout: Optional timeout duration for tool execution
3947 """
4048 super ().__init__ ()
4149 logger .debug ("tool_name=<%s> | creating mcp agent tool" , mcp_tool .name )
4250 self .mcp_tool = mcp_tool
4351 self .mcp_client = mcp_client
4452 self ._agent_tool_name = name_override or mcp_tool .name
53+ self ._timeout_seconds = timeout
4554
4655 @property
4756 def tool_name (self ) -> str :
@@ -105,5 +114,6 @@ async def stream(self, tool_use: ToolUse, invocation_state: dict[str, Any], **kw
105114 tool_use_id = tool_use ["toolUseId" ],
106115 name = self .mcp_tool .name , # Use original MCP name for server communication
107116 arguments = tool_use ["input" ],
117+ read_timeout_seconds = self ._timeout_seconds ,
108118 )
109119 yield ToolResultEvent (result )
0 commit comments