Merge pull request #67183 from NousResearch/bb/mcp-poll-loop-oom

fix(mcp): stop gateway OOM from poll loop swallowing a completed future's real TimeoutError (supersedes #63918, #64072, #63903, #66039)

tests/tools/test_mcp_poll_loop_oom_integration.py

@@ -0,0 +1,68 @@
+"""End-to-end coverage for the MCP poll-loop OOM spin (#63892).
+
+The unit tests in ``test_mcp_tool.py`` hand-construct completed futures to lock
+the ``_run_on_mcp_loop`` contract deterministically. This complements them by
+exercising the actual field trigger through the *live* MCP loop: an inner
+``asyncio.wait_for`` expiry produces a real ``TimeoutError`` stored on a real
+future scheduled via ``run_coroutine_threadsafe``. On Python >= 3.8 that class
+is the builtin ``TimeoutError``, so before the fix the poll loop swallowed the
+completed future's stored exception and spun with no sleep -- burning CPU,
+growing the exception traceback ~108 MB/s until the gateway OOM'd, and finally
+masking the real error behind the generic "MCP call timed out after <full
+timeout>" wrapper.
+
+The fixed loop must surface the real exception once, promptly -- long before
+the outer deadline.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import time
+
+import pytest
+
+
+@pytest.fixture
+def mcp_loop():
+    import tools.mcp_tool as mcp_tool
+
+    mcp_tool._ensure_mcp_loop()
+    yield mcp_tool
+    mcp_tool._stop_mcp_loop()
+
+
+def test_inner_wait_for_timeout_surfaces_promptly_without_spinning(mcp_loop):
+    async def inner():
+        # Real TimeoutError from wait_for, completing far before the outer
+        # deadline below -- the exact shape of an MCP call_tool wrapped in the
+        # server's configured mcp_servers.<srv>.timeout.
+        await asyncio.wait_for(asyncio.sleep(60), timeout=0.05)
+
+    start = time.monotonic()
+    with pytest.raises(TimeoutError) as exc:
+        mcp_loop._run_on_mcp_loop(inner, timeout=10)
+    elapsed = time.monotonic() - start
+
+    # Fixed: the real inner TimeoutError surfaces within a couple poll ticks.
+    # Broken: the loop swallows it and spins with no sleep until the 10s
+    # deadline, then raises the generic wrapper message instead. A generous
+    # (>= 2s) bound keeps this stable on a slow CI runner while still failing
+    # hard on a spin-to-deadline regression.
+    assert elapsed < 5.0, (
+        f"poll loop spun for {elapsed:.1f}s instead of resolving the completed "
+        "future once (#63892 regression)"
+    )
+    assert "MCP call timed out after" not in str(exc.value), (
+        "the real inner TimeoutError must surface, not the outer poll deadline"
+    )
+
+
+def test_successful_call_still_returns_through_real_loop(mcp_loop):
+    """Guard the done-future path against regressing normal success returns."""
+
+    async def inner():
+        await asyncio.sleep(0.25)  # first polls time out while still pending
+        return {"ok": True}
+
+    assert mcp_loop._run_on_mcp_loop(inner, timeout=10) == {"ok": True}

tests/tools/test_mcp_tool.py

@@ -4,6 +4,7 @@ All tests use mocks -- no real MCP servers or subprocesses are started.
 """
 
 import asyncio
+import concurrent.futures
 import json
 import threading
 import time
@@ -837,6 +838,23 @@ class TestToolHandler:
 
 
 class TestRunOnMCPLoopInterrupts:
+    @staticmethod
+    def _run_with_future(mcp_mod, future):
+        loop = MagicMock()
+        loop.is_running.return_value = True
+
+        async def _unused_call():
+            return "unused"
+
+        def _schedule(coro, scheduled_loop, **_kwargs):
+            assert scheduled_loop is loop
+            coro.close()
+            return future
+
+        with patch.object(mcp_mod, "_mcp_loop", loop):
+            with patch("agent.async_utils.safe_schedule_threadsafe", side_effect=_schedule):
+                return mcp_mod._run_on_mcp_loop(_unused_call(), timeout=1)
+
     def test_interrupt_cancels_waiting_mcp_call(self):
         import tools.mcp_tool as mcp_mod
         from tools.interrupt import set_interrupt
@@ -922,6 +940,54 @@ class TestRunOnMCPLoopInterrupts:
             mcp_mod._mcp_loop = old_loop
             mcp_mod._mcp_thread = old_thread
 
+    def test_completed_future_timeout_is_propagated_once(self):
+        import tools.mcp_tool as mcp_mod
+
+        inner_error = TimeoutError("inner MCP timeout")
+
+        class CompletedWithTimeout(concurrent.futures.Future):
+            def __init__(self):
+                super().__init__()
+                self.result_timeouts = []
+                self.set_exception(inner_error)
+
+            def result(self, timeout=None):
+                self.result_timeouts.append(timeout)
+                return super().result(timeout=timeout)
+
+        future = CompletedWithTimeout()
+
+        with pytest.raises(TimeoutError, match="inner MCP timeout") as exc_info:
+            self._run_with_future(mcp_mod, future)
+
+        assert exc_info.value is inner_error
+        assert len(future.result_timeouts) == 2
+        assert future.result_timeouts[0] is not None
+        assert future.result_timeouts[1] is None
+
+    def test_poll_timeout_racing_success_returns_completed_result(self):
+        import tools.mcp_tool as mcp_mod
+
+        class PollThenSuccess(concurrent.futures.Future):
+            def __init__(self):
+                super().__init__()
+                self.result_timeouts = []
+
+            def result(self, timeout=None):
+                self.result_timeouts.append(timeout)
+                if len(self.result_timeouts) == 1:
+                    self.set_result("completed")
+                    raise concurrent.futures.TimeoutError
+
+                return super().result(timeout=timeout)
+
+        future = PollThenSuccess()
+
+        assert self._run_with_future(mcp_mod, future) == "completed"
+        assert len(future.result_timeouts) == 2
+        assert future.result_timeouts[0] is not None
+        assert future.result_timeouts[1] is None
+
 
 # ---------------------------------------------------------------------------
 # Tool registration (discovery + register)

tools/mcp_tool.py

@@ -3909,6 +3909,13 @@ def _run_on_mcp_loop(coro_or_factory, timeout: float = 30):
         try:
             return future.result(timeout=wait_timeout)
         except concurrent.futures.TimeoutError:
+            # On supported Python versions, concurrent.futures.TimeoutError
+            # aliases the built-in TimeoutError, so result(timeout=...) also
+            # raises it for a coroutine's own timeout.
+            # Resolve a done future without a timeout to propagate its stored
+            # outcome, including completion racing with this polling timeout.
+            if future.done():
+                return future.result()
             continue