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_actor.py 56 kB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332
  1. """Tests for McpSessionActor to cover missing test coverage lines."""
  2. import asyncio
  3. import atexit
  4. import json
  5. from contextlib import asynccontextmanager
  6. from pathlib import Path
  7. from typing import Any, AsyncGenerator, Callable, Generator
  8. from unittest.mock import AsyncMock, MagicMock, patch
  9. import pytest
  10. from autogen_core import Image
  11. from autogen_core.models import (
  12. CreateResult,
  13. ModelInfo,
  14. RequestUsage,
  15. UserMessage,
  16. )
  17. from autogen_ext.tools.mcp import StdioServerParams
  18. from autogen_ext.tools.mcp._actor import (
  19. McpSessionActor,
  20. _parse_sampling_content, # pyright: ignore[reportPrivateUsage]
  21. _parse_sampling_message, # pyright: ignore[reportPrivateUsage]
  22. )
  23. from mcp import types as mcp_types
  24. from mcp.shared.context import RequestContext
  25. # Monkey patch to prevent atexit handlers from being registered during tests
  26. # This prevents the test suite from hanging during shutdown
  27. original_atexit_register = atexit.register
  28. def mock_atexit_register(func: Callable[[], None]) -> None:
  29. """Mock atexit.register to prevent registration during tests."""
  30. pass
  31. # Apply the monkey patch
  32. atexit.register = mock_atexit_register # type: ignore[assignment]
  33. @pytest.fixture
  34. def mcp_server_params() -> StdioServerParams:
  35. """Create server parameters that will launch the real MCP server subprocess."""
  36. # Get the path to the simple MCP server
  37. server_path = Path(__file__).parent.parent / "mcp_server_comprehensive.py"
  38. return StdioServerParams(
  39. command="uv",
  40. args=["run", "python", str(server_path)],
  41. read_timeout_seconds=10,
  42. )
  43. @pytest.fixture
  44. def mock_model_client() -> MagicMock:
  45. """Mock model client for testing."""
  46. model_client = MagicMock()
  47. model_client.model_info = {
  48. "vision": False,
  49. "function_calling": False,
  50. "json_output": False,
  51. "family": "test-model",
  52. "structured_output": False,
  53. }
  54. model_client.create = AsyncMock(
  55. return_value=CreateResult(
  56. content="Mock response",
  57. finish_reason="stop",
  58. usage=RequestUsage(prompt_tokens=10, completion_tokens=5),
  59. cached=False,
  60. )
  61. )
  62. return model_client
  63. @pytest.fixture
  64. def mock_model_client_with_vision() -> MagicMock:
  65. """Mock model client with vision support for testing."""
  66. model_client = MagicMock()
  67. model_client.model_info = {
  68. "vision": True,
  69. "function_calling": False,
  70. "json_output": False,
  71. "family": "test-vision-model",
  72. "structured_output": False,
  73. }
  74. model_client.create = AsyncMock(
  75. return_value=CreateResult(
  76. content="Mock response",
  77. finish_reason="stop",
  78. usage=RequestUsage(prompt_tokens=10, completion_tokens=5),
  79. cached=False,
  80. )
  81. )
  82. return model_client
  83. def test_parse_sampling_content_unsupported_image_without_vision() -> None:
  84. """Test _parse_sampling_content raises error for image content when model doesn't support vision (line 56)."""
  85. model_info: ModelInfo = {
  86. "vision": False,
  87. "function_calling": False,
  88. "json_output": False,
  89. "family": "test-model",
  90. "structured_output": False,
  91. }
  92. image_content = mcp_types.ImageContent(
  93. type="image",
  94. data="iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg==",
  95. mimeType="image/png",
  96. )
  97. with pytest.raises(ValueError, match="Sampling model does not support image content"):
  98. _parse_sampling_content(image_content, model_info)
  99. def test_parse_sampling_content_with_vision_support() -> None:
  100. """Test _parse_sampling_content works with vision-enabled model."""
  101. model_info: ModelInfo = {
  102. "vision": True,
  103. "function_calling": False,
  104. "json_output": False,
  105. "family": "test-model",
  106. "structured_output": False,
  107. }
  108. image_content = mcp_types.ImageContent(
  109. type="image",
  110. data="iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg==",
  111. mimeType="image/png",
  112. )
  113. result = _parse_sampling_content(image_content, model_info)
  114. assert isinstance(result, Image)
  115. def test_parse_sampling_content_text() -> None:
  116. """Test _parse_sampling_content with text content."""
  117. model_info: ModelInfo = {
  118. "vision": False,
  119. "function_calling": False,
  120. "json_output": False,
  121. "family": "test-model",
  122. "structured_output": False,
  123. }
  124. text_content = mcp_types.TextContent(type="text", text="Hello world")
  125. result = _parse_sampling_content(text_content, model_info)
  126. assert result == "Hello world"
  127. def test_parse_sampling_content_unknown_type() -> None:
  128. """Test _parse_sampling_content raises error for unknown content type."""
  129. model_info: ModelInfo = {
  130. "vision": False,
  131. "function_calling": False,
  132. "json_output": False,
  133. "family": "test-model",
  134. "structured_output": False,
  135. }
  136. # Create a mock content with unknown type
  137. unknown_content = MagicMock()
  138. unknown_content.type = "unknown"
  139. with pytest.raises(ValueError, match="Unsupported content type"):
  140. _parse_sampling_content(unknown_content, model_info)
  141. def test_parse_sampling_message_unrecognized_role() -> None:
  142. """Test _parse_sampling_message raises error for unrecognized role (lines 67-74)."""
  143. model_info: ModelInfo = {
  144. "vision": False,
  145. "function_calling": False,
  146. "json_output": False,
  147. "family": "test-model",
  148. "structured_output": False,
  149. }
  150. # Create a mock message with invalid role by bypassing type checking
  151. message = MagicMock()
  152. message.role = "system" # Invalid role that should trigger the error
  153. message.content = mcp_types.TextContent(type="text", text="Hello")
  154. with pytest.raises(ValueError, match="Unrecognized message role: system"):
  155. _parse_sampling_message(message, model_info)
  156. def test_parse_sampling_message_assistant_with_non_string_content() -> None:
  157. """Test _parse_sampling_message with assistant message containing non-string content."""
  158. model_info: ModelInfo = {
  159. "vision": True,
  160. "function_calling": False,
  161. "json_output": False,
  162. "family": "test-model",
  163. "structured_output": False,
  164. }
  165. # Create image content for assistant message (which should fail)
  166. image_content = mcp_types.ImageContent(
  167. type="image",
  168. data="iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg==",
  169. mimeType="image/png",
  170. )
  171. message = mcp_types.SamplingMessage(role="assistant", content=image_content)
  172. # This should raise an AssertionError because assistant messages only support string content
  173. with pytest.raises(AssertionError, match="Assistant messages only support string content"):
  174. _parse_sampling_message(message, model_info)
  175. @pytest.mark.asyncio
  176. async def test_call_when_not_active() -> None:
  177. """Test call method raises error when actor is not active (line 110)."""
  178. actor = McpSessionActor(StdioServerParams(command="echo", args=["test"]))
  179. with pytest.raises(RuntimeError, match="MCP Actor not running, call initialize\\(\\) first"):
  180. await actor.call("list_tools")
  181. @pytest.mark.asyncio
  182. async def test_call_when_actor_task_crashed() -> None:
  183. """Test call method raises error when actor task has crashed (line 119)."""
  184. actor = McpSessionActor(StdioServerParams(command="echo", args=["test"]))
  185. # Create a crashed task
  186. async def failing_task() -> Any:
  187. raise ValueError("Task crashed")
  188. actor._actor_task = asyncio.create_task(failing_task()) # type: ignore[reportPrivateUsage] # type: ignore[reportPrivateUsage]
  189. try:
  190. await asyncio.sleep(0.01) # Let the task crash
  191. actor._active = True # type: ignore[reportPrivateUsage]
  192. with pytest.raises(RuntimeError, match="MCP actor task crashed"):
  193. await actor.call("list_tools")
  194. finally:
  195. # Clean up the task
  196. if not actor._actor_task.done(): # type: ignore[reportPrivateUsage]
  197. actor._actor_task.cancel() # type: ignore[reportPrivateUsage]
  198. try:
  199. await actor._actor_task # type: ignore[reportPrivateUsage]
  200. except (asyncio.CancelledError, ValueError):
  201. pass
  202. @pytest.mark.asyncio
  203. async def test_call_without_required_args() -> None:
  204. """Test call method raises error when args are required but not provided (line 121)."""
  205. actor = McpSessionActor(StdioServerParams(command="echo", args=["test"]))
  206. actor._active = True # type: ignore[reportPrivateUsage]
  207. actor._actor_task = MagicMock() # type: ignore[reportPrivateUsage]
  208. actor._actor_task.done.return_value = False # type: ignore[reportPrivateUsage]
  209. with pytest.raises(ValueError, match="args is required for call_tool"):
  210. await actor.call("call_tool")
  211. @pytest.mark.asyncio
  212. async def test_call_tool_without_name() -> None:
  213. """Test call_tool raises error when name is not provided (line 128)."""
  214. actor = McpSessionActor(StdioServerParams(command="echo", args=["test"]))
  215. actor._active = True # type: ignore[reportPrivateUsage]
  216. actor._actor_task = MagicMock() # type: ignore[reportPrivateUsage]
  217. actor._actor_task.done.return_value = False # type: ignore[reportPrivateUsage]
  218. with pytest.raises(ValueError, match="name is required for call_tool"):
  219. await actor.call("call_tool", {"name": None, "kargs": {}})
  220. @pytest.mark.asyncio
  221. async def test_read_resource_without_uri() -> None:
  222. """Test read_resource raises error when uri is not provided (line 132)."""
  223. actor = McpSessionActor(StdioServerParams(command="echo", args=["test"]))
  224. actor._active = True # type: ignore[reportPrivateUsage]
  225. actor._actor_task = MagicMock() # type: ignore[reportPrivateUsage]
  226. actor._actor_task.done.return_value = False # type: ignore[reportPrivateUsage]
  227. with pytest.raises(ValueError, match="uri is required for read_resource"):
  228. await actor.call("read_resource", {"name": None, "kargs": {}})
  229. # @pytest.mark.asyncio
  230. # async def test_read_resource_command_queuing() -> None:
  231. # """Test read_resource command queuing (lines 134-137)."""
  232. # actor = McpSessionActor(StdioServerParams(command="echo", args=["test"]))
  233. # actor._active = True # type: ignore[reportPrivateUsage]
  234. # actor._actor_task = MagicMock() # type: ignore[reportPrivateUsage]
  235. # actor._actor_task.done.return_value = False # type: ignore[reportPrivateUsage]
  236. # # Mock the command queue to capture what gets put in it
  237. # original_queue = actor._command_queue # type: ignore[reportPrivateUsage]
  238. # actor._command_queue = MagicMock() # type: ignore[reportPrivateUsage]
  239. # actor._command_queue.put = AsyncMock() # type: ignore[reportPrivateUsage]
  240. # # Create a task that will be cancelled to avoid hanging
  241. # call_task = asyncio.create_task(actor.call("read_resource", {"name": None, "kargs": {"uri": "file:///test.txt"}}))
  242. # # Give it a brief moment to queue the command
  243. # await asyncio.sleep(0.001)
  244. # # Cancel the task to avoid hanging
  245. # call_task.cancel()
  246. # # Wait for the cancellation to complete
  247. # try:
  248. # await call_task
  249. # except asyncio.CancelledError:
  250. # pass # Expected
  251. # # Verify the command was queued correctly (this covers lines 134-137)
  252. # actor._command_queue.put.assert_called_once() # type: ignore[reportPrivateUsage]
  253. # call_args = actor._command_queue.put.call_args[0][0] # type: ignore[reportPrivateUsage]
  254. # assert call_args["type"] == "read_resource"
  255. # assert call_args["uri"] == "file:///test.txt"
  256. # assert "future" in call_args
  257. # # Restore original queue
  258. # actor._command_queue = original_queue # type: ignore[reportPrivateUsage]
  259. # @pytest.mark.asyncio
  260. # async def test_get_prompt_without_name() -> None:
  261. # """Test get_prompt raises error when name is not provided (lines 139-142)."""
  262. # actor = McpSessionActor(StdioServerParams(command="echo", args=["test"]))
  263. # actor._active = True # type: ignore[reportPrivateUsage]
  264. # actor._actor_task = MagicMock() # type: ignore[reportPrivateUsage]
  265. # actor._actor_task.done.return_value = False # type: ignore[reportPrivateUsage]
  266. # with pytest.raises(ValueError, match="name is required for get_prompt"):
  267. # await actor.call("get_prompt", {"name": None, "kargs": {}})
  268. # @pytest.mark.asyncio
  269. # async def test_get_prompt_command_queuing() -> None:
  270. # """Test get_prompt command queuing (lines 139-142)."""
  271. # actor = McpSessionActor(StdioServerParams(command="echo", args=["test"]))
  272. # actor._active = True # type: ignore[reportPrivateUsage]
  273. # actor._actor_task = MagicMock() # type: ignore[reportPrivateUsage]
  274. # actor._actor_task.done.return_value = False # type: ignore[reportPrivateUsage]
  275. # # Mock the command queue to capture what gets put in it
  276. # original_queue = actor._command_queue # type: ignore[reportPrivateUsage]
  277. # actor._command_queue = MagicMock() # type: ignore[reportPrivateUsage]
  278. # actor._command_queue.put = AsyncMock() # type: ignore[reportPrivateUsage]
  279. # # This will fail when trying to await the future, but we're testing the command queuing logic
  280. # try:
  281. # await actor.call("get_prompt", {"name": "test_prompt", "kargs": {"arguments": {"arg1": "value1"}}})
  282. # except Exception:
  283. # pass # Expected to fail, we're just testing the queuing logic
  284. # # Verify the command was queued correctly (this covers lines 139-142)
  285. # actor._command_queue.put.assert_called_once() # type: ignore[reportPrivateUsage]
  286. # call_args = actor._command_queue.put.call_args[0][0] # type: ignore[reportPrivateUsage]
  287. # assert call_args["type"] == "get_prompt"
  288. # assert call_args["name"] == "test_prompt"
  289. # assert call_args["args"] == {"arg1": "value1"}
  290. # assert "future" in call_args
  291. # # Restore original queue
  292. # actor._command_queue = original_queue # type: ignore[reportPrivateUsage]
  293. @pytest.mark.asyncio
  294. async def test_call_unknown_command_type() -> None:
  295. """Test call method raises error for unknown command type (line 147)."""
  296. actor = McpSessionActor(StdioServerParams(command="echo", args=["test"]))
  297. actor._active = True # type: ignore[reportPrivateUsage]
  298. actor._actor_task = MagicMock() # type: ignore[reportPrivateUsage]
  299. actor._actor_task.done.return_value = False # type: ignore[reportPrivateUsage]
  300. with pytest.raises(ValueError, match="Unknown command type: unknown_command"):
  301. await actor.call("unknown_command")
  302. @pytest.mark.asyncio
  303. async def test_close_when_not_active() -> None:
  304. """Test close method early return when not active (line 152)."""
  305. actor = McpSessionActor(StdioServerParams(command="echo", args=["test"]))
  306. actor._active = False # type: ignore[reportPrivateUsage]
  307. actor._actor_task = None # type: ignore[reportPrivateUsage]
  308. # This should return early without doing anything
  309. await actor.close()
  310. assert actor._shutdown_future is None # type: ignore[reportPrivateUsage]
  311. @pytest.mark.asyncio
  312. async def test_sampling_callback_without_model_client() -> None:
  313. """Test sampling callback returns error when no model client is available (line 194)."""
  314. actor = McpSessionActor(StdioServerParams(command="echo", args=["test"]), model_client=None)
  315. mock_context = MagicMock(spec=RequestContext)
  316. params = mcp_types.CreateMessageRequestParams(
  317. messages=[mcp_types.SamplingMessage(role="user", content=mcp_types.TextContent(type="text", text="Hello"))],
  318. maxTokens=100,
  319. )
  320. result = await actor._sampling_callback(mock_context, params) # type: ignore[reportPrivateUsage]
  321. assert isinstance(result, mcp_types.ErrorData)
  322. assert result.code == mcp_types.INVALID_REQUEST
  323. assert "No model client available" in result.message
  324. @pytest.mark.asyncio
  325. async def test_sampling_callback_message_processing_error() -> None:
  326. """Test sampling callback handles message processing errors (lines 226-227, 229-233)."""
  327. # Create a model client for the actor
  328. model_client = MagicMock()
  329. model_client.model_info = {
  330. "vision": False,
  331. "function_calling": False,
  332. "json_output": False,
  333. "family": "test-model",
  334. "structured_output": False,
  335. }
  336. actor = McpSessionActor(StdioServerParams(command="echo", args=["test"]), model_client=model_client)
  337. mock_context = MagicMock(spec=RequestContext)
  338. # Create a valid SamplingMessage but with invalid role that will cause parsing error
  339. # We'll patch the message after creation to bypass Pydantic validation
  340. valid_message = mcp_types.SamplingMessage(role="user", content=mcp_types.TextContent(type="text", text="Hello"))
  341. # Now change the role to something invalid that will cause _parse_sampling_message to fail
  342. valid_message.role = "invalid_role" # type: ignore
  343. params = mcp_types.CreateMessageRequestParams(
  344. messages=[valid_message],
  345. maxTokens=100,
  346. )
  347. result = await actor._sampling_callback(mock_context, params) # type: ignore[reportPrivateUsage]
  348. assert isinstance(result, mcp_types.ErrorData)
  349. assert result.code == mcp_types.INVALID_PARAMS
  350. assert "Error processing sampling messages" in result.message
  351. @pytest.mark.asyncio
  352. async def test_sampling_callback_model_client_error() -> None:
  353. """Test sampling callback handles model client errors (lines 235-239)."""
  354. failing_model_client = MagicMock()
  355. failing_model_client.model_info = {"vision": False, "family": "test-model"}
  356. failing_model_client.create = AsyncMock(side_effect=Exception("Model API error"))
  357. actor = McpSessionActor(StdioServerParams(command="echo", args=["test"]), model_client=failing_model_client)
  358. mock_context = MagicMock(spec=RequestContext)
  359. params = mcp_types.CreateMessageRequestParams(
  360. messages=[mcp_types.SamplingMessage(role="user", content=mcp_types.TextContent(type="text", text="Hello"))],
  361. maxTokens=100,
  362. )
  363. result = await actor._sampling_callback(mock_context, params) # type: ignore[reportPrivateUsage]
  364. assert isinstance(result, mcp_types.ErrorData)
  365. assert result.code == mcp_types.INTERNAL_ERROR
  366. assert "Error sampling from model client" in result.message
  367. assert "Model API error" in str(result.data)
  368. @pytest.mark.asyncio
  369. async def test_run_actor_exception_handling() -> None:
  370. """Test _run_actor exception handling for various command types (lines 244-268)."""
  371. # This test focuses on covering the exception handling code paths in _run_actor
  372. # We'll test this by creating a simplified scenario that exercises those paths
  373. actor = McpSessionActor(StdioServerParams(command="echo", args=["test"]))
  374. # Create a mock session that will raise exceptions by modifying our mock session
  375. @asynccontextmanager
  376. async def mock_failing_session(
  377. server_params: Any, sampling_callback: Any = None
  378. ) -> AsyncGenerator[MagicMock, None]:
  379. mock_session = MagicMock()
  380. mock_session.initialize = AsyncMock(
  381. return_value=mcp_types.InitializeResult(
  382. protocolVersion="1.0",
  383. capabilities=mcp_types.ServerCapabilities(),
  384. serverInfo=mcp_types.Implementation(name="test", version="1.0"),
  385. )
  386. )
  387. # Make sure all the session methods raise exceptions to cover the exception handling
  388. mock_session.call_tool = MagicMock(side_effect=Exception("Tool error"))
  389. mock_session.list_tools = MagicMock(side_effect=Exception("List tools error"))
  390. yield mock_session
  391. with patch("autogen_ext.tools.mcp._actor.create_mcp_server_session", mock_failing_session): # type: ignore[reportPrivateUsage]
  392. # Start the actor task
  393. actor._active = True # type: ignore[reportPrivateUsage]
  394. actor_task = asyncio.create_task(actor._run_actor()) # type: ignore[reportPrivateUsage]
  395. try:
  396. # Give it a moment to initialize
  397. await asyncio.sleep(0.05)
  398. # Test one command that will trigger exception handling (covers lines 244-268)
  399. future: asyncio.Future[Any] = asyncio.Future()
  400. cmd: dict[str, Any] = {"type": "call_tool", "name": "test_tool", "args": {}, "future": future}
  401. await actor._command_queue.put(cmd) # type: ignore[reportPrivateUsage]
  402. # Wait a bit for command to be processed
  403. await asyncio.sleep(0.05)
  404. # Send shutdown command
  405. shutdown_future: asyncio.Future[Any] = asyncio.Future()
  406. await actor._command_queue.put({"type": "shutdown", "future": shutdown_future}) # type: ignore[reportPrivateUsage]
  407. # Wait for actor to finish
  408. try:
  409. await asyncio.wait_for(actor_task, timeout=1.0)
  410. except asyncio.TimeoutError:
  411. pass # Expected if task doesn't finish properly
  412. # The key test: verify that the future was set with an exception
  413. # This proves the exception handling code (lines 244-268) was executed
  414. assert future.done()
  415. assert future.exception() is not None
  416. assert "Tool error" in str(future.exception())
  417. finally:
  418. # Ensure the task is cancelled and cleaned up
  419. if not actor_task.done():
  420. actor_task.cancel()
  421. try:
  422. await actor_task
  423. except asyncio.CancelledError:
  424. pass
  425. @pytest.mark.asyncio
  426. async def test_run_actor_session_exception() -> None:
  427. """Test _run_actor handles session creation exceptions (lines 274-288)."""
  428. actor = McpSessionActor(StdioServerParams(command="echo", args=["test"]))
  429. # Mock session creation to raise an exception
  430. with patch("autogen_ext.tools.mcp._actor.create_mcp_server_session", side_effect=Exception("Session error")): # type: ignore[reportPrivateUsage]
  431. actor._active = True # type: ignore[reportPrivateUsage]
  432. actor_task = asyncio.create_task(actor._run_actor()) # type: ignore[reportPrivateUsage]
  433. try:
  434. # Wait for the task to complete
  435. await asyncio.wait_for(actor_task, timeout=1.0)
  436. except asyncio.TimeoutError:
  437. # If it doesn't complete, cancel it
  438. actor_task.cancel()
  439. try:
  440. await actor_task
  441. except asyncio.CancelledError:
  442. pass
  443. # Check that the actor is no longer active
  444. assert not actor._active # type: ignore[reportPrivateUsage]
  445. assert actor._actor_task is None # type: ignore[reportPrivateUsage]
  446. @pytest.mark.asyncio
  447. async def test_run_actor_shutdown_future_exception() -> None:
  448. """Test _run_actor sets exception on shutdown future when session fails."""
  449. actor = McpSessionActor(StdioServerParams(command="echo", args=["test"]))
  450. actor._shutdown_future = asyncio.Future() # type: ignore[reportPrivateUsage]
  451. # Mock session creation to raise an exception
  452. with patch("autogen_ext.tools.mcp._actor.create_mcp_server_session", side_effect=Exception("Session error")): # type: ignore[reportPrivateUsage]
  453. actor._active = True # type: ignore[reportPrivateUsage]
  454. actor_task = asyncio.create_task(actor._run_actor()) # type: ignore[reportPrivateUsage]
  455. try:
  456. # Wait for the task to complete
  457. await asyncio.wait_for(actor_task, timeout=1.0)
  458. except asyncio.TimeoutError:
  459. # If it doesn't complete, cancel it
  460. actor_task.cancel()
  461. try:
  462. await actor_task
  463. except asyncio.CancelledError:
  464. pass
  465. # Check that shutdown future has the exception
  466. assert actor._shutdown_future.done() # type: ignore[reportPrivateUsage]
  467. assert actor._shutdown_future.exception() is not None # type: ignore[reportPrivateUsage]
  468. def test_sync_shutdown_when_not_active() -> None:
  469. """Test _sync_shutdown early return when not active (line 297)."""
  470. actor = McpSessionActor(StdioServerParams(command="echo", args=["test"]))
  471. actor._active = False # type: ignore[reportPrivateUsage]
  472. actor._actor_task = None # type: ignore[reportPrivateUsage]
  473. # This should return early without doing anything
  474. actor._sync_shutdown() # type: ignore[reportPrivateUsage]
  475. def test_sync_shutdown_no_event_loop() -> None:
  476. """Test _sync_shutdown handles RuntimeError when no event loop (line 310)."""
  477. actor = McpSessionActor(StdioServerParams(command="echo", args=["test"]))
  478. actor._active = True # type: ignore[reportPrivateUsage]
  479. actor._actor_task = MagicMock() # type: ignore[reportPrivateUsage]
  480. # Mock get_event_loop to raise RuntimeError
  481. with patch("asyncio.get_event_loop", side_effect=RuntimeError("No event loop")):
  482. # This should return early due to the RuntimeError
  483. actor._sync_shutdown() # type: ignore[reportPrivateUsage]
  484. def test_sync_shutdown_closed_loop() -> None:
  485. """Test _sync_shutdown handles closed event loop."""
  486. actor = McpSessionActor(StdioServerParams(command="echo", args=["test"]))
  487. actor._active = True # type: ignore[reportPrivateUsage]
  488. actor._actor_task = MagicMock() # type: ignore[reportPrivateUsage]
  489. # Mock event loop that is closed
  490. mock_loop = MagicMock()
  491. mock_loop.is_closed.return_value = True
  492. with patch("asyncio.get_event_loop", return_value=mock_loop):
  493. # This should return early due to the closed loop
  494. actor._sync_shutdown() # type: ignore[reportPrivateUsage]
  495. def test_sync_shutdown_running_loop() -> None:
  496. """Test _sync_shutdown creates task when loop is running."""
  497. actor = McpSessionActor(StdioServerParams(command="echo", args=["test"]))
  498. actor._active = True # type: ignore[reportPrivateUsage]
  499. actor._actor_task = MagicMock() # type: ignore[reportPrivateUsage]
  500. # Mock event loop that is running
  501. mock_loop = MagicMock()
  502. mock_loop.is_closed.return_value = False
  503. mock_loop.is_running.return_value = True
  504. mock_loop.create_task = MagicMock()
  505. with patch("asyncio.get_event_loop", return_value=mock_loop):
  506. actor._sync_shutdown() # type: ignore[reportPrivateUsage]
  507. # Should create a task to close the actor
  508. mock_loop.create_task.assert_called_once()
  509. def test_sync_shutdown_non_running_loop() -> None:
  510. """Test _sync_shutdown runs until complete when loop is not running."""
  511. actor = McpSessionActor(StdioServerParams(command="echo", args=["test"]))
  512. actor._active = True # type: ignore[reportPrivateUsage]
  513. actor._actor_task = MagicMock() # type: ignore[reportPrivateUsage]
  514. # Mock event loop that is not running
  515. mock_loop = MagicMock()
  516. mock_loop.is_closed.return_value = False
  517. mock_loop.is_running.return_value = False
  518. mock_loop.run_until_complete = MagicMock()
  519. with patch("asyncio.get_event_loop", return_value=mock_loop):
  520. actor._sync_shutdown() # type: ignore[reportPrivateUsage]
  521. # Should run until complete
  522. mock_loop.run_until_complete.assert_called_once()
  523. def test_to_config() -> None:
  524. """Test _to_config method."""
  525. server_params = StdioServerParams(command="echo", args=["test"])
  526. actor = McpSessionActor(server_params)
  527. config = actor._to_config() # type: ignore[reportPrivateUsage]
  528. assert config.server_params == server_params
  529. def test_from_config() -> None:
  530. """Test _from_config class method."""
  531. from autogen_ext.tools.mcp._actor import McpSessionActorConfig # type: ignore[reportPrivateUsage]
  532. server_params = StdioServerParams(command="echo", args=["test"])
  533. config = McpSessionActorConfig(server_params=server_params)
  534. actor = McpSessionActor._from_config(config) # type: ignore[reportPrivateUsage]
  535. assert actor.server_params == server_params
  536. @pytest.mark.asyncio
  537. async def test_initialize_result_property() -> None:
  538. """Test initialize_result property."""
  539. actor = McpSessionActor(StdioServerParams(command="echo", args=["test"]))
  540. # Initially should be None
  541. assert actor.initialize_result is None
  542. # Set a mock result
  543. mock_result = mcp_types.InitializeResult(
  544. protocolVersion="1.0",
  545. capabilities=mcp_types.ServerCapabilities(),
  546. serverInfo=mcp_types.Implementation(name="test", version="1.0"),
  547. )
  548. actor._initialize_result = mock_result # type: ignore[reportPrivateUsage]
  549. assert actor.initialize_result == mock_result
  550. @pytest.mark.asyncio
  551. async def test_actor_initialization() -> None:
  552. """Test actor initialization sets up correctly."""
  553. server_params = StdioServerParams(command="echo", args=["test"])
  554. model_client = MagicMock()
  555. actor = McpSessionActor(server_params, model_client=model_client)
  556. # Check initial state
  557. assert actor.server_params == server_params
  558. assert actor._model_client == model_client # type: ignore[reportPrivateUsage]
  559. assert actor.name == "mcp_session_actor"
  560. assert actor.description == "MCP session actor"
  561. assert not actor._active # type: ignore[reportPrivateUsage]
  562. assert actor._actor_task is None # type: ignore[reportPrivateUsage]
  563. assert actor._shutdown_future is None # type: ignore[reportPrivateUsage]
  564. assert actor._initialize_result is None # type: ignore[reportPrivateUsage]
  565. @pytest.mark.asyncio
  566. async def test_initialize_method(mcp_server_params: Any) -> None:
  567. """Test initialize method."""
  568. actor = McpSessionActor(mcp_server_params)
  569. await actor.initialize()
  570. assert actor._active # type: ignore[reportPrivateUsage]
  571. assert actor._actor_task is not None # type: ignore[reportPrivateUsage]
  572. # Clean up
  573. await actor.close()
  574. @pytest.mark.asyncio
  575. async def test_call_with_valid_list_commands() -> None:
  576. """Test call method with valid list commands."""
  577. actor = McpSessionActor(StdioServerParams(command="echo", args=["test"]))
  578. actor._active = True # type: ignore[reportPrivateUsage]
  579. actor._actor_task = MagicMock() # type: ignore[reportPrivateUsage]
  580. actor._actor_task.done.return_value = False # type: ignore[reportPrivateUsage]
  581. # Mock the command queue to capture what gets put in it
  582. original_queue = actor._command_queue # type: ignore[reportPrivateUsage]
  583. actor._command_queue = MagicMock() # type: ignore[reportPrivateUsage]
  584. actor._command_queue.put = AsyncMock() # type: ignore[reportPrivateUsage]
  585. # Test all valid list commands - we don't await the result, just test the queuing
  586. test_commands = ["list_tools", "list_prompts", "list_resources", "list_resource_templates", "shutdown"]
  587. created_tasks: list[asyncio.Task[Any]] = []
  588. try:
  589. for cmd_type in test_commands:
  590. # Create a task but don't await it (to avoid hanging)
  591. call_task = asyncio.create_task(actor.call(cmd_type))
  592. created_tasks.append(call_task)
  593. # Give it a brief moment to queue the command
  594. await asyncio.sleep(0.001)
  595. # Verify the command was queued correctly
  596. actor._command_queue.put.assert_called() # type: ignore[reportPrivateUsage]
  597. call_args = actor._command_queue.put.call_args[0][0] # type: ignore[reportPrivateUsage]
  598. assert call_args["type"] == cmd_type
  599. assert "future" in call_args
  600. actor._command_queue.put.reset_mock() # type: ignore[reportPrivateUsage]
  601. finally:
  602. # Clean up all created tasks
  603. for task in created_tasks:
  604. if not task.done():
  605. task.cancel()
  606. try:
  607. await task
  608. except asyncio.CancelledError:
  609. pass
  610. # Restore original queue
  611. actor._command_queue = original_queue # type: ignore[reportPrivateUsage]
  612. # Integration tests using the real MCP server
  613. @pytest.mark.asyncio
  614. async def test_actor_basic_functionality(mcp_server_params: Any) -> None:
  615. """Test basic actor functionality with real MCP server."""
  616. actor = McpSessionActor(mcp_server_params)
  617. try:
  618. # Initialize the actor
  619. await actor.initialize()
  620. assert actor._active # type: ignore[reportPrivateUsage]
  621. assert actor._actor_task is not None # type: ignore[reportPrivateUsage]
  622. # Test listing tools
  623. tools_future = await actor.call("list_tools")
  624. tools_result: mcp_types.ListToolsResult = await tools_future # type: ignore
  625. assert len(tools_result.tools) == 2 # echo and get_time
  626. tool_names = [tool.name for tool in tools_result.tools]
  627. assert "echo" in tool_names
  628. assert "get_time" in tool_names
  629. # Test calling a tool
  630. call_future = await actor.call("call_tool", {"name": "echo", "kargs": {"text": "Hello World"}})
  631. call_result: mcp_types.CallToolResult = await call_future # type: ignore
  632. assert call_result.content[0].text == "Echo: Hello World" # type: ignore
  633. finally:
  634. await actor.close()
  635. @pytest.mark.asyncio
  636. async def test_actor_prompt_operations(mcp_server_params: Any) -> None:
  637. """Test actor prompt operations with real MCP server."""
  638. actor = McpSessionActor(mcp_server_params)
  639. try:
  640. await actor.initialize()
  641. await asyncio.sleep(0.1)
  642. # Test listing prompts
  643. prompts_future = await actor.call("list_prompts")
  644. prompts_result: mcp_types.ListPromptsResult = await prompts_future # type: ignore
  645. assert len(prompts_result.prompts) == 2 # code_review and documentation
  646. prompt_names = [prompt.name for prompt in prompts_result.prompts]
  647. assert "code_review" in prompt_names
  648. assert "documentation" in prompt_names
  649. # Test getting a prompt with arguments
  650. prompt_future = await actor.call(
  651. "get_prompt",
  652. {"name": "code_review", "kargs": {"arguments": {"code": "print('hello')", "language": "python"}}},
  653. )
  654. prompt_result: mcp_types.GetPromptResult = await prompt_future # type: ignore
  655. assert prompt_result.description is not None and "python" in prompt_result.description
  656. assert "print('hello')" in prompt_result.messages[0].content.text # type: ignore
  657. finally:
  658. await actor.close()
  659. @pytest.mark.asyncio
  660. async def test_actor_resource_operations(mcp_server_params: Any) -> None:
  661. """Test actor resource operations with real MCP server."""
  662. actor = McpSessionActor(mcp_server_params)
  663. try:
  664. await actor.initialize()
  665. await asyncio.sleep(0.1)
  666. # Test listing resources
  667. resources_future = await actor.call("list_resources")
  668. resources_result: mcp_types.ListResourcesResult = await resources_future # type: ignore
  669. assert len(resources_result.resources) == 2 # users and projects
  670. resource_names = [resource.name for resource in resources_result.resources]
  671. assert "Company Users" in resource_names
  672. assert "Active Projects" in resource_names
  673. # Test reading a resource
  674. read_future = await actor.call("read_resource", {"name": None, "kargs": {"uri": "file:///company/users.json"}})
  675. read_result: mcp_types.ReadResourceResult = await read_future # type: ignore
  676. users_data = json.loads(read_result.contents[0].text) # type: ignore
  677. assert isinstance(users_data, list)
  678. assert len(users_data) == 3 # type: ignore[reportUnknownArgumentType]
  679. assert users_data[0]["name"] == "Alice"
  680. finally:
  681. await actor.close()
  682. @pytest.mark.asyncio
  683. async def test_actor_tool_failure_handling(mcp_server_params: Any) -> None:
  684. """Test actor handles tool failures correctly."""
  685. actor = McpSessionActor(mcp_server_params)
  686. try:
  687. await actor.initialize()
  688. await asyncio.sleep(0.1)
  689. # Test calling an unknown tool - the server should return an error result
  690. call_future = await actor.call("call_tool", {"name": "unknown_tool", "kargs": {}})
  691. call_result: mcp_types.CallToolResult = await call_future # type: ignore
  692. # The server returns an error but doesn't raise an exception
  693. assert call_result.isError is True # type: ignore
  694. assert "Unknown tool" in call_result.content[0].text # type: ignore
  695. finally:
  696. await actor.close()
  697. @pytest.mark.asyncio
  698. async def test_actor_with_model_client_sampling(mcp_server_params: Any, mock_model_client: Any) -> None:
  699. """Test actor with model client for sampling operations."""
  700. actor = McpSessionActor(mcp_server_params, model_client=mock_model_client)
  701. try:
  702. await actor.initialize()
  703. await asyncio.sleep(0.1)
  704. # Test sampling callback functionality
  705. mock_context = MagicMock(spec=RequestContext)
  706. params = mcp_types.CreateMessageRequestParams(
  707. messages=[
  708. mcp_types.SamplingMessage(
  709. role="user", content=mcp_types.TextContent(type="text", text="Hello from test")
  710. )
  711. ],
  712. maxTokens=100,
  713. )
  714. result = await actor._sampling_callback(mock_context, params) # type: ignore[reportPrivateUsage]
  715. assert isinstance(result, mcp_types.CreateMessageResult)
  716. assert result.role == "assistant"
  717. assert isinstance(result.content, mcp_types.TextContent)
  718. assert result.content.text == "Mock response"
  719. assert result.model == "test-model"
  720. finally:
  721. await actor.close()
  722. # Integration tests with real MCP server
  723. @pytest.mark.asyncio
  724. async def test_actor(mcp_server_params: Any) -> None:
  725. """Test actor with real MCP server subprocess."""
  726. actor = McpSessionActor(mcp_server_params)
  727. try:
  728. # Initialize the actor
  729. await actor.initialize()
  730. assert actor._active # type: ignore[reportPrivateUsage]
  731. assert actor._actor_task is not None # type: ignore[reportPrivateUsage]
  732. # Test listing tools
  733. tools_future = await actor.call("list_tools")
  734. tools_result: mcp_types.ListToolsResult = await tools_future # type: ignore
  735. assert len(tools_result.tools) == 2 # echo and get_time
  736. tool_names = [tool.name for tool in tools_result.tools]
  737. assert "echo" in tool_names
  738. assert "get_time" in tool_names
  739. # Test calling the echo tool
  740. call_future = await actor.call("call_tool", {"name": "echo", "kargs": {"text": "Hello World"}})
  741. call_result: mcp_types.CallToolResult = await call_future # type: ignore
  742. assert call_result.content[0].text == "Echo: Hello World" # type: ignore
  743. # Test calling the get_time tool
  744. time_future = await actor.call("call_tool", {"name": "get_time", "kargs": {}})
  745. time_result: mcp_types.CallToolResult = await time_future # type: ignore
  746. assert "Current time:" in time_result.content[0].text # type: ignore
  747. finally:
  748. await actor.close()
  749. @pytest.mark.asyncio
  750. async def test_actor_prompts(mcp_server_params: Any) -> None:
  751. """Test actor prompt operations with real MCP server."""
  752. actor = McpSessionActor(mcp_server_params)
  753. try:
  754. await actor.initialize()
  755. # Test listing prompts
  756. prompts_future = await actor.call("list_prompts")
  757. prompts_result: mcp_types.ListPromptsResult = await prompts_future # type: ignore
  758. assert len(prompts_result.prompts) == 2 # code_review and documentation
  759. prompt_names = [prompt.name for prompt in prompts_result.prompts]
  760. assert "code_review" in prompt_names
  761. assert "documentation" in prompt_names
  762. # Test getting a prompt
  763. prompt_future = await actor.call(
  764. "get_prompt",
  765. {"name": "code_review", "kargs": {"arguments": {"code": "print('hello')", "language": "python"}}},
  766. )
  767. prompt_result: mcp_types.GetPromptResult = await prompt_future # type: ignore
  768. assert prompt_result.description is not None and "python" in prompt_result.description
  769. assert "print('hello')" in prompt_result.messages[0].content.text # type: ignore
  770. finally:
  771. await actor.close()
  772. @pytest.mark.asyncio
  773. async def test_actor_resources(mcp_server_params: Any) -> None:
  774. """Test actor resource operations with real MCP server."""
  775. actor = McpSessionActor(mcp_server_params)
  776. try:
  777. await actor.initialize()
  778. # Test listing resources
  779. resources_future = await actor.call("list_resources")
  780. resources_result: mcp_types.ListResourcesResult = await resources_future # type: ignore
  781. assert len(resources_result.resources) == 2 # users and projects
  782. resource_names = [resource.name for resource in resources_result.resources]
  783. assert "Company Users" in resource_names
  784. assert "Active Projects" in resource_names
  785. # Test reading a resource
  786. read_future = await actor.call("read_resource", {"name": None, "kargs": {"uri": "file:///company/users.json"}})
  787. read_result: mcp_types.ReadResourceResult = await read_future # type: ignore
  788. # The real server returns content in the ReadResourceResult
  789. users_data = json.loads(read_result.contents[0].text) # type: ignore
  790. assert isinstance(users_data, list)
  791. assert len(users_data) == 3 # type: ignore[reportUnknownArgumentType]
  792. assert users_data[0]["name"] == "Alice"
  793. finally:
  794. await actor.close()
  795. @pytest.mark.asyncio
  796. async def test_actor_unknown_tool(mcp_server_params: Any) -> None:
  797. """Test actor handles unknown tools with real MCP server."""
  798. actor = McpSessionActor(mcp_server_params)
  799. try:
  800. await actor.initialize()
  801. # Test calling an unknown tool - the server should return an error result
  802. call_future = await actor.call("call_tool", {"name": "unknown_tool", "kargs": {}})
  803. call_result: mcp_types.CallToolResult = await call_future # type: ignore
  804. # The server returns an error but doesn't raise an exception
  805. assert call_result.isError is True # type: ignore
  806. assert "Unknown tool" in call_result.content[0].text # type: ignore
  807. finally:
  808. await actor.close()
  809. @pytest.fixture
  810. def clean_actor() -> Generator[Callable[..., McpSessionActor], None, None]:
  811. """Fixture to track and clean up actors created in tests."""
  812. actors: list[McpSessionActor] = []
  813. def create_actor(*args: Any, **kwargs: Any) -> McpSessionActor:
  814. actor = McpSessionActor(*args, **kwargs)
  815. actors.append(actor)
  816. return actor
  817. yield create_actor
  818. # Clean up all actors
  819. for actor in actors:
  820. if hasattr(actor, "_active") and actor._active: # type: ignore[reportPrivateUsage]
  821. try:
  822. # Try to close the actor properly
  823. loop = asyncio.get_event_loop()
  824. if loop.is_running():
  825. loop.create_task(actor.close())
  826. else:
  827. loop.run_until_complete(actor.close())
  828. except Exception:
  829. # If we can't close it properly, at least deactivate it
  830. actor._active = False # type: ignore[reportPrivateUsage]
  831. if hasattr(actor, "_actor_task") and actor._actor_task: # type: ignore[reportPrivateUsage]
  832. actor._actor_task.cancel() # type: ignore[reportPrivateUsage]
  833. @pytest.mark.asyncio
  834. async def test_sampling_callback_with_system_prompt() -> None:
  835. """Test sampling callback with systemPrompt parameter (line 177)."""
  836. model_client = MagicMock()
  837. model_client.model_info = {
  838. "vision": False,
  839. "function_calling": False,
  840. "json_output": False,
  841. "family": "test-model",
  842. "structured_output": False,
  843. }
  844. model_client.create = AsyncMock(
  845. return_value=CreateResult(
  846. content="Mock response",
  847. finish_reason="stop",
  848. usage=RequestUsage(prompt_tokens=10, completion_tokens=5),
  849. cached=False,
  850. )
  851. )
  852. actor = McpSessionActor(StdioServerParams(command="echo", args=["test"]), model_client=model_client)
  853. mock_context = MagicMock(spec=RequestContext)
  854. params = mcp_types.CreateMessageRequestParams(
  855. messages=[mcp_types.SamplingMessage(role="user", content=mcp_types.TextContent(type="text", text="Hello"))],
  856. maxTokens=100,
  857. systemPrompt="You are a helpful assistant.",
  858. )
  859. result = await actor._sampling_callback(mock_context, params) # type: ignore[reportPrivateUsage]
  860. assert isinstance(result, mcp_types.CreateMessageResult)
  861. assert result.role == "assistant"
  862. assert isinstance(result.content, mcp_types.TextContent)
  863. assert result.content.text == "Mock response"
  864. # Verify that the model client was called with the system prompt
  865. model_client.create.assert_called_once()
  866. call_args = model_client.create.call_args[1]
  867. messages = call_args["messages"]
  868. assert len(messages) == 2 # SystemMessage + UserMessage
  869. assert messages[0].content == "You are a helpful assistant."
  870. @pytest.mark.asyncio
  871. async def test_sampling_callback_with_non_string_content() -> None:
  872. """Test sampling callback when model returns non-string content (line 194)."""
  873. model_client = MagicMock()
  874. model_client.model_info = {
  875. "vision": False,
  876. "function_calling": False,
  877. "json_output": False,
  878. "family": "test-model",
  879. "structured_output": False,
  880. }
  881. # Mock the model to return a non-string content
  882. non_string_content = {"data": "complex object"}
  883. model_client.create = AsyncMock(
  884. return_value=CreateResult(
  885. content=str(non_string_content), # Convert to string for valid content
  886. finish_reason="stop",
  887. usage=RequestUsage(prompt_tokens=10, completion_tokens=5),
  888. cached=False,
  889. )
  890. )
  891. actor = McpSessionActor(StdioServerParams(command="echo", args=["test"]), model_client=model_client)
  892. mock_context = MagicMock(spec=RequestContext)
  893. params = mcp_types.CreateMessageRequestParams(
  894. messages=[mcp_types.SamplingMessage(role="user", content=mcp_types.TextContent(type="text", text="Hello"))],
  895. maxTokens=100,
  896. )
  897. result = await actor._sampling_callback(mock_context, params) # type: ignore[reportPrivateUsage]
  898. assert isinstance(result, mcp_types.CreateMessageResult)
  899. assert result.role == "assistant"
  900. assert isinstance(result.content, mcp_types.TextContent)
  901. # Should be converted to string
  902. assert result.content.text == "{'data': 'complex object'}"
  903. @pytest.mark.asyncio
  904. async def test_run_actor_all_command_types_exception_handling() -> None:
  905. """Test _run_actor exception handling for all command types (lines 232-233, 238-239, 244-245, 250-251, 256-263)."""
  906. actor = McpSessionActor(StdioServerParams(command="echo", args=["test"]))
  907. # Create a mock session that will raise exceptions for all command types
  908. @asynccontextmanager
  909. async def mock_failing_session(
  910. server_params: Any, sampling_callback: Any = None
  911. ) -> AsyncGenerator[MagicMock, None]:
  912. mock_session = MagicMock()
  913. mock_session.initialize = AsyncMock(
  914. return_value=mcp_types.InitializeResult(
  915. protocolVersion="1.0",
  916. capabilities=mcp_types.ServerCapabilities(),
  917. serverInfo=mcp_types.Implementation(name="test", version="1.0"),
  918. )
  919. )
  920. # Make all session methods raise exceptions
  921. mock_session.call_tool = MagicMock(side_effect=Exception("call_tool error"))
  922. mock_session.read_resource = MagicMock(side_effect=Exception("read_resource error"))
  923. mock_session.get_prompt = MagicMock(side_effect=Exception("get_prompt error"))
  924. mock_session.list_tools = MagicMock(side_effect=Exception("list_tools error"))
  925. mock_session.list_prompts = MagicMock(side_effect=Exception("list_prompts error"))
  926. mock_session.list_resources = MagicMock(side_effect=Exception("list_resources error"))
  927. mock_session.list_resource_templates = MagicMock(side_effect=Exception("list_resource_templates error"))
  928. yield mock_session
  929. with patch("autogen_ext.tools.mcp._actor.create_mcp_server_session", mock_failing_session): # type: ignore[reportPrivateUsage]
  930. # Start the actor task
  931. actor._active = True # type: ignore[reportPrivateUsage]
  932. actor_task = asyncio.create_task(actor._run_actor()) # type: ignore[reportPrivateUsage]
  933. try:
  934. # Give it a moment to initialize
  935. await asyncio.sleep(0.05)
  936. # Test all command types that can raise exceptions
  937. commands_to_test: list[dict[str, Any]] = [
  938. {"type": "call_tool", "name": "test_tool", "args": {}},
  939. {"type": "read_resource", "uri": "test://resource"},
  940. {"type": "get_prompt", "name": "test_prompt", "args": {}},
  941. {"type": "list_tools"},
  942. {"type": "list_prompts"},
  943. {"type": "list_resources"},
  944. {"type": "list_resource_templates"},
  945. ]
  946. futures: list[asyncio.Future[Any]] = []
  947. for cmd in commands_to_test:
  948. future: asyncio.Future[Any] = asyncio.Future()
  949. cmd["future"] = future
  950. await actor._command_queue.put(cmd) # type: ignore[reportPrivateUsage]
  951. futures.append(future)
  952. # Wait a bit for commands to be processed
  953. await asyncio.sleep(0.1)
  954. # Send shutdown command
  955. shutdown_future: asyncio.Future[Any] = asyncio.Future()
  956. await actor._command_queue.put({"type": "shutdown", "future": shutdown_future}) # type: ignore[reportPrivateUsage]
  957. # Wait for actor to finish
  958. try:
  959. await asyncio.wait_for(actor_task, timeout=1.0)
  960. except asyncio.TimeoutError:
  961. pass # Expected if task doesn't finish properly
  962. # Verify that all futures were set with exceptions
  963. for i, future in enumerate(futures):
  964. assert future.done(), f"Future {i} was not completed"
  965. assert future.exception() is not None, f"Future {i} should have an exception"
  966. finally:
  967. # Ensure the task is cancelled and cleaned up
  968. if not actor_task.done():
  969. actor_task.cancel()
  970. try:
  971. await actor_task
  972. except asyncio.CancelledError:
  973. pass
  974. @pytest.mark.asyncio
  975. async def test_close_with_shutdown_await() -> None:
  976. """Test close method waits for shutdown future (line 140)."""
  977. actor = McpSessionActor(StdioServerParams(command="echo", args=["test"]))
  978. actor._active = True # type: ignore[reportPrivateUsage]
  979. actor._actor_task = MagicMock() # type: ignore[reportPrivateUsage]
  980. actor._actor_task.done.return_value = False # type: ignore[reportPrivateUsage]
  981. # Mock the command queue and actor task
  982. original_queue = actor._command_queue # type: ignore[reportPrivateUsage]
  983. actor._command_queue = MagicMock() # type: ignore[reportPrivateUsage]
  984. actor._command_queue.put = AsyncMock() # type: ignore[reportPrivateUsage]
  985. # Create a shutdown future that will be set
  986. shutdown_future: asyncio.Future[Any] = asyncio.Future()
  987. # Mock the close method to set the shutdown future after a delay
  988. async def mock_close() -> None:
  989. actor._shutdown_future = shutdown_future # type: ignore[reportPrivateUsage]
  990. await actor._command_queue.put({"type": "shutdown", "future": shutdown_future}) # type: ignore[reportPrivateUsage]
  991. # Simulate the shutdown completion
  992. await asyncio.sleep(0.01)
  993. shutdown_future.set_result("ok")
  994. # Replace the close method temporarily
  995. original_close = actor.close
  996. actor.close = mock_close # type: ignore[method-assign]
  997. try:
  998. # This should complete without hanging
  999. await actor.close()
  1000. assert shutdown_future.done()
  1001. assert shutdown_future.result() == "ok"
  1002. finally:
  1003. # Restore original queue and close method
  1004. actor._command_queue = original_queue # type: ignore[reportPrivateUsage]
  1005. actor.close = original_close # type: ignore[method-assign]
  1006. def test_parse_sampling_message_user_role() -> None:
  1007. """Test _parse_sampling_message with user role (line 69)."""
  1008. model_info: ModelInfo = {
  1009. "vision": False,
  1010. "function_calling": False,
  1011. "json_output": False,
  1012. "family": "test-model",
  1013. "structured_output": False,
  1014. }
  1015. message = mcp_types.SamplingMessage(role="user", content=mcp_types.TextContent(type="text", text="Hello user"))
  1016. result = _parse_sampling_message(message, model_info)
  1017. assert isinstance(result, UserMessage)
  1018. assert result.content == ["Hello user"]
  1019. @pytest.mark.asyncio
  1020. async def test_call_tool_command_queuing() -> None:
  1021. """Test call_tool command queuing (line 140)."""
  1022. actor = McpSessionActor(StdioServerParams(command="echo", args=["test"]))
  1023. actor._active = True # type: ignore[reportPrivateUsage]
  1024. actor._actor_task = MagicMock() # type: ignore[reportPrivateUsage]
  1025. actor._actor_task.done.return_value = False # type: ignore[reportPrivateUsage]
  1026. # Mock the command queue to capture what gets put in it
  1027. original_queue = actor._command_queue # type: ignore[reportPrivateUsage]
  1028. actor._command_queue = MagicMock() # type: ignore[reportPrivateUsage]
  1029. actor._command_queue.put = AsyncMock() # type: ignore[reportPrivateUsage]
  1030. # Create a task that will be cancelled to avoid hanging
  1031. call_task = asyncio.create_task(actor.call("call_tool", {"name": "test_tool", "kargs": {"param": "value"}}))
  1032. # Give it a brief moment to queue the command
  1033. await asyncio.sleep(0.001)
  1034. # Cancel the task to avoid hanging
  1035. call_task.cancel()
  1036. # Wait for the cancellation to complete
  1037. try:
  1038. await call_task
  1039. except asyncio.CancelledError:
  1040. pass # Expected
  1041. # Verify the command was queued correctly (this covers line 140)
  1042. actor._command_queue.put.assert_called_once() # type: ignore[reportPrivateUsage]
  1043. call_args = actor._command_queue.put.call_args[0][0] # type: ignore[reportPrivateUsage]
  1044. assert call_args["type"] == "call_tool"
  1045. assert call_args["name"] == "test_tool"
  1046. assert call_args["args"] == {"param": "value"}
  1047. assert "future" in call_args
  1048. # Restore original queue
  1049. actor._command_queue = original_queue # type: ignore[reportPrivateUsage]
  1050. def test_parse_sampling_message_user_role_branch() -> None:
  1051. """Test _parse_sampling_message with user role (line 69)."""
  1052. model_info: ModelInfo = {
  1053. "vision": False,
  1054. "function_calling": False,
  1055. "json_output": False,
  1056. "family": "test-model",
  1057. "structured_output": False,
  1058. }
  1059. message = mcp_types.SamplingMessage(role="user", content=mcp_types.TextContent(type="text", text="Hello user test"))
  1060. result = _parse_sampling_message(message, model_info)
  1061. # This specifically tests the user role branch (line 69)
  1062. assert isinstance(result, UserMessage)
  1063. assert result.content == ["Hello user test"]