Skip to content

Commit c26d894

Browse files
authored
fix/mcp_session_auto_close_when_Mcpworkbench_deleted (#6497)
<!-- Thank you for your contribution! Please review https://microsoft.github.io/autogen/docs/Contribute before opening a pull request. --> <!-- Please add a reviewer to the assignee section when you create a PR. If you don't have the access to it, we will shortly find a reviewer and assign them to your PR. --> ## Why are these changes needed? As-is: Deleting `McpWorkbench` does not close the `McpSession`. To-be: Deleting `McpWorkbench` now properly closes the `McpSession`. <!-- Please give a short summary of the change and the problem this solves. --> ## Related issue number <!-- For example: "Closes #1234" --> ## Checks - [ ] I've included any doc changes needed for <https://microsoft.github.io/autogen/>. See <https://github.com/microsoft/autogen/blob/main/CONTRIBUTING.md> to build and test documentation locally. - [x] I've added tests (if relevant) corresponding to the changes introduced in this PR. - [x] I've made sure all auto checks have passed.
1 parent 6427c07 commit c26d894

File tree

2 files changed

+65
-1
lines changed

2 files changed

+65
-1
lines changed

python/packages/autogen-ext/src/autogen_ext/tools/mcp/_workbench.py

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import asyncio
12
import builtins
23
import warnings
34
from typing import Any, List, Literal, Mapping
@@ -152,6 +153,7 @@ def __init__(self, server_params: McpServerParams) -> None:
152153
self._server_params = server_params
153154
# self._session: ClientSession | None = None
154155
self._actor: McpSessionActor | None = None
156+
self._actor_loop: asyncio.AbstractEventLoop | None = None
155157
self._read = None
156158
self._write = None
157159

@@ -253,6 +255,7 @@ async def start(self) -> None:
253255
if isinstance(self._server_params, (StdioServerParams, SseServerParams)):
254256
self._actor = McpSessionActor(self._server_params)
255257
await self._actor.initialize()
258+
self._actor_loop = asyncio.get_event_loop()
256259
else:
257260
raise ValueError(f"Unsupported server params type: {type(self._server_params)}")
258261

@@ -282,4 +285,10 @@ def _from_config(cls, config: McpWorkbenchConfig) -> Self:
282285

283286
def __del__(self) -> None:
284287
# Ensure the actor is stopped when the workbench is deleted
285-
pass
288+
if self._actor and self._actor_loop:
289+
loop = self._actor_loop
290+
if loop.is_running() and not loop.is_closed():
291+
loop.call_soon_threadsafe(lambda: asyncio.create_task(self.stop()))
292+
else:
293+
msg = "Cannot safely stop actor at [McpWorkbench.__del__]: loop is closed or not running"
294+
warnings.warn(msg, RuntimeWarning, stacklevel=2)

python/packages/autogen-ext/tests/tools/test_mcp_tools.py

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,17 @@
1+
import asyncio
12
import logging
23
import os
4+
import threading
5+
from typing import cast
36
from unittest.mock import AsyncMock, MagicMock
47

58
import pytest
9+
from _pytest.logging import LogCaptureFixture # type: ignore[import]
610
from autogen_core import CancellationToken
711
from autogen_core.tools import Workbench
812
from autogen_core.utils import schema_to_pydantic_model
913
from autogen_ext.tools.mcp import (
14+
McpSessionActor,
1015
McpWorkbench,
1116
SseMcpToolAdapter,
1217
SseServerParams,
@@ -594,4 +599,54 @@ async def test_lazy_init_and_finalize_cleanup() -> None:
594599
assert workbench._actor is not None # type: ignore[reportPrivateUsage]
595600
assert workbench._actor._active is True # type: ignore[reportPrivateUsage]
596601

602+
actor = workbench._actor # type: ignore[reportPrivateUsage]
597603
del workbench
604+
await asyncio.sleep(0.1)
605+
assert actor._active is False
606+
607+
608+
@pytest.mark.asyncio
609+
async def test_del_to_new_event_loop_when_get_event_loop_fails() -> None:
610+
params = StdioServerParams(
611+
command="npx",
612+
args=[
613+
"-y",
614+
"@modelcontextprotocol/server-filesystem",
615+
".",
616+
],
617+
read_timeout_seconds=60,
618+
)
619+
workbench = McpWorkbench(server_params=params)
620+
621+
await workbench.list_tools()
622+
assert workbench._actor is not None # type: ignore[reportPrivateUsage]
623+
assert workbench._actor._active is True # type: ignore[reportPrivateUsage]
624+
625+
actor = workbench._actor # type: ignore[reportPrivateUsage]
626+
627+
def cleanup() -> None:
628+
nonlocal workbench
629+
del workbench
630+
631+
t = threading.Thread(target=cleanup)
632+
t.start()
633+
t.join()
634+
635+
await asyncio.sleep(0.1)
636+
assert actor._active is False # type: ignore[reportPrivateUsage]
637+
638+
639+
def test_del_raises_when_loop_closed() -> None:
640+
loop = asyncio.new_event_loop()
641+
asyncio.set_event_loop(loop)
642+
643+
params = StdioServerParams(command="echo", args=["ok"])
644+
workbench = McpWorkbench(server_params=params)
645+
646+
workbench._actor_loop = loop # type: ignore[reportPrivateUsage]
647+
workbench._actor = cast(McpSessionActor, object()) # type: ignore[reportPrivateUsage]
648+
649+
loop.close()
650+
651+
with pytest.warns(RuntimeWarning, match="loop is closed or not running"):
652+
del workbench

0 commit comments

Comments
 (0)