fix(mcp): propagate stored timeouts from completed futures

This commit is contained in:
Jupiter363 2026-07-14 02:44:27 +08:00 committed by Brooklyn Nicholson
parent 1310ceb07b
commit 3df8bd3478
2 changed files with 73 additions and 0 deletions

View file

@ -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)

View file

@ -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