diff --git a/tests/tools/test_mcp_tool.py b/tests/tools/test_mcp_tool.py index 663895c0b..883bbe318 100644 --- a/tests/tools/test_mcp_tool.py +++ b/tests/tools/test_mcp_tool.py @@ -6,6 +6,8 @@ All tests use mocks -- no real MCP servers or subprocesses are started. import asyncio import json import os +import threading +import time from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch @@ -255,6 +257,77 @@ class TestToolHandler: finally: _servers.pop("test_srv", None) + def test_interrupted_call_returns_interrupted_error(self): + from tools.mcp_tool import _make_tool_handler, _servers + + mock_session = MagicMock() + server = _make_mock_server("test_srv", session=mock_session) + _servers["test_srv"] = server + + try: + handler = _make_tool_handler("test_srv", "greet", 120) + def _interrupting_run(coro, timeout=30): + coro.close() + raise InterruptedError("User sent a new message") + with patch( + "tools.mcp_tool._run_on_mcp_loop", + side_effect=_interrupting_run, + ): + result = json.loads(handler({})) + assert result == {"error": "MCP call interrupted: user sent a new message"} + finally: + _servers.pop("test_srv", None) + + +class TestRunOnMCPLoopInterrupts: + def test_interrupt_cancels_waiting_mcp_call(self): + import tools.mcp_tool as mcp_mod + from tools.interrupt import set_interrupt + + loop = asyncio.new_event_loop() + thread = threading.Thread(target=loop.run_forever, daemon=True) + thread.start() + + cancelled = threading.Event() + + async def _slow_call(): + try: + await asyncio.sleep(5) + return "done" + except asyncio.CancelledError: + cancelled.set() + raise + + old_loop = mcp_mod._mcp_loop + old_thread = mcp_mod._mcp_thread + mcp_mod._mcp_loop = loop + mcp_mod._mcp_thread = thread + + waiter_tid = threading.current_thread().ident + + def _interrupt_soon(): + time.sleep(0.2) + set_interrupt(True, waiter_tid) + + interrupter = threading.Thread(target=_interrupt_soon, daemon=True) + interrupter.start() + + try: + with pytest.raises(InterruptedError, match="User sent a new message"): + mcp_mod._run_on_mcp_loop(_slow_call(), timeout=2) + + deadline = time.time() + 2 + while time.time() < deadline and not cancelled.is_set(): + time.sleep(0.05) + assert cancelled.is_set() + finally: + set_interrupt(False, waiter_tid) + loop.call_soon_threadsafe(loop.stop) + thread.join(timeout=2) + loop.close() + mcp_mod._mcp_loop = old_loop + mcp_mod._mcp_thread = old_thread + # --------------------------------------------------------------------------- # Tool registration (discovery + register) diff --git a/tools/mcp_tool.py b/tools/mcp_tool.py index e953998cc..2356830c4 100644 --- a/tools/mcp_tool.py +++ b/tools/mcp_tool.py @@ -70,6 +70,7 @@ Thread safety: """ import asyncio +import concurrent.futures import inspect import json import logging @@ -1167,13 +1168,43 @@ def _ensure_mcp_loop(): def _run_on_mcp_loop(coro, timeout: float = 30): - """Schedule a coroutine on the MCP event loop and block until done.""" + """Schedule a coroutine on the MCP event loop and block until done. + + Poll in short intervals so the calling agent thread can honor user + interrupts while the MCP work is still running on the background loop. + """ + from tools.interrupt import is_interrupted + with _lock: loop = _mcp_loop if loop is None or not loop.is_running(): raise RuntimeError("MCP event loop is not running") future = asyncio.run_coroutine_threadsafe(coro, loop) - return future.result(timeout=timeout) + deadline = None if timeout is None else time.monotonic() + timeout + + while True: + if is_interrupted(): + future.cancel() + raise InterruptedError("User sent a new message") + + wait_timeout = 0.1 + if deadline is not None: + remaining = deadline - time.monotonic() + if remaining <= 0: + return future.result(timeout=0) + wait_timeout = min(wait_timeout, remaining) + + try: + return future.result(timeout=wait_timeout) + except concurrent.futures.TimeoutError: + continue + + +def _interrupted_call_result() -> str: + """Standardized JSON error for a user-interrupted MCP tool call.""" + return json.dumps({ + "error": "MCP call interrupted: user sent a new message" + }) # --------------------------------------------------------------------------- @@ -1299,6 +1330,8 @@ def _make_tool_handler(server_name: str, tool_name: str, tool_timeout: float): try: return _run_on_mcp_loop(_call(), timeout=tool_timeout) + except InterruptedError: + return _interrupted_call_result() except Exception as exc: logger.error( "MCP tool %s/%s call failed: %s", @@ -1342,6 +1375,8 @@ def _make_list_resources_handler(server_name: str, tool_timeout: float): try: return _run_on_mcp_loop(_call(), timeout=tool_timeout) + except InterruptedError: + return _interrupted_call_result() except Exception as exc: logger.error( "MCP %s/list_resources failed: %s", server_name, exc, @@ -1386,6 +1421,8 @@ def _make_read_resource_handler(server_name: str, tool_timeout: float): try: return _run_on_mcp_loop(_call(), timeout=tool_timeout) + except InterruptedError: + return _interrupted_call_result() except Exception as exc: logger.error( "MCP %s/read_resource failed: %s", server_name, exc, @@ -1433,6 +1470,8 @@ def _make_list_prompts_handler(server_name: str, tool_timeout: float): try: return _run_on_mcp_loop(_call(), timeout=tool_timeout) + except InterruptedError: + return _interrupted_call_result() except Exception as exc: logger.error( "MCP %s/list_prompts failed: %s", server_name, exc, @@ -1488,6 +1527,8 @@ def _make_get_prompt_handler(server_name: str, tool_timeout: float): try: return _run_on_mcp_loop(_call(), timeout=tool_timeout) + except InterruptedError: + return _interrupted_call_result() except Exception as exc: logger.error( "MCP %s/get_prompt failed: %s", server_name, exc,