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.
This commit is contained in:
EeS 2025-05-12 22:31:01 +09:00 committed by GitHub
parent 6427c07f5c
commit c26d894c34
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 65 additions and 1 deletions

View File

@ -1,3 +1,4 @@
import asyncio
import builtins
import warnings
from typing import Any, List, Literal, Mapping
@ -152,6 +153,7 @@ class McpWorkbench(Workbench, Component[McpWorkbenchConfig]):
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
@ -253,6 +255,7 @@ class McpWorkbench(Workbench, Component[McpWorkbenchConfig]):
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)}")
@ -282,4 +285,10 @@ class McpWorkbench(Workbench, Component[McpWorkbenchConfig]):
def __del__(self) -> None:
# Ensure the actor is stopped when the workbench is deleted
pass
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:
msg = "Cannot safely stop actor at [McpWorkbench.__del__]: loop is closed or not running"
warnings.warn(msg, RuntimeWarning, stacklevel=2)

View File

@ -1,12 +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,
@ -594,4 +599,54 @@ 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
@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