mirror of
https://github.com/microsoft/autogen.git
synced 2025-12-17 18:18:54 +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 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)
|
||||||
|
|||||||
@ -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
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user