mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-05-03 00:02:36 +00:00
feat: allow custom client_timeout for MCPAdapter (#409)
This commit is contained in:
@@ -50,13 +50,17 @@ class MCPServerAdapter:
|
|||||||
with MCPServerAdapter(..., "tool1", "tool2") as filtered_tools:
|
with MCPServerAdapter(..., "tool1", "tool2") as filtered_tools:
|
||||||
# only tool1 and tool2 are available
|
# only tool1 and tool2 are available
|
||||||
|
|
||||||
|
# context manager with custom connect timeout (60 seconds)
|
||||||
|
with MCPServerAdapter(..., connect_timeout=60) as tools:
|
||||||
|
# tools is now available with longer timeout
|
||||||
|
|
||||||
# manually stop mcp server
|
# manually stop mcp server
|
||||||
try:
|
try:
|
||||||
mcp_server = MCPServerAdapter(...)
|
mcp_server = MCPServerAdapter(...)
|
||||||
tools = mcp_server.tools # all tools
|
tools = mcp_server.tools # all tools
|
||||||
|
|
||||||
# or with filtered tools
|
# or with filtered tools and custom timeout
|
||||||
mcp_server = MCPServerAdapter(..., "tool1", "tool2")
|
mcp_server = MCPServerAdapter(..., "tool1", "tool2", connect_timeout=45)
|
||||||
filtered_tools = mcp_server.tools # only tool1 and tool2
|
filtered_tools = mcp_server.tools # only tool1 and tool2
|
||||||
...
|
...
|
||||||
finally:
|
finally:
|
||||||
@@ -70,6 +74,7 @@ class MCPServerAdapter:
|
|||||||
self,
|
self,
|
||||||
serverparams: StdioServerParameters | dict[str, Any],
|
serverparams: StdioServerParameters | dict[str, Any],
|
||||||
*tool_names: str,
|
*tool_names: str,
|
||||||
|
connect_timeout: int = 30,
|
||||||
):
|
):
|
||||||
"""Initialize the MCP Server
|
"""Initialize the MCP Server
|
||||||
|
|
||||||
@@ -78,6 +83,7 @@ class MCPServerAdapter:
|
|||||||
`StdioServerParameters` or a `dict` respectively for STDIO and SSE.
|
`StdioServerParameters` or a `dict` respectively for STDIO and SSE.
|
||||||
*tool_names: Optional names of tools to filter. If provided, only tools with
|
*tool_names: Optional names of tools to filter. If provided, only tools with
|
||||||
matching names will be available.
|
matching names will be available.
|
||||||
|
connect_timeout: Connection timeout in seconds to the MCP server (default is 30s).
|
||||||
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
@@ -106,7 +112,7 @@ class MCPServerAdapter:
|
|||||||
|
|
||||||
try:
|
try:
|
||||||
self._serverparams = serverparams
|
self._serverparams = serverparams
|
||||||
self._adapter = MCPAdapt(self._serverparams, CrewAIAdapter())
|
self._adapter = MCPAdapt(self._serverparams, CrewAIAdapter(), connect_timeout)
|
||||||
self.start()
|
self.start()
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
|
|||||||
@@ -1,4 +1,5 @@
|
|||||||
from textwrap import dedent
|
from textwrap import dedent
|
||||||
|
from unittest.mock import MagicMock, patch
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
from mcp import StdioServerParameters
|
from mcp import StdioServerParameters
|
||||||
@@ -187,3 +188,43 @@ def test_filter_with_only_nonexistent_tools(echo_server_script):
|
|||||||
# Should return an empty tool collection
|
# Should return an empty tool collection
|
||||||
assert isinstance(tools, ToolCollection)
|
assert isinstance(tools, ToolCollection)
|
||||||
assert len(tools) == 0
|
assert len(tools) == 0
|
||||||
|
|
||||||
|
def test_connect_timeout_parameter(echo_server_script):
|
||||||
|
serverparams = StdioServerParameters(
|
||||||
|
command="uv", args=["run", "python", "-c", echo_server_script]
|
||||||
|
)
|
||||||
|
with MCPServerAdapter(serverparams, connect_timeout=60) as tools:
|
||||||
|
assert isinstance(tools, ToolCollection)
|
||||||
|
assert len(tools) == 2
|
||||||
|
assert tools[0].name == "echo_tool"
|
||||||
|
assert tools[1].name == "calc_tool"
|
||||||
|
assert tools[0].run(text="hello") == "Echo: hello"
|
||||||
|
|
||||||
|
def test_connect_timeout_with_filtered_tools(echo_server_script):
|
||||||
|
serverparams = StdioServerParameters(
|
||||||
|
command="uv", args=["run", "python", "-c", echo_server_script]
|
||||||
|
)
|
||||||
|
with MCPServerAdapter(serverparams, "echo_tool", connect_timeout=45) as tools:
|
||||||
|
assert isinstance(tools, ToolCollection)
|
||||||
|
assert len(tools) == 1
|
||||||
|
assert tools[0].name == "echo_tool"
|
||||||
|
assert tools[0].run(text="timeout test") == "Echo: timeout test"
|
||||||
|
|
||||||
|
@patch('crewai_tools.adapters.mcp_adapter.MCPAdapt')
|
||||||
|
def test_connect_timeout_passed_to_mcpadapt(mock_mcpadapt):
|
||||||
|
mock_adapter_instance = MagicMock()
|
||||||
|
mock_mcpadapt.return_value = mock_adapter_instance
|
||||||
|
|
||||||
|
serverparams = StdioServerParameters(
|
||||||
|
command="uv", args=["run", "echo", "test"]
|
||||||
|
)
|
||||||
|
|
||||||
|
MCPServerAdapter(serverparams)
|
||||||
|
mock_mcpadapt.assert_called_once()
|
||||||
|
assert mock_mcpadapt.call_args[0][2] == 30
|
||||||
|
|
||||||
|
mock_mcpadapt.reset_mock()
|
||||||
|
|
||||||
|
MCPServerAdapter(serverparams, connect_timeout=5)
|
||||||
|
mock_mcpadapt.assert_called_once()
|
||||||
|
assert mock_mcpadapt.call_args[0][2] == 5
|
||||||
|
|||||||
Reference in New Issue
Block a user