You can not select more than 25 topics Topics must start with a chinese character,a letter or number, can include dashes ('-') and can be up to 35 characters long.

test_mcp_tools.py 35 kB

Make shared session possible for MCP tool (#6312) Resolves #6232, #6198 This PR introduces an optional parameter `session` to `mcp_server_tools` to support reuse of the same session. ```python import asyncio from autogen_agentchat.agents import AssistantAgent from autogen_agentchat.conditions import TextMentionTermination from autogen_agentchat.teams import RoundRobinGroupChat from autogen_agentchat.ui import Console from autogen_ext.models.openai import OpenAIChatCompletionClient from autogen_ext.tools.mcp import StdioServerParams, create_mcp_server_session, mcp_server_tools async def main() -> None: model_client = OpenAIChatCompletionClient(model="gpt-4o", parallel_tool_calls=False) # type: ignore params = StdioServerParams( command="npx", args=["@playwright/mcp@latest"], read_timeout_seconds=60, ) async with create_mcp_server_session(params) as session: await session.initialize() tools = await mcp_server_tools(server_params=params, session=session) print(f"Tools: {[tool.name for tool in tools]}") agent = AssistantAgent( name="Assistant", model_client=model_client, tools=tools, # type: ignore ) termination = TextMentionTermination("TERMINATE") team = RoundRobinGroupChat([agent], termination_condition=termination) await Console( team.run_stream( task="Go to https://ekzhu.com/, visit the first link in the page, then tell me about the linked page." ) ) asyncio.run(main()) ``` Based on discussion in this thread: #6284, we will consider serialization and deserialization of MCP server tools when used in this manner in a separate issue. This PR also replaces the `json_schema_to_pydantic` dependency with built-in utils.
1 year ago
Make shared session possible for MCP tool (#6312) Resolves #6232, #6198 This PR introduces an optional parameter `session` to `mcp_server_tools` to support reuse of the same session. ```python import asyncio from autogen_agentchat.agents import AssistantAgent from autogen_agentchat.conditions import TextMentionTermination from autogen_agentchat.teams import RoundRobinGroupChat from autogen_agentchat.ui import Console from autogen_ext.models.openai import OpenAIChatCompletionClient from autogen_ext.tools.mcp import StdioServerParams, create_mcp_server_session, mcp_server_tools async def main() -> None: model_client = OpenAIChatCompletionClient(model="gpt-4o", parallel_tool_calls=False) # type: ignore params = StdioServerParams( command="npx", args=["@playwright/mcp@latest"], read_timeout_seconds=60, ) async with create_mcp_server_session(params) as session: await session.initialize() tools = await mcp_server_tools(server_params=params, session=session) print(f"Tools: {[tool.name for tool in tools]}") agent = AssistantAgent( name="Assistant", model_client=model_client, tools=tools, # type: ignore ) termination = TextMentionTermination("TERMINATE") team = RoundRobinGroupChat([agent], termination_condition=termination) await Console( team.run_stream( task="Go to https://ekzhu.com/, visit the first link in the page, then tell me about the linked page." ) ) asyncio.run(main()) ``` Based on discussion in this thread: #6284, we will consider serialization and deserialization of MCP server tools when used in this manner in a separate issue. This PR also replaces the `json_schema_to_pydantic` dependency with built-in utils.
1 year ago
Make shared session possible for MCP tool (#6312) Resolves #6232, #6198 This PR introduces an optional parameter `session` to `mcp_server_tools` to support reuse of the same session. ```python import asyncio from autogen_agentchat.agents import AssistantAgent from autogen_agentchat.conditions import TextMentionTermination from autogen_agentchat.teams import RoundRobinGroupChat from autogen_agentchat.ui import Console from autogen_ext.models.openai import OpenAIChatCompletionClient from autogen_ext.tools.mcp import StdioServerParams, create_mcp_server_session, mcp_server_tools async def main() -> None: model_client = OpenAIChatCompletionClient(model="gpt-4o", parallel_tool_calls=False) # type: ignore params = StdioServerParams( command="npx", args=["@playwright/mcp@latest"], read_timeout_seconds=60, ) async with create_mcp_server_session(params) as session: await session.initialize() tools = await mcp_server_tools(server_params=params, session=session) print(f"Tools: {[tool.name for tool in tools]}") agent = AssistantAgent( name="Assistant", model_client=model_client, tools=tools, # type: ignore ) termination = TextMentionTermination("TERMINATE") team = RoundRobinGroupChat([agent], termination_condition=termination) await Console( team.run_stream( task="Go to https://ekzhu.com/, visit the first link in the page, then tell me about the linked page." ) ) asyncio.run(main()) ``` Based on discussion in this thread: #6284, we will consider serialization and deserialization of MCP server tools when used in this manner in a separate issue. This PR also replaces the `json_schema_to_pydantic` dependency with built-in utils.
1 year ago
Make shared session possible for MCP tool (#6312) Resolves #6232, #6198 This PR introduces an optional parameter `session` to `mcp_server_tools` to support reuse of the same session. ```python import asyncio from autogen_agentchat.agents import AssistantAgent from autogen_agentchat.conditions import TextMentionTermination from autogen_agentchat.teams import RoundRobinGroupChat from autogen_agentchat.ui import Console from autogen_ext.models.openai import OpenAIChatCompletionClient from autogen_ext.tools.mcp import StdioServerParams, create_mcp_server_session, mcp_server_tools async def main() -> None: model_client = OpenAIChatCompletionClient(model="gpt-4o", parallel_tool_calls=False) # type: ignore params = StdioServerParams( command="npx", args=["@playwright/mcp@latest"], read_timeout_seconds=60, ) async with create_mcp_server_session(params) as session: await session.initialize() tools = await mcp_server_tools(server_params=params, session=session) print(f"Tools: {[tool.name for tool in tools]}") agent = AssistantAgent( name="Assistant", model_client=model_client, tools=tools, # type: ignore ) termination = TextMentionTermination("TERMINATE") team = RoundRobinGroupChat([agent], termination_condition=termination) await Console( team.run_stream( task="Go to https://ekzhu.com/, visit the first link in the page, then tell me about the linked page." ) ) asyncio.run(main()) ``` Based on discussion in this thread: #6284, we will consider serialization and deserialization of MCP server tools when used in this manner in a separate issue. This PR also replaces the `json_schema_to_pydantic` dependency with built-in utils.
1 year ago
Make shared session possible for MCP tool (#6312) Resolves #6232, #6198 This PR introduces an optional parameter `session` to `mcp_server_tools` to support reuse of the same session. ```python import asyncio from autogen_agentchat.agents import AssistantAgent from autogen_agentchat.conditions import TextMentionTermination from autogen_agentchat.teams import RoundRobinGroupChat from autogen_agentchat.ui import Console from autogen_ext.models.openai import OpenAIChatCompletionClient from autogen_ext.tools.mcp import StdioServerParams, create_mcp_server_session, mcp_server_tools async def main() -> None: model_client = OpenAIChatCompletionClient(model="gpt-4o", parallel_tool_calls=False) # type: ignore params = StdioServerParams( command="npx", args=["@playwright/mcp@latest"], read_timeout_seconds=60, ) async with create_mcp_server_session(params) as session: await session.initialize() tools = await mcp_server_tools(server_params=params, session=session) print(f"Tools: {[tool.name for tool in tools]}") agent = AssistantAgent( name="Assistant", model_client=model_client, tools=tools, # type: ignore ) termination = TextMentionTermination("TERMINATE") team = RoundRobinGroupChat([agent], termination_condition=termination) await Console( team.run_stream( task="Go to https://ekzhu.com/, visit the first link in the page, then tell me about the linked page." ) ) asyncio.run(main()) ``` Based on discussion in this thread: #6284, we will consider serialization and deserialization of MCP server tools when used in this manner in a separate issue. This PR also replaces the `json_schema_to_pydantic` dependency with built-in utils.
1 year ago
Make shared session possible for MCP tool (#6312) Resolves #6232, #6198 This PR introduces an optional parameter `session` to `mcp_server_tools` to support reuse of the same session. ```python import asyncio from autogen_agentchat.agents import AssistantAgent from autogen_agentchat.conditions import TextMentionTermination from autogen_agentchat.teams import RoundRobinGroupChat from autogen_agentchat.ui import Console from autogen_ext.models.openai import OpenAIChatCompletionClient from autogen_ext.tools.mcp import StdioServerParams, create_mcp_server_session, mcp_server_tools async def main() -> None: model_client = OpenAIChatCompletionClient(model="gpt-4o", parallel_tool_calls=False) # type: ignore params = StdioServerParams( command="npx", args=["@playwright/mcp@latest"], read_timeout_seconds=60, ) async with create_mcp_server_session(params) as session: await session.initialize() tools = await mcp_server_tools(server_params=params, session=session) print(f"Tools: {[tool.name for tool in tools]}") agent = AssistantAgent( name="Assistant", model_client=model_client, tools=tools, # type: ignore ) termination = TextMentionTermination("TERMINATE") team = RoundRobinGroupChat([agent], termination_condition=termination) await Console( team.run_stream( task="Go to https://ekzhu.com/, visit the first link in the page, then tell me about the linked page." ) ) asyncio.run(main()) ``` Based on discussion in this thread: #6284, we will consider serialization and deserialization of MCP server tools when used in this manner in a separate issue. This PR also replaces the `json_schema_to_pydantic` dependency with built-in utils.
1 year ago
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930
  1. import asyncio
  2. import logging
  3. import os
  4. import threading
  5. from typing import cast
  6. from unittest.mock import AsyncMock, MagicMock
  7. import pytest
  8. from _pytest.logging import LogCaptureFixture # type: ignore[import]
  9. from autogen_core import CancellationToken
  10. from autogen_core.tools import Workbench
  11. from autogen_core.utils import schema_to_pydantic_model
  12. from autogen_ext.tools.mcp import (
  13. McpSessionActor,
  14. McpWorkbench,
  15. SseMcpToolAdapter,
  16. SseServerParams,
  17. StdioMcpToolAdapter,
  18. StdioServerParams,
  19. StreamableHttpMcpToolAdapter,
  20. StreamableHttpServerParams,
  21. create_mcp_server_session,
  22. mcp_server_tools,
  23. )
  24. from mcp import ClientSession, Tool
  25. from mcp.types import (
  26. Annotations,
  27. EmbeddedResource,
  28. ImageContent,
  29. ResourceLink,
  30. TextContent,
  31. TextResourceContents,
  32. )
  33. from pydantic.networks import AnyUrl
  34. @pytest.fixture
  35. def sample_tool() -> Tool:
  36. return Tool(
  37. name="test_tool",
  38. description="A test tool",
  39. inputSchema={
  40. "type": "object",
  41. "properties": {"test_param": {"type": "string"}},
  42. "required": ["test_param"],
  43. },
  44. )
  45. @pytest.fixture
  46. def sample_server_params() -> StdioServerParams:
  47. return StdioServerParams(command="echo", args=["test"])
  48. @pytest.fixture
  49. def sample_sse_tool() -> Tool:
  50. return Tool(
  51. name="test_sse_tool",
  52. description="A test SSE tool",
  53. inputSchema={
  54. "type": "object",
  55. "properties": {"test_param": {"type": "string"}},
  56. "required": ["test_param"],
  57. },
  58. )
  59. @pytest.fixture
  60. def sample_streamable_http_tool() -> Tool:
  61. return Tool(
  62. name="test_streamable_http_tool",
  63. description="A test StreamableHttp tool",
  64. inputSchema={
  65. "type": "object",
  66. "properties": {"test_param": {"type": "string"}},
  67. "required": ["test_param"],
  68. },
  69. )
  70. @pytest.fixture
  71. def mock_sse_session() -> AsyncMock:
  72. session = AsyncMock(spec=ClientSession)
  73. session.initialize = AsyncMock()
  74. session.call_tool = AsyncMock()
  75. session.list_tools = AsyncMock()
  76. return session
  77. @pytest.fixture
  78. def mock_streamable_http_session() -> AsyncMock:
  79. session = AsyncMock(spec=ClientSession)
  80. session.initialize = AsyncMock()
  81. session.call_tool = AsyncMock()
  82. session.list_tools = AsyncMock()
  83. return session
  84. @pytest.fixture
  85. def mock_session() -> AsyncMock:
  86. session = AsyncMock(spec=ClientSession)
  87. session.initialize = AsyncMock()
  88. session.call_tool = AsyncMock()
  89. session.list_tools = AsyncMock()
  90. return session
  91. @pytest.fixture
  92. def mock_tool_response() -> MagicMock:
  93. response = MagicMock()
  94. response.isError = False
  95. response.content = [
  96. TextContent(
  97. text="test_output",
  98. type="text",
  99. annotations=Annotations(audience=["user", "assistant"], priority=0.7),
  100. ),
  101. ]
  102. return response
  103. @pytest.fixture
  104. def cancellation_token() -> CancellationToken:
  105. return CancellationToken()
  106. @pytest.fixture
  107. def mock_error_tool_response() -> MagicMock:
  108. response = MagicMock()
  109. response.isError = True
  110. response.content = [TextContent(text="error output", type="text")]
  111. return response
  112. def test_adapter_config_serialization(sample_tool: Tool, sample_server_params: StdioServerParams) -> None:
  113. """Test that adapter can be saved to and loaded from config."""
  114. original_adapter = StdioMcpToolAdapter(server_params=sample_server_params, tool=sample_tool)
  115. config = original_adapter.dump_component()
  116. loaded_adapter = StdioMcpToolAdapter.load_component(config)
  117. # Test that the loaded adapter has the same properties
  118. assert loaded_adapter.name == "test_tool"
  119. assert loaded_adapter.description == "A test tool"
  120. # Verify schema structure
  121. schema = loaded_adapter.schema
  122. assert "parameters" in schema, "Schema must have parameters"
  123. params_schema = schema["parameters"]
  124. assert isinstance(params_schema, dict), "Parameters must be a dict"
  125. assert "type" in params_schema, "Parameters must have type"
  126. assert "required" in params_schema, "Parameters must have required fields"
  127. assert "properties" in params_schema, "Parameters must have properties"
  128. # Compare schema content
  129. assert params_schema["type"] == sample_tool.inputSchema["type"]
  130. assert params_schema["required"] == sample_tool.inputSchema["required"]
  131. assert (
  132. params_schema["properties"]["test_param"]["type"] == sample_tool.inputSchema["properties"]["test_param"]["type"]
  133. )
  134. @pytest.mark.asyncio
  135. async def test_mcp_tool_execution(
  136. sample_tool: Tool,
  137. sample_server_params: StdioServerParams,
  138. mock_session: AsyncMock,
  139. mock_tool_response: MagicMock,
  140. cancellation_token: CancellationToken,
  141. monkeypatch: pytest.MonkeyPatch,
  142. caplog: pytest.LogCaptureFixture,
  143. ) -> None:
  144. """Test that adapter properly executes tools through ClientSession."""
  145. mock_context = AsyncMock()
  146. mock_context.__aenter__.return_value = mock_session
  147. monkeypatch.setattr(
  148. "autogen_ext.tools.mcp._base.create_mcp_server_session",
  149. lambda *args, **kwargs: mock_context, # type: ignore
  150. )
  151. mock_session.call_tool.return_value = mock_tool_response
  152. with caplog.at_level(logging.INFO):
  153. adapter = StdioMcpToolAdapter(server_params=sample_server_params, tool=sample_tool)
  154. result = await adapter.run_json(
  155. args=schema_to_pydantic_model(sample_tool.inputSchema)(**{"test_param": "test"}).model_dump(),
  156. cancellation_token=cancellation_token,
  157. )
  158. assert result == mock_tool_response.content
  159. mock_session.initialize.assert_called_once()
  160. mock_session.call_tool.assert_called_once()
  161. # Check log.
  162. assert "test_output" in caplog.text
  163. @pytest.mark.asyncio
  164. async def test_adapter_from_server_params(
  165. sample_tool: Tool,
  166. sample_server_params: StdioServerParams,
  167. mock_session: AsyncMock,
  168. monkeypatch: pytest.MonkeyPatch,
  169. ) -> None:
  170. """Test that adapter can be created from server parameters."""
  171. mock_context = AsyncMock()
  172. mock_context.__aenter__.return_value = mock_session
  173. monkeypatch.setattr(
  174. "autogen_ext.tools.mcp._base.create_mcp_server_session",
  175. lambda *args, **kwargs: mock_context, # type: ignore
  176. )
  177. mock_session.list_tools.return_value.tools = [sample_tool]
  178. adapter = await StdioMcpToolAdapter.from_server_params(sample_server_params, "test_tool")
  179. assert isinstance(adapter, StdioMcpToolAdapter)
  180. assert adapter.name == "test_tool"
  181. assert adapter.description == "A test tool"
  182. # Verify schema structure
  183. schema = adapter.schema
  184. assert "parameters" in schema, "Schema must have parameters"
  185. params_schema = schema["parameters"]
  186. assert isinstance(params_schema, dict), "Parameters must be a dict"
  187. assert "type" in params_schema, "Parameters must have type"
  188. assert "required" in params_schema, "Parameters must have required fields"
  189. assert "properties" in params_schema, "Parameters must have properties"
  190. # Compare schema content
  191. assert params_schema["type"] == sample_tool.inputSchema["type"]
  192. assert params_schema["required"] == sample_tool.inputSchema["required"]
  193. assert (
  194. params_schema["properties"]["test_param"]["type"] == sample_tool.inputSchema["properties"]["test_param"]["type"]
  195. )
  196. @pytest.mark.asyncio
  197. async def test_adapter_from_server_params_with_return_value_as_string(
  198. sample_tool: Tool,
  199. sample_server_params: StdioServerParams,
  200. mock_session: AsyncMock,
  201. monkeypatch: pytest.MonkeyPatch,
  202. ) -> None:
  203. """Test that adapter can be created from server parameters."""
  204. mock_context = AsyncMock()
  205. mock_context.__aenter__.return_value = mock_session
  206. monkeypatch.setattr(
  207. "autogen_ext.tools.mcp._base.create_mcp_server_session",
  208. lambda *args, **kwargs: mock_context, # type: ignore
  209. )
  210. mock_session.list_tools.return_value.tools = [sample_tool]
  211. adapter = await StdioMcpToolAdapter.from_server_params(sample_server_params, "test_tool")
  212. assert (
  213. adapter.return_value_as_string(
  214. [
  215. TextContent(
  216. text="this is a sample text",
  217. type="text",
  218. annotations=Annotations(audience=["user", "assistant"], priority=0.7),
  219. ),
  220. ImageContent(
  221. data="this is a sample base64 encoded image",
  222. mimeType="image/png",
  223. type="image",
  224. annotations=None,
  225. ),
  226. EmbeddedResource(
  227. type="resource",
  228. resource=TextResourceContents(
  229. text="this is a sample text",
  230. uri=AnyUrl(url="http://example.com/test"),
  231. ),
  232. annotations=Annotations(audience=["user"], priority=0.3),
  233. ),
  234. ]
  235. )
  236. == '[{"type": "text", "text": "this is a sample text", "annotations": {"audience": ["user", "assistant"], "priority": 0.7}}, {"type": "image", "data": "this is a sample base64 encoded image", "mimeType": "image/png", "annotations": null}, {"type": "resource", "resource": {"uri": "http://example.com/test", "mimeType": null, "text": "this is a sample text"}, "annotations": {"audience": ["user"], "priority": 0.3}}]'
  237. )
  238. @pytest.mark.asyncio
  239. async def test_adapter_from_factory(
  240. sample_tool: Tool,
  241. sample_server_params: StdioServerParams,
  242. mock_session: AsyncMock,
  243. monkeypatch: pytest.MonkeyPatch,
  244. ) -> None:
  245. """Test that factory function returns a list of tools."""
  246. mock_context = AsyncMock()
  247. mock_context.__aenter__.return_value = mock_session
  248. monkeypatch.setattr(
  249. "autogen_ext.tools.mcp._factory.create_mcp_server_session",
  250. lambda *args, **kwargs: mock_context, # type: ignore
  251. )
  252. mock_session.list_tools.return_value.tools = [sample_tool]
  253. tools = await mcp_server_tools(server_params=sample_server_params)
  254. assert tools is not None
  255. assert len(tools) > 0
  256. assert isinstance(tools[0], StdioMcpToolAdapter)
  257. @pytest.mark.asyncio
  258. async def test_adapter_from_factory_existing_session(
  259. sample_tool: Tool,
  260. sample_server_params: StdioServerParams,
  261. mock_session: AsyncMock,
  262. monkeypatch: pytest.MonkeyPatch,
  263. ) -> None:
  264. """Test that factory function returns a list of tools with an existing session."""
  265. mock_context = AsyncMock()
  266. mock_context.__aenter__.return_value = mock_session
  267. monkeypatch.setattr(
  268. "autogen_ext.tools.mcp._factory.create_mcp_server_session",
  269. lambda *args, **kwargs: mock_context, # type: ignore
  270. )
  271. mock_session.list_tools.return_value.tools = [sample_tool]
  272. tools = await mcp_server_tools(server_params=sample_server_params, session=mock_session)
  273. assert tools is not None
  274. assert len(tools) > 0
  275. assert isinstance(tools[0], StdioMcpToolAdapter)
  276. @pytest.mark.asyncio
  277. async def test_sse_adapter_config_serialization(sample_sse_tool: Tool) -> None:
  278. """Test that SSE adapter can be saved to and loaded from config."""
  279. params = SseServerParams(url="http://test-url")
  280. original_adapter = SseMcpToolAdapter(server_params=params, tool=sample_sse_tool)
  281. config = original_adapter.dump_component()
  282. loaded_adapter = SseMcpToolAdapter.load_component(config)
  283. # Test that the loaded adapter has the same properties
  284. assert loaded_adapter.name == "test_sse_tool"
  285. assert loaded_adapter.description == "A test SSE tool"
  286. # Verify schema structure
  287. schema = loaded_adapter.schema
  288. assert "parameters" in schema, "Schema must have parameters"
  289. params_schema = schema["parameters"]
  290. assert isinstance(params_schema, dict), "Parameters must be a dict"
  291. assert "type" in params_schema, "Parameters must have type"
  292. assert "required" in params_schema, "Parameters must have required fields"
  293. assert "properties" in params_schema, "Parameters must have properties"
  294. # Compare schema content
  295. assert params_schema["type"] == sample_sse_tool.inputSchema["type"]
  296. assert params_schema["required"] == sample_sse_tool.inputSchema["required"]
  297. assert (
  298. params_schema["properties"]["test_param"]["type"]
  299. == sample_sse_tool.inputSchema["properties"]["test_param"]["type"]
  300. )
  301. @pytest.mark.asyncio
  302. async def test_sse_tool_execution(
  303. sample_sse_tool: Tool,
  304. mock_sse_session: AsyncMock,
  305. monkeypatch: pytest.MonkeyPatch,
  306. caplog: pytest.LogCaptureFixture,
  307. ) -> None:
  308. """Test that SSE adapter properly executes tools through ClientSession."""
  309. params = SseServerParams(url="http://test-url")
  310. mock_context = AsyncMock()
  311. mock_context.__aenter__.return_value = mock_sse_session
  312. mock_sse_session.call_tool.return_value = MagicMock(
  313. isError=False,
  314. content=[
  315. TextContent(
  316. text="test_output",
  317. type="text",
  318. annotations=Annotations(audience=["user", "assistant"], priority=0.7),
  319. ),
  320. ],
  321. )
  322. monkeypatch.setattr(
  323. "autogen_ext.tools.mcp._base.create_mcp_server_session",
  324. lambda *args, **kwargs: mock_context, # type: ignore
  325. )
  326. with caplog.at_level(logging.INFO):
  327. adapter = SseMcpToolAdapter(server_params=params, tool=sample_sse_tool)
  328. result = await adapter.run_json(
  329. args=schema_to_pydantic_model(sample_sse_tool.inputSchema)(**{"test_param": "test"}).model_dump(),
  330. cancellation_token=CancellationToken(),
  331. )
  332. assert result == mock_sse_session.call_tool.return_value.content
  333. mock_sse_session.initialize.assert_called_once()
  334. mock_sse_session.call_tool.assert_called_once()
  335. # Check log.
  336. assert "test_output" in caplog.text
  337. @pytest.mark.asyncio
  338. async def test_sse_adapter_from_server_params(
  339. sample_sse_tool: Tool,
  340. mock_sse_session: AsyncMock,
  341. monkeypatch: pytest.MonkeyPatch,
  342. ) -> None:
  343. """Test that SSE adapter can be created from server parameters."""
  344. params = SseServerParams(url="http://test-url")
  345. mock_context = AsyncMock()
  346. mock_context.__aenter__.return_value = mock_sse_session
  347. monkeypatch.setattr(
  348. "autogen_ext.tools.mcp._base.create_mcp_server_session",
  349. lambda *args, **kwargs: mock_context, # type: ignore
  350. )
  351. mock_sse_session.list_tools.return_value.tools = [sample_sse_tool]
  352. adapter = await SseMcpToolAdapter.from_server_params(params, "test_sse_tool")
  353. assert isinstance(adapter, SseMcpToolAdapter)
  354. assert adapter.name == "test_sse_tool"
  355. assert adapter.description == "A test SSE tool"
  356. # Verify schema structure
  357. schema = adapter.schema
  358. assert "parameters" in schema, "Schema must have parameters"
  359. params_schema = schema["parameters"]
  360. assert isinstance(params_schema, dict), "Parameters must be a dict"
  361. assert "type" in params_schema, "Parameters must have type"
  362. assert "required" in params_schema, "Parameters must have required fields"
  363. assert "properties" in params_schema, "Parameters must have properties"
  364. # Compare schema content
  365. assert params_schema["type"] == sample_sse_tool.inputSchema["type"]
  366. assert params_schema["required"] == sample_sse_tool.inputSchema["required"]
  367. assert (
  368. params_schema["properties"]["test_param"]["type"]
  369. == sample_sse_tool.inputSchema["properties"]["test_param"]["type"]
  370. )
  371. @pytest.mark.asyncio
  372. async def test_streamable_http_adapter_config_serialization(sample_streamable_http_tool: Tool) -> None:
  373. """Test that StreamableHttp adapter can be saved to and loaded from config."""
  374. params = StreamableHttpServerParams(url="http://test-url")
  375. original_adapter = StreamableHttpMcpToolAdapter(server_params=params, tool=sample_streamable_http_tool)
  376. config = original_adapter.dump_component()
  377. loaded_adapter = StreamableHttpMcpToolAdapter.load_component(config)
  378. # Test that the loaded adapter has the same properties
  379. assert loaded_adapter.name == "test_streamable_http_tool"
  380. assert loaded_adapter.description == "A test StreamableHttp tool"
  381. # Verify schema structure
  382. schema = loaded_adapter.schema
  383. assert "parameters" in schema, "Schema must have parameters"
  384. params_schema = schema["parameters"]
  385. assert isinstance(params_schema, dict), "Parameters must be a dict"
  386. assert "type" in params_schema, "Parameters must have type"
  387. assert "required" in params_schema, "Parameters must have required fields"
  388. assert "properties" in params_schema, "Parameters must have properties"
  389. # Compare schema content
  390. assert params_schema["type"] == sample_streamable_http_tool.inputSchema["type"]
  391. assert params_schema["required"] == sample_streamable_http_tool.inputSchema["required"]
  392. assert (
  393. params_schema["properties"]["test_param"]["type"]
  394. == sample_streamable_http_tool.inputSchema["properties"]["test_param"]["type"]
  395. )
  396. @pytest.mark.asyncio
  397. async def test_streamable_http_tool_execution(
  398. sample_streamable_http_tool: Tool,
  399. mock_streamable_http_session: AsyncMock,
  400. monkeypatch: pytest.MonkeyPatch,
  401. caplog: pytest.LogCaptureFixture,
  402. ) -> None:
  403. """Test that StreamableHttp adapter properly executes tools through ClientSession."""
  404. params = StreamableHttpServerParams(url="http://test-url")
  405. mock_context = AsyncMock()
  406. mock_context.__aenter__.return_value = mock_streamable_http_session
  407. mock_streamable_http_session.call_tool.return_value = MagicMock(
  408. isError=False,
  409. content=[
  410. TextContent(
  411. text="test_output",
  412. type="text",
  413. annotations=Annotations(audience=["user", "assistant"], priority=0.7),
  414. ),
  415. ],
  416. )
  417. monkeypatch.setattr(
  418. "autogen_ext.tools.mcp._base.create_mcp_server_session",
  419. lambda *args, **kwargs: mock_context, # type: ignore
  420. )
  421. with caplog.at_level(logging.INFO):
  422. adapter = StreamableHttpMcpToolAdapter(server_params=params, tool=sample_streamable_http_tool)
  423. result = await adapter.run_json(
  424. args=schema_to_pydantic_model(sample_streamable_http_tool.inputSchema)(
  425. **{"test_param": "test"}
  426. ).model_dump(),
  427. cancellation_token=CancellationToken(),
  428. )
  429. assert result == mock_streamable_http_session.call_tool.return_value.content
  430. mock_streamable_http_session.initialize.assert_called_once()
  431. mock_streamable_http_session.call_tool.assert_called_once()
  432. # Check log.
  433. assert "test_output" in caplog.text
  434. @pytest.mark.asyncio
  435. async def test_streamable_http_adapter_from_server_params(
  436. sample_streamable_http_tool: Tool,
  437. mock_streamable_http_session: AsyncMock,
  438. monkeypatch: pytest.MonkeyPatch,
  439. ) -> None:
  440. """Test that StreamableHttp adapter can be created from server parameters."""
  441. params = StreamableHttpServerParams(url="http://test-url")
  442. mock_context = AsyncMock()
  443. mock_context.__aenter__.return_value = mock_streamable_http_session
  444. monkeypatch.setattr(
  445. "autogen_ext.tools.mcp._base.create_mcp_server_session",
  446. lambda *args, **kwargs: mock_context, # type: ignore
  447. )
  448. mock_streamable_http_session.list_tools.return_value.tools = [sample_streamable_http_tool]
  449. adapter = await StreamableHttpMcpToolAdapter.from_server_params(params, "test_streamable_http_tool")
  450. assert isinstance(adapter, StreamableHttpMcpToolAdapter)
  451. assert adapter.name == "test_streamable_http_tool"
  452. assert adapter.description == "A test StreamableHttp tool"
  453. # Verify schema structure
  454. schema = adapter.schema
  455. assert "parameters" in schema, "Schema must have parameters"
  456. params_schema = schema["parameters"]
  457. assert isinstance(params_schema, dict), "Parameters must be a dict"
  458. assert "type" in params_schema, "Parameters must have type"
  459. assert "required" in params_schema, "Parameters must have required fields"
  460. assert "properties" in params_schema, "Parameters must have properties"
  461. # Compare schema content
  462. assert params_schema["type"] == sample_streamable_http_tool.inputSchema["type"]
  463. assert params_schema["required"] == sample_streamable_http_tool.inputSchema["required"]
  464. assert (
  465. params_schema["properties"]["test_param"]["type"]
  466. == sample_streamable_http_tool.inputSchema["properties"]["test_param"]["type"]
  467. )
  468. @pytest.mark.asyncio
  469. async def test_mcp_server_fetch() -> None:
  470. params = StdioServerParams(
  471. command="uvx",
  472. args=["mcp-server-fetch"],
  473. read_timeout_seconds=60,
  474. )
  475. tools = await mcp_server_tools(server_params=params)
  476. assert tools is not None
  477. assert tools[0].name == "fetch"
  478. result = await tools[0].run_json({"url": "https://github.com/"}, CancellationToken())
  479. assert result is not None
  480. @pytest.mark.asyncio
  481. async def test_mcp_server_filesystem() -> None:
  482. params = StdioServerParams(
  483. command="npx",
  484. args=[
  485. "-y",
  486. "@modelcontextprotocol/server-filesystem",
  487. ".",
  488. ],
  489. read_timeout_seconds=60,
  490. )
  491. tools = await mcp_server_tools(server_params=params)
  492. assert tools is not None
  493. tools = [tool for tool in tools if tool.name == "read_file"]
  494. assert len(tools) == 1
  495. tool = tools[0]
  496. result = await tool.run_json({"path": "README.md"}, CancellationToken())
  497. assert result is not None
  498. @pytest.mark.asyncio
  499. async def test_mcp_server_git() -> None:
  500. params = StdioServerParams(
  501. command="uvx",
  502. args=["mcp-server-git"],
  503. read_timeout_seconds=60,
  504. )
  505. tools = await mcp_server_tools(server_params=params)
  506. assert tools is not None
  507. tools = [tool for tool in tools if tool.name == "git_log"]
  508. assert len(tools) == 1
  509. tool = tools[0]
  510. repo_path = os.path.join(os.path.dirname(__file__), "..", "..", "..", "..", "..")
  511. result = await tool.run_json({"repo_path": repo_path}, CancellationToken())
  512. assert result is not None
  513. @pytest.mark.asyncio
  514. async def test_mcp_server_git_existing_session() -> None:
  515. params = StdioServerParams(
  516. command="uvx",
  517. args=["mcp-server-git"],
  518. read_timeout_seconds=60,
  519. )
  520. async with create_mcp_server_session(params) as session:
  521. await session.initialize()
  522. tools = await mcp_server_tools(server_params=params, session=session)
  523. assert tools is not None
  524. git_log = [tool for tool in tools if tool.name == "git_log"][0]
  525. repo_path = os.path.join(os.path.dirname(__file__), "..", "..", "..", "..", "..")
  526. result = await git_log.run_json({"repo_path": repo_path}, CancellationToken())
  527. assert result is not None
  528. git_status = [tool for tool in tools if tool.name == "git_status"][0]
  529. result = await git_status.run_json({"repo_path": repo_path}, CancellationToken())
  530. assert result is not None
  531. @pytest.mark.asyncio
  532. async def test_mcp_server_github() -> None:
  533. # Check if GITHUB_TOKEN is set.
  534. if "GITHUB_TOKEN" not in os.environ:
  535. pytest.skip("GITHUB_TOKEN environment variable is not set. Skipping test.")
  536. params = StdioServerParams(
  537. command="npx",
  538. args=[
  539. "-y",
  540. "@modelcontextprotocol/server-github",
  541. ],
  542. env={"GITHUB_PERSONAL_ACCESS_TOKEN": os.environ["GITHUB_TOKEN"]},
  543. read_timeout_seconds=60,
  544. )
  545. tools = await mcp_server_tools(server_params=params)
  546. assert tools is not None
  547. tools = [tool for tool in tools if tool.name == "get_file_contents"]
  548. assert len(tools) == 1
  549. tool = tools[0]
  550. result = await tool.run_json(
  551. {"owner": "microsoft", "repo": "autogen", "path": "python", "branch": "main"},
  552. CancellationToken(),
  553. )
  554. assert result is not None
  555. @pytest.mark.asyncio
  556. async def test_mcp_workbench_start_stop() -> None:
  557. params = StdioServerParams(
  558. command="uvx",
  559. args=["mcp-server-fetch"],
  560. read_timeout_seconds=60,
  561. )
  562. workbench = McpWorkbench(params)
  563. assert workbench is not None
  564. assert workbench.server_params == params
  565. await workbench.start()
  566. assert workbench._actor is not None # type: ignore[reportPrivateUsage]
  567. await workbench.stop()
  568. assert workbench._actor is None # type: ignore[reportPrivateUsage]
  569. @pytest.mark.asyncio
  570. async def test_mcp_workbench_server_fetch() -> None:
  571. params = StdioServerParams(
  572. command="uvx",
  573. args=["mcp-server-fetch"],
  574. read_timeout_seconds=60,
  575. )
  576. workbench = McpWorkbench(server_params=params)
  577. await workbench.start()
  578. tools = await workbench.list_tools()
  579. assert tools is not None
  580. assert tools[0]["name"] == "fetch"
  581. result = await workbench.call_tool(tools[0]["name"], {"url": "https://github.com/"}, CancellationToken())
  582. assert result is not None
  583. await workbench.stop()
  584. @pytest.mark.asyncio
  585. async def test_mcp_workbench_server_filesystem() -> None:
  586. params = StdioServerParams(
  587. command="npx",
  588. args=[
  589. "-y",
  590. "@modelcontextprotocol/server-filesystem",
  591. ".",
  592. ],
  593. read_timeout_seconds=60,
  594. )
  595. workbench = McpWorkbench(server_params=params)
  596. await workbench.start()
  597. tools = await workbench.list_tools()
  598. assert tools is not None
  599. tools = [tool for tool in tools if tool["name"] == "read_file"]
  600. assert len(tools) == 1
  601. tool = tools[0]
  602. result = await workbench.call_tool(tool["name"], {"path": "README.md"}, CancellationToken())
  603. assert result is not None
  604. await workbench.stop()
  605. # Serialize the workbench.
  606. config = workbench.dump_component()
  607. # Deserialize the workbench.
  608. async with Workbench.load_component(config) as new_workbench:
  609. tools = await new_workbench.list_tools()
  610. assert tools is not None
  611. tools = [tool for tool in tools if tool["name"] == "read_file"]
  612. assert len(tools) == 1
  613. tool = tools[0]
  614. result = await new_workbench.call_tool(tool["name"], {"path": "README.md"}, CancellationToken())
  615. assert result is not None
  616. @pytest.mark.asyncio
  617. async def test_lazy_init_and_finalize_cleanup() -> None:
  618. params = StdioServerParams(
  619. command="npx",
  620. args=[
  621. "-y",
  622. "@modelcontextprotocol/server-filesystem",
  623. ".",
  624. ],
  625. read_timeout_seconds=60,
  626. )
  627. workbench = McpWorkbench(server_params=params)
  628. # Before any call, actor should not be initialized
  629. assert workbench._actor is None # type: ignore[reportPrivateUsage]
  630. # Trigger list_tools → lazy start
  631. await workbench.list_tools()
  632. assert workbench._actor is not None # type: ignore[reportPrivateUsage]
  633. assert workbench._actor._active is True # type: ignore[reportPrivateUsage]
  634. actor = workbench._actor # type: ignore[reportPrivateUsage]
  635. del workbench
  636. await asyncio.sleep(0.1)
  637. assert actor._active is False
  638. @pytest.mark.asyncio
  639. async def test_del_to_new_event_loop_when_get_event_loop_fails() -> None:
  640. params = StdioServerParams(
  641. command="npx",
  642. args=[
  643. "-y",
  644. "@modelcontextprotocol/server-filesystem",
  645. ".",
  646. ],
  647. read_timeout_seconds=60,
  648. )
  649. workbench = McpWorkbench(server_params=params)
  650. await workbench.list_tools()
  651. assert workbench._actor is not None # type: ignore[reportPrivateUsage]
  652. assert workbench._actor._active is True # type: ignore[reportPrivateUsage]
  653. actor = workbench._actor # type: ignore[reportPrivateUsage]
  654. def cleanup() -> None:
  655. nonlocal workbench
  656. del workbench
  657. t = threading.Thread(target=cleanup)
  658. t.start()
  659. t.join()
  660. await asyncio.sleep(0.1)
  661. assert actor._active is False # type: ignore[reportPrivateUsage]
  662. def test_del_raises_when_loop_closed() -> None:
  663. loop = asyncio.new_event_loop()
  664. asyncio.set_event_loop(loop)
  665. params = StdioServerParams(command="echo", args=["ok"])
  666. workbench = McpWorkbench(server_params=params)
  667. workbench._actor_loop = loop # type: ignore[reportPrivateUsage]
  668. workbench._actor = cast(McpSessionActor, object()) # type: ignore[reportPrivateUsage]
  669. loop.close()
  670. with pytest.warns(RuntimeWarning, match="loop is closed or not running"):
  671. del workbench
  672. def test_mcp_tool_adapter_normalize_payload(sample_tool: Tool, sample_server_params: StdioServerParams) -> None:
  673. """Test the _normalize_payload_to_content_list method of McpToolAdapter."""
  674. adapter = StdioMcpToolAdapter(server_params=sample_server_params, tool=sample_tool)
  675. # Case 1: Payload is already a list of valid content items
  676. valid_content_list: list[TextContent | ImageContent | EmbeddedResource] = [
  677. TextContent(text="hello", type="text"),
  678. ImageContent(data="base64data", mimeType="image/png", type="image"),
  679. EmbeddedResource(
  680. type="resource",
  681. resource=TextResourceContents(text="embedded text", uri=AnyUrl(url="http://example.com/resource")),
  682. ),
  683. ]
  684. assert adapter._normalize_payload_to_content_list(valid_content_list) == valid_content_list # type: ignore[reportPrivateUsage]
  685. # Case 2: Payload is a single TextContent
  686. single_text_content = TextContent(text="single text", type="text")
  687. assert adapter._normalize_payload_to_content_list(single_text_content) == [single_text_content] # type: ignore[reportPrivateUsage, arg-type]
  688. # Case 3: Payload is a single ImageContent
  689. single_image_content = ImageContent(data="imagedata", mimeType="image/jpeg", type="image")
  690. assert adapter._normalize_payload_to_content_list(single_image_content) == [single_image_content] # type: ignore[reportPrivateUsage, arg-type]
  691. # Case 4: Payload is a single EmbeddedResource
  692. single_embedded_resource = EmbeddedResource(
  693. type="resource",
  694. resource=TextResourceContents(text="other embedded", uri=AnyUrl(url="http://example.com/other")),
  695. )
  696. assert adapter._normalize_payload_to_content_list(single_embedded_resource) == [single_embedded_resource] # type: ignore[reportPrivateUsage, arg-type]
  697. # Case 5: Payload is a string
  698. string_payload = "This is a string payload."
  699. expected_from_string = [TextContent(text=string_payload, type="text")]
  700. assert adapter._normalize_payload_to_content_list(string_payload) == expected_from_string # type: ignore[reportPrivateUsage, arg-type]
  701. # Case 6: Payload is an integer
  702. int_payload = 12345
  703. expected_from_int = [TextContent(text=str(int_payload), type="text")]
  704. assert adapter._normalize_payload_to_content_list(int_payload) == expected_from_int # type: ignore[reportPrivateUsage, arg-type]
  705. # Case 7: Payload is a dictionary
  706. dict_payload = {"key": "value", "number": 42}
  707. expected_from_dict = [TextContent(text=str(dict_payload), type="text")]
  708. assert adapter._normalize_payload_to_content_list(dict_payload) == expected_from_dict # type: ignore[reportPrivateUsage, arg-type]
  709. # Case 8: Payload is an empty list (should still be a list of valid items, so returns as is)
  710. empty_list_payload: list[TextContent | ImageContent | EmbeddedResource] = []
  711. assert adapter._normalize_payload_to_content_list(empty_list_payload) == empty_list_payload # type: ignore[reportPrivateUsage]
  712. # Case 9: Payload is None (should be stringified)
  713. none_payload = None
  714. expected_from_none = [TextContent(text=str(none_payload), type="text")]
  715. assert adapter._normalize_payload_to_content_list(none_payload) == expected_from_none # type: ignore[reportPrivateUsage, arg-type]
  716. @pytest.mark.asyncio
  717. async def test_mcp_tool_adapter_run_error(
  718. sample_tool: Tool,
  719. sample_server_params: StdioServerParams,
  720. mock_session: AsyncMock,
  721. mock_error_tool_response: MagicMock,
  722. cancellation_token: CancellationToken,
  723. ) -> None:
  724. """Test McpToolAdapter._run when tool returns an error."""
  725. adapter = StdioMcpToolAdapter(server_params=sample_server_params, tool=sample_tool, session=mock_session)
  726. mock_session.call_tool.return_value = mock_error_tool_response
  727. args = {"test_param": "test_value"}
  728. with pytest.raises(Exception) as excinfo:
  729. await adapter._run(args=args, cancellation_token=cancellation_token, session=mock_session) # type: ignore[reportPrivateUsage]
  730. mock_session.call_tool.assert_called_once_with(name=sample_tool.name, arguments=args)
  731. assert adapter.return_value_as_string([TextContent(text="error output", type="text")]) in str(excinfo.value)
  732. @pytest.mark.asyncio
  733. async def test_mcp_tool_adapter_run_cancelled_before_call(
  734. sample_tool: Tool,
  735. sample_server_params: StdioServerParams,
  736. mock_session: AsyncMock,
  737. cancellation_token: CancellationToken,
  738. ) -> None:
  739. """Test McpToolAdapter._run when operation is cancelled before tool call."""
  740. adapter = StdioMcpToolAdapter(server_params=sample_server_params, tool=sample_tool, session=mock_session)
  741. cancellation_token.cancel() # Cancel before the call
  742. args = {"test_param": "test_value"}
  743. with pytest.raises(asyncio.CancelledError):
  744. await adapter._run(args=args, cancellation_token=cancellation_token, session=mock_session) # type: ignore[reportPrivateUsage]
  745. mock_session.call_tool.assert_not_called()
  746. @pytest.mark.asyncio
  747. async def test_mcp_tool_adapter_run_cancelled_during_call(
  748. sample_tool: Tool,
  749. sample_server_params: StdioServerParams,
  750. mock_session: AsyncMock,
  751. cancellation_token: CancellationToken,
  752. ) -> None:
  753. """Test McpToolAdapter._run when operation is cancelled during tool call."""
  754. adapter = StdioMcpToolAdapter(server_params=sample_server_params, tool=sample_tool, session=mock_session)
  755. mock_session.call_tool.side_effect = asyncio.CancelledError("Tool call cancelled")
  756. args = {"test_param": "test_value"}
  757. with pytest.raises(asyncio.CancelledError):
  758. await adapter._run(args=args, cancellation_token=cancellation_token, session=mock_session) # type: ignore[reportPrivateUsage]
  759. mock_session.call_tool.assert_called_once_with(name=sample_tool.name, arguments=args)
  760. def test_return_value_as_string_with_resource_link(sample_tool: Tool, sample_server_params: StdioServerParams) -> None:
  761. """Test return_value_as_string handles ResourceLink objects correctly."""
  762. adapter = StdioMcpToolAdapter(server_params=sample_server_params, tool=sample_tool)
  763. # Test ResourceLink with meta field
  764. resource_link = ResourceLink(
  765. name="test_link",
  766. type="resource_link",
  767. uri=AnyUrl(url="http://example.com"),
  768. )
  769. result = adapter.return_value_as_string([resource_link])
  770. # Verify the JSON serialization contains expected fields
  771. assert '"type": "resource_link"' in result
  772. assert '"name": "test_link"' in result
  773. assert '"uri": "http://example.com/"' in result # AnyUrl normalizes with trailing slash