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 builtins
import warnings import warnings
from typing import Any, List, Literal, Mapping from typing import Any, List, Literal, Mapping
@ -152,6 +153,7 @@ class McpWorkbench(Workbench, Component[McpWorkbenchConfig]):
self._server_params = server_params self._server_params = server_params
# self._session: ClientSession | None = None # self._session: ClientSession | None = None
self._actor: McpSessionActor | None = None self._actor: McpSessionActor | None = None
self._actor_loop: asyncio.AbstractEventLoop | None = None
self._read = None self._read = None
self._write = None self._write = None
@ -253,6 +255,7 @@ class McpWorkbench(Workbench, Component[McpWorkbenchConfig]):
if isinstance(self._server_params, (StdioServerParams, SseServerParams)): if isinstance(self._server_params, (StdioServerParams, SseServerParams)):
self._actor = McpSessionActor(self._server_params) self._actor = McpSessionActor(self._server_params)
await self._actor.initialize() await self._actor.initialize()
self._actor_loop = asyncio.get_event_loop()
else: else:
raise ValueError(f"Unsupported server params type: {type(self._server_params)}") raise ValueError(f"Unsupported server params type: {type(self._server_params)}")
@ -282,4 +285,10 @@ class McpWorkbench(Workbench, Component[McpWorkbenchConfig]):
def __del__(self) -> None: def __del__(self) -> None:
# Ensure the actor is stopped when the workbench is deleted # 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 logging
import os import os
import threading
from typing import cast
from unittest.mock import AsyncMock, MagicMock from unittest.mock import AsyncMock, MagicMock
import pytest import pytest
from _pytest.logging import LogCaptureFixture # type: ignore[import]
from autogen_core import CancellationToken from autogen_core import CancellationToken
from autogen_core.tools import Workbench from autogen_core.tools import Workbench
from autogen_core.utils import schema_to_pydantic_model from autogen_core.utils import schema_to_pydantic_model
from autogen_ext.tools.mcp import ( from autogen_ext.tools.mcp import (
McpSessionActor,
McpWorkbench, McpWorkbench,
SseMcpToolAdapter, SseMcpToolAdapter,
SseServerParams, 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 is not None # type: ignore[reportPrivateUsage]
assert workbench._actor._active is True # type: ignore[reportPrivateUsage] assert workbench._actor._active is True # type: ignore[reportPrivateUsage]
actor = workbench._actor # type: ignore[reportPrivateUsage]
del workbench 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