mirror of
https://github.com/microsoft/autogen.git
synced 2025-12-13 07:51:21 +00:00
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:
parent
6427c07f5c
commit
c26d894c34
@ -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)
|
||||
|
||||
@ -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
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user