diff --git a/tests/tools/test_mcp_poll_loop_oom_integration.py b/tests/tools/test_mcp_poll_loop_oom_integration.py new file mode 100644 index 000000000000..ec2b1766f667 --- /dev/null +++ b/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 " 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..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} diff --git a/tests/tools/test_mcp_tool.py b/tests/tools/test_mcp_tool.py index c2208165fa4b..a66235040284 100644 --- a/tests/tools/test_mcp_tool.py +++ b/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) diff --git a/tools/mcp_tool.py b/tools/mcp_tool.py index 3e6844de3d0b..7da89baff599 100644 --- a/tools/mcp_tool.py +++ b/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