diff --git a/tests/tools/test_mcp_tool.py b/tests/tools/test_mcp_tool.py index a2cf61e04a5..cc0c3e4089f 100644 --- a/tests/tools/test_mcp_tool.py +++ b/tests/tools/test_mcp_tool.py @@ -1795,6 +1795,79 @@ class TestShutdown: assert len(_servers) == 0 mock_server.shutdown.assert_called_once() + def test_shutdown_drains_parked_server_after_bounded_wait_expires(self): + """The public shutdown path drains a parked server if graceful shutdown stalls. + + This exercises the production ownership path: a real ``MCPServerTask`` + is registered, its parked waiter owns child tasks on the shared loop, + and the bounded wait for the scheduled shutdown expires. The loop owner + must still cancel and drain that waiter before closing the loop. + """ + import tools.mcp_tool as mcp_mod + from tools.mcp_tool import MCPServerTask, shutdown_mcp_servers + + shutdown_started = threading.Event() + parked_task_done = threading.Event() + scheduled_shutdown = {} + + class StalledShutdownServer(MCPServerTask): + async def shutdown(self): + shutdown_started.set() + await asyncio.Event().wait() + + server = StalledShutdownServer("parked") + with mcp_mod._lock: + mcp_mod._servers.clear() + mcp_mod._server_connecting.clear() + mcp_mod._ensure_mcp_loop() + with mcp_mod._lock: + loop = mcp_mod._mcp_loop + assert loop is not None + + async def install_parked_waiter(): + task = asyncio.create_task(server._wait_for_reconnect_or_shutdown()) + server._task = task + task.add_done_callback(lambda _task: parked_task_done.set()) + await asyncio.sleep(0) + return task + + parked_task = asyncio.run_coroutine_threadsafe( + install_parked_waiter(), loop + ).result(timeout=2) + with mcp_mod._lock: + mcp_mod._servers[server.name] = server + + def schedule_then_report_timeout(coro, target_loop, **_kwargs): + future = asyncio.run_coroutine_threadsafe(coro, target_loop) + scheduled_shutdown["future"] = future + + class TimedOutFuture: + def result(self, timeout): + assert timeout > 0 + assert shutdown_started.wait(timeout=2) + raise TimeoutError("simulated bounded MCP shutdown timeout") + + return TimedOutFuture() + + try: + with patch( + "agent.async_utils.safe_schedule_threadsafe", + side_effect=schedule_then_report_timeout, + ): + shutdown_mcp_servers() + + assert loop.is_closed() + assert parked_task_done.is_set(), ( + "parked MCPServerTask was not drained before its loop closed" + ) + assert parked_task.done() + assert scheduled_shutdown["future"].done() + finally: + with mcp_mod._lock: + mcp_mod._servers.clear() + mcp_mod._server_connecting.clear() + mcp_mod._stop_mcp_loop() + def test_shutdown_deregisters_registered_tools(self): """shutdown_mcp_servers removes MCP tools and their raw alias.""" import tools.mcp_tool as mcp_mod