From 588377ee3ee5e556247d2d3f7932bef85e31b3ed Mon Sep 17 00:00:00 2001 From: "chiyoung.song" Date: Fri, 9 May 2025 20:56:35 +0900 Subject: [PATCH 1/3] FIX/McpWorkbench_del_with_closed_session --- .../src/autogen_ext/tools/mcp/_workbench.py | 13 ++++++++++++- .../autogen-ext/tests/tools/test_mcp_tools.py | 4 ++++ 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_workbench.py b/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_workbench.py index 419ae34bfe2e..4aa1a48dfc74 100644 --- a/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_workbench.py +++ b/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_workbench.py @@ -1,3 +1,4 @@ +import asyncio import builtins import warnings from typing import Any, List, Literal, Mapping @@ -282,4 +283,14 @@ def _from_config(cls, config: McpWorkbenchConfig) -> Self: def __del__(self) -> None: # Ensure the actor is stopped when the workbench is deleted - pass + if self._actor: + try: + loop = asyncio.get_event_loop() + except RuntimeError: + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + + if loop.is_running(): + loop.create_task(self.stop()) + else: + loop.run_until_complete(self.stop()) diff --git a/python/packages/autogen-ext/tests/tools/test_mcp_tools.py b/python/packages/autogen-ext/tests/tools/test_mcp_tools.py index de9fd5cb513c..3fcea2636a3c 100644 --- a/python/packages/autogen-ext/tests/tools/test_mcp_tools.py +++ b/python/packages/autogen-ext/tests/tools/test_mcp_tools.py @@ -1,5 +1,6 @@ import logging import os +import asyncio from unittest.mock import AsyncMock, MagicMock import pytest @@ -594,4 +595,7 @@ async def test_lazy_init_and_finalize_cleanup() -> None: assert workbench._actor is not None # type: ignore[reportPrivateUsage] assert workbench._actor._active is True # type: ignore[reportPrivateUsage] + actor = workbench._actor # type: ignore[reportPrivateUsage] del workbench + await asyncio.sleep(0.1) + assert actor._active is False \ No newline at end of file From 8e9fc0b955e6c2f515b042c4b48d268b4f6165e7 Mon Sep 17 00:00:00 2001 From: "chiyoung.song" Date: Fri, 9 May 2025 21:08:49 +0900 Subject: [PATCH 2/3] clean --- python/packages/autogen-ext/tests/tools/test_mcp_tools.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/python/packages/autogen-ext/tests/tools/test_mcp_tools.py b/python/packages/autogen-ext/tests/tools/test_mcp_tools.py index 3fcea2636a3c..bb6c52841864 100644 --- a/python/packages/autogen-ext/tests/tools/test_mcp_tools.py +++ b/python/packages/autogen-ext/tests/tools/test_mcp_tools.py @@ -1,6 +1,6 @@ +import asyncio import logging import os -import asyncio from unittest.mock import AsyncMock, MagicMock import pytest @@ -597,5 +597,5 @@ async def test_lazy_init_and_finalize_cleanup() -> None: actor = workbench._actor # type: ignore[reportPrivateUsage] del workbench - await asyncio.sleep(0.1) - assert actor._active is False \ No newline at end of file + await asyncio.sleep(0.1) + assert actor._active is False From 0232b7de16e84eb0d768d2e029bcaa98c2cf7d27 Mon Sep 17 00:00:00 2001 From: "chiyoung.song" Date: Sat, 10 May 2025 21:51:59 +0900 Subject: [PATCH 3/3] FIX: improve codecov and fix logic for testcase --- .../src/autogen_ext/tools/mcp/_workbench.py | 18 +++---- .../autogen-ext/tests/tools/test_mcp_tools.py | 51 +++++++++++++++++++ 2 files changed, 59 insertions(+), 10 deletions(-) diff --git a/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_workbench.py b/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_workbench.py index 4aa1a48dfc74..fd43ff852f3d 100644 --- a/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_workbench.py +++ b/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_workbench.py @@ -153,6 +153,7 @@ def __init__(self, server_params: McpServerParams) -> None: self._server_params = server_params # self._session: ClientSession | None = None self._actor: McpSessionActor | None = None + self._actor_loop: asyncio.AbstractEventLoop | None = None self._read = None self._write = None @@ -254,6 +255,7 @@ async def start(self) -> None: if isinstance(self._server_params, (StdioServerParams, SseServerParams)): self._actor = McpSessionActor(self._server_params) await self._actor.initialize() + self._actor_loop = asyncio.get_event_loop() else: raise ValueError(f"Unsupported server params type: {type(self._server_params)}") @@ -283,14 +285,10 @@ def _from_config(cls, config: McpWorkbenchConfig) -> Self: def __del__(self) -> None: # Ensure the actor is stopped when the workbench is deleted - if self._actor: - try: - loop = asyncio.get_event_loop() - except RuntimeError: - loop = asyncio.new_event_loop() - asyncio.set_event_loop(loop) - - if loop.is_running(): - loop.create_task(self.stop()) + if self._actor and self._actor_loop: + loop = self._actor_loop + if loop.is_running() and not loop.is_closed(): + loop.call_soon_threadsafe(lambda: asyncio.create_task(self.stop())) else: - loop.run_until_complete(self.stop()) + msg = "Cannot safely stop actor at [McpWorkbench.__del__]: loop is closed or not running" + warnings.warn(msg, RuntimeWarning, stacklevel=2) diff --git a/python/packages/autogen-ext/tests/tools/test_mcp_tools.py b/python/packages/autogen-ext/tests/tools/test_mcp_tools.py index bb6c52841864..1230395a77f3 100644 --- a/python/packages/autogen-ext/tests/tools/test_mcp_tools.py +++ b/python/packages/autogen-ext/tests/tools/test_mcp_tools.py @@ -1,13 +1,17 @@ import asyncio import logging import os +import threading +from typing import cast from unittest.mock import AsyncMock, MagicMock import pytest +from _pytest.logging import LogCaptureFixture # type: ignore[import] from autogen_core import CancellationToken from autogen_core.tools import Workbench from autogen_core.utils import schema_to_pydantic_model from autogen_ext.tools.mcp import ( + McpSessionActor, McpWorkbench, SseMcpToolAdapter, SseServerParams, @@ -599,3 +603,50 @@ async def test_lazy_init_and_finalize_cleanup() -> None: del workbench await asyncio.sleep(0.1) assert actor._active is False + + +@pytest.mark.asyncio +async def test_del_to_new_event_loop_when_get_event_loop_fails() -> None: + params = StdioServerParams( + command="npx", + args=[ + "-y", + "@modelcontextprotocol/server-filesystem", + ".", + ], + read_timeout_seconds=60, + ) + workbench = McpWorkbench(server_params=params) + + await workbench.list_tools() + assert workbench._actor is not None # type: ignore[reportPrivateUsage] + assert workbench._actor._active is True # type: ignore[reportPrivateUsage] + + actor = workbench._actor # type: ignore[reportPrivateUsage] + + def cleanup() -> None: + nonlocal workbench + del workbench + + t = threading.Thread(target=cleanup) + t.start() + t.join() + + await asyncio.sleep(0.1) + assert actor._active is False # type: ignore[reportPrivateUsage] + + +def test_del_raises_when_loop_closed() -> None: + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + + params = StdioServerParams(command="echo", args=["ok"]) + workbench = McpWorkbench(server_params=params) + + workbench._actor_loop = loop # type: ignore[reportPrivateUsage] + workbench._actor = cast(McpSessionActor, object()) # type: ignore[reportPrivateUsage] + + loop.close() + + with pytest.warns(RuntimeWarning, match="loop is closed or not running"): + del workbench