From f1df0aa01b9c75925b096edffff1ba1335f31bb4 Mon Sep 17 00:00:00 2001 From: kshitijk4poor <82637225+kshitijk4poor@users.noreply.github.com> Date: Thu, 17 Sep 2026 20:29:06 +0530 Subject: [PATCH] fix(mcp): an idle stdio server is proven only after a full default interval The no-keepalive proof branch ran on ANY timeout wake. With idle_timeout_seconds=60 the wait is min(180, 60); an overlapping RPC hides the recycle deadline so _recycle_if_due() is False and the session was marked proven after 60 s, clearing the #62212 rapid-drop budget early. Fix a proof_at deadline before the loop and gate the branch on it; the unproven timeout is the remaining distance to proof_at. Tests: the stdio test now shows the first (idle-limit) wake does not prove; the interval test covers an explicit keepalive_interval on stdio. --- tests/tools/test_mcp_capability_gating.py | 30 +++++++++++++++++------ tools/mcp_tool_server_run.py | 18 +++++++++----- 2 files changed, 35 insertions(+), 13 deletions(-) diff --git a/tests/tools/test_mcp_capability_gating.py b/tests/tools/test_mcp_capability_gating.py index 6f06054f74..794a584e85 100644 --- a/tests/tools/test_mcp_capability_gating.py +++ b/tests/tools/test_mcp_capability_gating.py @@ -147,19 +147,23 @@ class TestKeepaliveInterval: the session alive instead of hitting an expired session on every idle call. """ - async def _run_lifecycle_cycles(self, config): + async def _run_lifecycle_cycles(self, config, idle_timeout=None): """Run two lifecycle cycles on a server with the *full* ``config``: the first - ``asyncio.wait`` times out (no event fired), the second is ended by shutdown. + ``asyncio.wait`` times out (no event fired) and the loop's clock is advanced by + that timeout, the second is ended by shutdown. Returns ``(task, [timeout_of_cycle_1, timeout_of_cycle_2])``.""" task = MCPServerTask("test") task._config = config + task._idle_timeout_seconds = idle_timeout task.session = SimpleNamespace(send_ping=AsyncMock()) timeouts = [] real_wait = asyncio.wait + elapsed = [0.0] async def fake_wait(tasks, timeout=None, return_when=None): timeouts.append(timeout) if len(timeouts) == 1: + elapsed[0] += timeout or 0.0 return set(), set(tasks) # simulate the timeout firing task._shutdown_event.set() return await real_wait( @@ -167,12 +171,14 @@ class TestKeepaliveInterval: ) import tools.mcp_tool as mcp_mod - orig = mcp_mod.asyncio.wait + import tools.mcp_tool_server_run as run_mod + orig, orig_time = mcp_mod.asyncio.wait, run_mod.time mcp_mod.asyncio.wait = fake_wait + run_mod.time = SimpleNamespace(monotonic=lambda: orig_time.monotonic() + elapsed[0]) try: assert await task._wait_for_lifecycle_event() == "shutdown" finally: - mcp_mod.asyncio.wait = orig + mcp_mod.asyncio.wait, run_mod.time = orig, orig_time return task, timeouts async def _captured_interval(self, config): @@ -185,7 +191,10 @@ class TestKeepaliveInterval: async def test_default_interval_when_unset(self): from tools.mcp_tool import _DEFAULT_KEEPALIVE_INTERVAL assert await self._captured_interval({}) == _DEFAULT_KEEPALIVE_INTERVAL - + # An explicit interval is honoured on stdio too (where there is no default). + _task, timeouts = await self._run_lifecycle_cycles( + {"command": "example-mcp", "keepalive_interval": 42}) + assert timeouts[0] == 42 @pytest.mark.asyncio async def test_interval_clamped_to_floor(self): @@ -200,14 +209,21 @@ class TestKeepaliveInterval: async def test_default_stdio_connection_waits_without_keepalive(self): """A healthy local pipe must not be probed solely because it is idle — but surviving a full default interval idle must still prove the session (clears the rapid-drop budget); - once proven, the loop waits without any timeout.""" + once proven, the loop waits without any timeout. An earlier wake caused by a shorter + idle limit is NOT proof: the budget stays armed until the full interval has passed.""" from tools.mcp_tool import _DEFAULT_KEEPALIVE_INTERVAL task, timeouts = await self._run_lifecycle_cycles({"command": "example-mcp"}) - assert timeouts == [_DEFAULT_KEEPALIVE_INTERVAL, None] + assert timeouts[0] == pytest.approx(_DEFAULT_KEEPALIVE_INTERVAL, abs=0.05) + assert timeouts[1] is None assert task._session_proven is True task.session.send_ping.assert_not_called() + task, timeouts = await self._run_lifecycle_cycles({"command": "example-mcp"}, idle_timeout=60) + assert timeouts[0] == pytest.approx(60, abs=0.05) # woke for the idle deadline, not the proof + assert task._session_proven is False + task.session.send_ping.assert_not_called() + @pytest.mark.asyncio async def test_stdio_recycle_wakes_after_active_rpc(self): """Lock release must expose an elapsed recycle deadline without sending a ping.""" diff --git a/tools/mcp_tool_server_run.py b/tools/mcp_tool_server_run.py index 253c48b56b..36a5ed8920 100644 --- a/tools/mcp_tool_server_run.py +++ b/tools/mcp_tool_server_run.py @@ -71,6 +71,12 @@ class MCPServerRunMixin: keepalive_interval = max( _core._MIN_KEEPALIVE_INTERVAL, float(_core._DEFAULT_KEEPALIVE_INTERVAL if configured is None else configured)) + # No keepalive, but an unproven stdio session must still get its chance to prove + # itself: it counts as proven only once a FULL default interval has elapsed — not on + # the first timeout wake, which a shorter recycle deadline may cause (see below). + proof_at = None + if keepalive_interval is None and not self._session_proven: + proof_at = time.monotonic() + _core._DEFAULT_KEEPALIVE_INTERVAL shutdown_task, reconnect_task = self._event_waiters() rpc_idle_task = None waiters = [shutdown_task, reconnect_task] @@ -79,10 +85,8 @@ class MCPServerRunMixin: if self._recycle_if_due(): return "recycle" timeout = keepalive_interval - if timeout is None and not self._session_proven: - # No keepalive, but an unproven stdio session must still get its chance to - # prove itself: wake once after the default interval (no ping) — see below. - timeout = float(_core._DEFAULT_KEEPALIVE_INTERVAL) + if timeout is None and not self._session_proven and proof_at is not None: + timeout = max(0.0, proof_at - time.monotonic()) recycle_deadline = self._next_stdio_recycle_deadline() if recycle_deadline is not None: recycle_timeout = max(0.0, recycle_deadline - time.monotonic()) @@ -106,8 +110,10 @@ class MCPServerRunMixin: if keepalive_interval is None: # Stdio without a keepalive: idling a full default interval with the child # still alive is the proof of health a successful ping gives remote - # transports — clear the rapid-drop budget without pinging (#62212). - if not self._session_proven and not self._stdio_children_dead(): + # transports — clear the rapid-drop budget without pinging (#62212). An + # earlier wake (a hidden-then-missed recycle deadline) is not that proof. + if (not self._session_proven and proof_at is not None + and time.monotonic() >= proof_at and not self._stdio_children_dead()): self._mark_session_proven() continue # Timeout: probe for a stale session — NEVER while an RPC is in flight (a