diff --git a/tests/tools/test_mcp_stability.py b/tests/tools/test_mcp_stability.py index b843439fff7..712a80c04eb 100644 --- a/tests/tools/test_mcp_stability.py +++ b/tests/tools/test_mcp_stability.py @@ -795,7 +795,23 @@ class TestMCPLoopDrainOnStop: time.sleep(0.01) assert state["started"], "task never started on the MCP loop" - mcp_mod._stop_mcp_loop() + stop_saw_cleanup = [] + call_soon_threadsafe = loop.call_soon_threadsafe + + def record_stop_order(callback, *args, **kwargs): + if ( + getattr(callback, "__self__", None) is loop + and getattr(callback, "__name__", None) == "stop" + ): + stop_saw_cleanup.append(state["cleanup_ran"]) + return call_soon_threadsafe(callback, *args, **kwargs) + + with patch.object( + loop, + "call_soon_threadsafe", + side_effect=record_stop_order, + ): + mcp_mod._stop_mcp_loop() assert state["task"] is not None assert state["task"].done(), "task left pending when the loop closed" @@ -803,8 +819,79 @@ class TestMCPLoopDrainOnStop: f"cleanup ran against a closed loop: {state['cleanup_error']!r}" ) assert state["cleanup_ran"], "task cleanup never ran" + assert stop_saw_cleanup == [True], "loop.stop ran before task cleanup" finally: with mcp_mod._lock: mcp_mod._servers.clear() mcp_mod._server_connecting.clear() mcp_mod._stop_mcp_loop() + + def test_drain_is_bounded_when_task_ignores_cancellation(self, caplog): + """A cancellation-resistant task must not hang final MCP shutdown.""" + import tools.mcp_tool as mcp_mod + + async def _run(): + release = asyncio.Event() + + async def cancellation_resistant(): + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + await release.wait() + + task = asyncio.create_task(cancellation_resistant()) + await asyncio.sleep(0) + try: + with caplog.at_level("WARNING", logger=mcp_mod.logger.name): + await mcp_mod._drain_mcp_loop_tasks(timeout=0.01) + assert not task.done(), "drain waited indefinitely for resistant task" + finally: + release.set() + if not task.done() and task.cancelling() == 0: + task.cancel() + await task + + asyncio.run(_run()) + + assert any( + "still pending after" in record.getMessage() + for record in caplog.records + ) + + def test_stop_warns_when_drain_wait_times_out(self, caplog): + """A failed final drain must be visible instead of silently no-oping.""" + import tools.mcp_tool as mcp_mod + + class TimedOutFuture: + def result(self, timeout): + assert timeout > 0 + raise TimeoutError("simulated drain timeout") + + def cancel(self): + return True + + def report_timeout(coro, _loop, **_kwargs): + coro.close() + return TimedOutFuture() + + with mcp_mod._lock: + mcp_mod._servers.clear() + mcp_mod._server_connecting.clear() + mcp_mod._ensure_mcp_loop() + try: + with caplog.at_level("WARNING", logger=mcp_mod.logger.name): + with patch( + "agent.async_utils.safe_schedule_threadsafe", + side_effect=report_timeout, + ): + mcp_mod._stop_mcp_loop() + finally: + with mcp_mod._lock: + mcp_mod._servers.clear() + mcp_mod._server_connecting.clear() + mcp_mod._stop_mcp_loop() + + assert any( + "Timed out waiting for MCP loop drain" in record.getMessage() + for record in caplog.records + ) diff --git a/tests/tools/test_mcp_tool.py b/tests/tools/test_mcp_tool.py index cc0c3e4089f..35eb365c9ad 100644 --- a/tests/tools/test_mcp_tool.py +++ b/tests/tools/test_mcp_tool.py @@ -1809,6 +1809,7 @@ class TestShutdown: shutdown_started = threading.Event() parked_task_done = threading.Event() scheduled_shutdown = {} + schedule_count = 0 class StalledShutdownServer(MCPServerTask): async def shutdown(self): @@ -1838,7 +1839,12 @@ class TestShutdown: mcp_mod._servers[server.name] = server def schedule_then_report_timeout(coro, target_loop, **_kwargs): + nonlocal schedule_count + schedule_count += 1 future = asyncio.run_coroutine_threadsafe(coro, target_loop) + if schedule_count > 1: + return future + scheduled_shutdown["future"] = future class TimedOutFuture: @@ -1862,6 +1868,7 @@ class TestShutdown: ) assert parked_task.done() assert scheduled_shutdown["future"].done() + assert schedule_count == 2 finally: with mcp_mod._lock: mcp_mod._servers.clear() diff --git a/tools/mcp_tool.py b/tools/mcp_tool.py index bc5a8d7f273..a5eb720a24a 100644 --- a/tools/mcp_tool.py +++ b/tools/mcp_tool.py @@ -363,6 +363,11 @@ def _jittered(seconds: float) -> float: _DEFAULT_KEEPALIVE_INTERVAL = 180 # seconds between liveness pings _MIN_KEEPALIVE_INTERVAL = 5 # clamp floor for configured intervals +# Final shutdown gives pending MCP-loop tasks one bounded cancellation cycle +# before closing their owning loop. Cooperative parked/reconnect waiters finish +# immediately; cancellation-resistant tasks must not hang process exit. +_MCP_LOOP_DRAIN_TIMEOUT = 3.0 + # Environment variables that are safe to pass to stdio subprocesses _SAFE_ENV_KEYS = frozenset({ "PATH", "HOME", "USER", "LANG", "LC_ALL", "TERM", "SHELL", "TMPDIR", @@ -6618,12 +6623,16 @@ def _stop_mcp_loop_if_idle() -> bool: return _stop_mcp_loop(only_if_idle=True) -async def _drain_mcp_loop_tasks() -> None: +async def _drain_mcp_loop_tasks( + *, + timeout: float = _MCP_LOOP_DRAIN_TIMEOUT, +) -> None: """Cancel every task still pending on the MCP loop and reap it. Cancelling is not enough on its own: ``Task.cancel()`` only schedules the - throw, so the task must be awaited before the loop goes away. Gather them - here — while the loop is still open — so each one runs its own ``finally``. + throw, so tasks need a cancellation cycle before the loop goes away. Wait + for them here — on their owning loop — but keep the final drain bounded so + a task that suppresses cancellation cannot hang process exit indefinitely. """ current = asyncio.current_task() pending = [t for t in asyncio.all_tasks() if t is not current and not t.done()] @@ -6632,7 +6641,23 @@ async def _drain_mcp_loop_tasks() -> None: logger.debug("Draining %d pending task(s) from the MCP loop", len(pending)) for task in pending: task.cancel() - await asyncio.gather(*pending, return_exceptions=True) + + done, still_pending = await asyncio.wait(pending, timeout=timeout) + for task in done: + if task.cancelled(): + continue + try: + task.exception() + except asyncio.CancelledError: + pass + except Exception as exc: + logger.debug("Pending MCP loop task ended during shutdown: %s", exc) + + if still_pending: + logger.warning( + "%d MCP loop task(s) still pending after %.1fs drain", + len(still_pending), timeout, + ) def _stop_mcp_loop(*, only_if_idle: bool = False) -> bool: @@ -6656,22 +6681,31 @@ def _stop_mcp_loop(*, only_if_idle: bool = False) -> bool: from agent.async_utils import safe_schedule_threadsafe future = safe_schedule_threadsafe( - _drain_mcp_loop_tasks(), loop, + _drain_mcp_loop_tasks(timeout=_MCP_LOOP_DRAIN_TIMEOUT), loop, logger=logger, log_message="MCP loop drain: failed to schedule", + log_level=logging.WARNING, ) if future is not None: try: - future.result(timeout=5) + future.result(timeout=_MCP_LOOP_DRAIN_TIMEOUT + 1) + except TimeoutError: + future.cancel() + logger.warning( + "Timed out waiting for MCP loop drain after %.1fs", + _MCP_LOOP_DRAIN_TIMEOUT + 1, + ) except BaseException as exc: - logger.debug("Error draining MCP loop tasks: %s", exc) + logger.warning("Error draining MCP loop tasks: %s", exc) loop.call_soon_threadsafe(loop.stop) if thread is not None: thread.join(timeout=5) + if thread.is_alive(): + logger.warning("MCP event loop thread did not stop within 5.0s") try: loop.close() - except Exception: - pass + except Exception as exc: + logger.warning("Unable to close MCP event loop cleanly: %s", exc) # After closing the loop, any stdio subprocesses that survived the # graceful shutdown are now orphaned — include active PIDs too # since the loop is gone and no session can still be in flight.