diff --git a/src/agents/mcp/manager.py b/src/agents/mcp/manager.py index 733e744363..cb2151117c 100644 --- a/src/agents/mcp/manager.py +++ b/src/agents/mcp/manager.py @@ -5,7 +5,7 @@ from collections.abc import Awaitable, Callable, Iterable from contextlib import AbstractAsyncContextManager from dataclasses import dataclass -from typing import Any +from typing import Any, cast from ..logger import log_tool_action_debug, log_tool_action_error, logger from ._logging import get_mcp_server_log_message @@ -35,12 +35,77 @@ class _ServerCommand: class _ServerWorker: - def __init__(self, server: MCPServer) -> None: + def __init__( + self, + server: MCPServer, + cleanup_timeout_seconds_fn: Callable[[], float | None], + ) -> None: self._server = server + # The emergency cleanup below is a cleanup, not a connect, so it must be + # bounded by the manager cleanup timeout rather than by the interrupted + # command's own timeout. The callable is read at cancel time so runtime + # reassignment of the manager property stays honored. + self._cleanup_timeout_seconds_fn = cleanup_timeout_seconds_fn self._queue: asyncio.Queue[_ServerCommand] = asyncio.Queue() - self._task = asyncio.create_task(self._run()) + # The worker needs to tell an external cancel from a CancelledError a + # server raised on its own. Task.cancelling() covers that only on + # Python 3.11+, and this package still supports 3.10, so every cancel() + # call on the worker task is counted here. + self._cancel_requests = 0 + self._task: asyncio.Task[None] = self._create_task() self._cleanup_future: asyncio.Future[None] | None = None + def _create_task(self) -> asyncio.Task[None]: + # Route worker creation through the event loop's task factory so + # applications supervising or instrumenting tasks through + # set_task_factory still see this worker. The temporary factory keeps + # the app factory in charge of the task type and only adds the cancel + # counting on top, because loop shutdown and supervisors cancel the + # task object directly and would otherwise bypass the count. + loop = asyncio.get_running_loop() + app_factory = loop.get_task_factory() + + def factory( + loop: asyncio.AbstractEventLoop, coro: Any, **kwargs: Any + ) -> asyncio.Task[None]: + task: asyncio.Task[None] = ( + cast(asyncio.Task[None], app_factory(loop, coro, **kwargs)) + if app_factory + else cast(asyncio.Task[None], asyncio.Task(coro, loop=loop, **kwargs)) + ) + original_cancel = task.cancel + + def counting_cancel(msg: Any = None) -> bool: + self._cancel_requests += 1 + return bool(original_cancel(msg)) + + task.cancel = counting_cancel # type: ignore[method-assign] + return task + + # The typeshed _TaskFactory protocol is a private generic stub; this + # factory matches its runtime call shape exactly. + loop.set_task_factory(factory) # type: ignore[arg-type] + try: + return asyncio.create_task(self._run()) + finally: + loop.set_task_factory(app_factory) + + def _record_cleanup_failure(self, exc: BaseException) -> None: + """Record an emergency cleanup failure on the cleanup result future. + + The done callback reads the cleanup_error property through that + future, so a recorded failure keeps the worker registered instead of + letting the manager treat the server as cleanly stopped. A later + cleanup() call awaits the same settled future and surfaces the failure + to the manager, which keeps reconnect from starting a second + connection over a transport that was never successfully cleaned up. + """ + if self._cleanup_future is None: + self._cleanup_future = asyncio.get_running_loop().create_future() + if self._cleanup_future.done(): + return + self._cleanup_future.set_exception(exc) + @property def is_done(self) -> bool: return self._task.done() @@ -91,23 +156,67 @@ async def _submit(self, action: str, timeout_seconds: float | None) -> None: await future async def _run(self) -> None: - while True: - command = await self._queue.get() - should_exit = command.action == "cleanup" - try: - if command.action == "connect": - await _run_with_timeout_in_task(self._server.connect, command.timeout_seconds) - elif command.action == "cleanup": - await _run_with_timeout_in_task(self._server.cleanup, command.timeout_seconds) - else: - raise ValueError(f"Unknown command: {command.action}") - if not command.future.cancelled(): - command.future.set_result(None) - except BaseException as exc: - if not command.future.cancelled(): - command.future.set_exception(exc) - if should_exit: - return + try: + while True: + command = await self._queue.get() + should_exit = command.action == "cleanup" + cancellation_baseline = self._cancel_requests + try: + if command.action == "connect": + await _run_with_timeout_in_task( + self._server.connect, command.timeout_seconds + ) + elif command.action == "cleanup": + await _run_with_timeout_in_task( + self._server.cleanup, command.timeout_seconds + ) + else: + raise ValueError(f"Unknown command: {command.action}") + if not command.future.cancelled(): + command.future.set_result(None) + except BaseException as exc: + external_cancel = ( + isinstance(exc, asyncio.CancelledError) + and self._cancel_requests > cancellation_baseline + ) + if external_cancel and command.action != "cleanup": + # The worker task itself was cancelled while the command + # was running. A partially acquired connection from the + # interrupted command still needs cleanup, and this task + # owns it (some transports require cleanup in the same + # task), so clean up here before the caller is released. + # The manager cleanup timeout bounds the emergency + # cleanup, not the interrupted command's own timeout. + try: + await _run_with_timeout_in_task( + self._server.cleanup, self._cleanup_timeout_seconds_fn() + ) + except BaseException as cleanup_exc: + # A failed or timed-out emergency cleanup must not be + # swallowed silently. Recording it on the cleanup + # result future keeps the worker registered, so a + # later reconnect surfaces the failure and never + # starts a second connection over the transport + # that was never cleaned up. + self._record_cleanup_failure(cleanup_exc) + if not command.future.cancelled(): + command.future.set_exception(exc) + if external_cancel: + # The future carries the error to the caller; re-raise + # here to end the task and honor the cancellation. + raise + if should_exit: + return + finally: + # The worker loop has ended, so commands still queued will never run. + # Settle their futures so waiting callers do not hang forever. + while True: + try: + pending = self._queue.get_nowait() + except asyncio.QueueEmpty: + break + if not pending.future.done(): + pending.future.cancel() async def _run_with_timeout_in_task( @@ -535,7 +644,10 @@ async def _get_worker(self, server: MCPServer) -> _ServerWorker: self._discard_worker(server, worker) worker = self._workers.get(server) if worker is None: - worker = _ServerWorker(server=server) + worker = _ServerWorker( + server=server, + cleanup_timeout_seconds_fn=lambda: self.cleanup_timeout_seconds, + ) self._workers[server] = worker worker.add_done_callback(lambda _task: self._handle_worker_done(server, worker)) return worker diff --git a/tests/mcp/test_mcp_server_manager.py b/tests/mcp/test_mcp_server_manager.py index 40476867a6..c5376c58ec 100644 --- a/tests/mcp/test_mcp_server_manager.py +++ b/tests/mcp/test_mcp_server_manager.py @@ -1,6 +1,6 @@ import asyncio import logging -from collections.abc import Awaitable, Callable +from collections.abc import Awaitable, Callable, Coroutine from typing import Any, cast import pytest @@ -1280,3 +1280,264 @@ async def test_manager_restores_one_shot_iterable_servers_after_a_failed_connect # drop_failed_servers=False keeps failed servers active, so the restored list must match # what an equivalent list argument produces. assert manager.active_servers == [server] + + +@pytest.mark.asyncio +async def test_worker_task_cancellation_stops_worker_and_fails_caller() -> None: + # This test pins an exact cancellation boundary: a connect that hangs until the + # worker task itself is cancelled. The scripted utilities cannot model a worker + # task receiving an external cancel, so a local double is used here. + class HangingConnectServer: + def __init__(self) -> None: + self.connect_started = asyncio.Event() + + async def connect(self) -> None: + self.connect_started.set() + await asyncio.sleep(3600) + + async def cleanup(self) -> None: + return None + + server = HangingConnectServer() + worker = manager_module._ServerWorker(cast(Any, server), lambda: None) + + caller = asyncio.create_task(worker.connect(timeout_seconds=None)) + await asyncio.wait_for(server.connect_started.wait(), timeout=TEST_TIMEOUT_SECONDS) + + worker._task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(caller, timeout=TEST_TIMEOUT_SECONDS) + + assert worker._task.done() + assert worker._task.cancelled() + + +@pytest.mark.asyncio +async def test_worker_survives_server_raised_cancelled_error_and_serves_next_command() -> None: + # This test pins the CancelledServer contract: a server raising CancelledError on + # its own must reach the caller through the command future without ending the + # worker, because later lifecycle commands still rely on that worker. + class SelfCancellingThenOkServer: + def __init__(self) -> None: + self.connect_calls = 0 + + async def connect(self) -> None: + self.connect_calls += 1 + if self.connect_calls == 1: + raise asyncio.CancelledError() + + async def cleanup(self) -> None: + return None + + server = SelfCancellingThenOkServer() + worker = manager_module._ServerWorker(cast(Any, server), lambda: None) + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(worker.connect(timeout_seconds=None), timeout=TEST_TIMEOUT_SECONDS) + + assert not worker.is_done + + await asyncio.wait_for(worker.connect(timeout_seconds=None), timeout=TEST_TIMEOUT_SECONDS) + assert server.connect_calls == 2 + + await asyncio.wait_for(worker.cleanup(timeout_seconds=None), timeout=TEST_TIMEOUT_SECONDS) + assert worker.is_done + + +@pytest.mark.asyncio +async def test_worker_cancel_settles_commands_queued_behind_the_in_flight_one() -> None: + # An external cancel ends the worker, so commands still queued behind the + # in-flight one will never run. Their futures must be settled on the way out, + # or the callers awaiting those futures hang forever. + class BlockingConnectServer: + def __init__(self) -> None: + self.connect_started = asyncio.Event() + self.cleanup_calls = 0 + + async def connect(self) -> None: + self.connect_started.set() + await asyncio.sleep(3600) + + async def cleanup(self) -> None: + self.cleanup_calls += 1 + + server = BlockingConnectServer() + worker = manager_module._ServerWorker(cast(Any, server), lambda: None) + + connect_caller = asyncio.create_task(worker.connect(timeout_seconds=None)) + await asyncio.wait_for(server.connect_started.wait(), timeout=TEST_TIMEOUT_SECONDS) + cleanup_caller = asyncio.create_task(worker.cleanup(timeout_seconds=None)) + await asyncio.sleep(0) + + worker._task.cancel() + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(cleanup_caller, timeout=TEST_TIMEOUT_SECONDS) + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(connect_caller, timeout=TEST_TIMEOUT_SECONDS) + + assert worker._task.done() + assert worker._task.cancelled() + + +@pytest.mark.asyncio +async def test_worker_cancel_cleans_up_partially_acquired_connection() -> None: + # An external cancel during a blocking connect must not leak the partially + # acquired connection: the worker owns it and is the only task allowed to + # clean it up, so the cleanup has to run in that task before the worker + # terminates. + class BlockingPartialConnectServer(TaskBoundServer): + def __init__(self) -> None: + super().__init__() + self.connect_started = asyncio.Event() + self.cleanup_calls = 0 + + async def connect(self) -> None: + self._connect_task = asyncio.current_task() + self.connect_started.set() + await asyncio.sleep(3600) + + async def cleanup(self) -> None: + self.cleanup_calls += 1 + await super().cleanup() + + server = BlockingPartialConnectServer() + manager = MCPServerManager([server], connect_in_parallel=True) + + connect_all = asyncio.create_task(manager.connect_all()) + await asyncio.wait_for(server.connect_started.wait(), timeout=TEST_TIMEOUT_SECONDS) + manager._workers[server]._task.cancel() + await asyncio.wait_for(connect_all, timeout=TEST_TIMEOUT_SECONDS) + + assert server.cleaned is True + assert server.cleanup_calls == 1 + await asyncio.wait_for(manager.cleanup_all(), timeout=TEST_TIMEOUT_SECONDS) + assert server.cleanup_calls == 1 + + +@pytest.mark.asyncio +async def test_worker_emergency_cleanup_uses_the_cleanup_timeout() -> None: + # An external cancel during a connect must bound the emergency cleanup by + # the manager cleanup timeout, not by the interrupted connect command's own + # timeout: when connect_timeout_seconds is None the command carries no + # bound, and an unbounded emergency cleanup can hang the worker (and loop + # shutdown) forever, while a shorter connect timeout could cut a legitimate + # cleanup short. + class HangingCleanupServer: + def __init__(self) -> None: + self.connect_started = asyncio.Event() + self.cleanup_started = asyncio.Event() + self.cleanup_cancelled = False + + async def connect(self) -> None: + self.connect_started.set() + await asyncio.sleep(3600) + + async def cleanup(self) -> None: + self.cleanup_started.set() + try: + await asyncio.sleep(3600) + except asyncio.CancelledError: + self.cleanup_cancelled = True + raise + + server = HangingCleanupServer() + worker = manager_module._ServerWorker(cast(Any, server), lambda: 0.2) + + caller = asyncio.create_task(worker.connect(timeout_seconds=None)) + await asyncio.wait_for(server.connect_started.wait(), timeout=TEST_TIMEOUT_SECONDS) + worker._task.cancel() + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(caller, timeout=TEST_TIMEOUT_SECONDS) + # the caller settles only after the emergency cleanup is bounded and done, + # so the worker must already be terminated here; awaiting a cancelled task + # would re-raise CancelledError in the test task instead + assert worker._task.done() + assert worker._task.cancelled() + assert server.cleanup_started.is_set() + assert server.cleanup_cancelled is True + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ["raise", "timeout"]) +async def test_manager_reconnect_does_not_retry_after_failed_emergency_cleanup( + mode: str, +) -> None: + # A manager-level regression for the emergency cleanup result: an external + # cancel during a partially started connect runs the emergency cleanup in + # the worker task, and when that cleanup raises or times out the failure + # must reach the existing cleanup result mechanism. The done callback then + # keeps the worker instead of discarding it, a later reconnect surfaces the + # failure in manager errors, and the still-uncleaned server is never + # started again as a second connection. + class EmergencyCleanupFailingServer(TaskBoundServer): + def __init__(self, mode: str) -> None: + super().__init__() + self.connect_started = asyncio.Event() + self.mode = mode + self.connect_calls = 0 + self.cleanup_calls = 0 + + async def connect(self) -> None: + self._connect_task = asyncio.current_task() + self.connect_calls += 1 + self.connect_started.set() + await asyncio.sleep(3600) + + async def cleanup(self) -> None: + self.cleanup_calls += 1 + await super().cleanup() + if self.mode == "raise": + raise RuntimeError("emergency cleanup failed") + await asyncio.sleep(3600) + + server = EmergencyCleanupFailingServer(mode) + manager = MCPServerManager([server], connect_in_parallel=True, cleanup_timeout_seconds=0.2) + + connect_all = asyncio.create_task(manager.connect_all()) + await asyncio.wait_for(server.connect_started.wait(), timeout=TEST_TIMEOUT_SECONDS) + manager._workers[server]._task.cancel() + await asyncio.wait_for(connect_all, timeout=TEST_TIMEOUT_SECONDS) + + assert server.connect_calls == 1 + assert server.cleanup_calls == 1 + + await manager.reconnect() + + assert server.connect_calls == 1 + assert manager.failed_servers == [server] + assert manager.active_servers == [] + if mode == "raise": + assert str(manager.errors[server]) == "emergency cleanup failed" + else: + assert isinstance(manager.errors[server], asyncio.TimeoutError) + worker = manager._workers[server] + assert worker.is_done + assert worker.cleanup_error is not None + + +@pytest.mark.asyncio +async def test_worker_task_is_created_through_the_loop_task_factory() -> None: + # Worker creation must honor the event loop's task factory so applications + # supervising or instrumenting tasks through set_task_factory can discover + # the connection worker. + loop = asyncio.get_running_loop() + created: list[asyncio.Task[object]] = [] + + def factory( + loop_: asyncio.AbstractEventLoop, coro: Coroutine[Any, Any, object], **kwargs: object + ) -> asyncio.Task[object]: + task = asyncio.Task(coro, loop=loop_, **kwargs) # type: ignore[arg-type] + created.append(task) + return task + + loop.set_task_factory(factory) + try: + server = TaskBoundServer() + manager = MCPServerManager([server], connect_in_parallel=True) + await asyncio.wait_for(manager.connect_all(), timeout=TEST_TIMEOUT_SECONDS) + worker = manager._workers[server] + assert worker._task in created + finally: + loop.set_task_factory(None)