diff --git a/strix/tools/mcp/session.py b/strix/tools/mcp/session.py index e76a3eb0..d3062b45 100644 --- a/strix/tools/mcp/session.py +++ b/strix/tools/mcp/session.py @@ -66,6 +66,11 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) +# How long a graceful (sentinel) shutdown waits for the serve loop to drain +# before the supervising task is cancelled instead. Bounds teardown so a slow or +# hung in-flight call cannot stall it forever. +_SHUTDOWN_TIMEOUT = 10.0 + class McpConnectionUnavailableError(RuntimeError): """A dead MCP connection could not be reached and did not come back. @@ -229,27 +234,42 @@ class SupervisedMcpSession: async def aclose(self) -> None: """Shut the connection down and clean up its session on its owning task. - For a supervised session this signals the supervising task with a sentinel - so ``cleanup()`` runs on the same task that ran ``connect()``. It never - cancels the task, so the supervisor can tell an orderly shutdown from a - session death. + For a connected supervised session this signals the supervising task with a + sentinel so ``cleanup()`` runs on the same task that ran ``connect()``, + giving an orderly shutdown the supervisor tells apart from a session death. + Teardown is always bounded: if the serve loop cannot drain the sentinel in + time (a slow or hung in-flight call), or the session never finished + connecting (including a connect cancelled mid-await), the task is cancelled + instead. ``_closing`` is set first, so the supervisor treats that + cancellation as shutdown and still cleans up on its own task. """ self._closing = True if self._supervised and self._task is not None: if not self._task.done(): + # A cancelled readiness future (the connect was cancelled mid-await) + # counts as "not connected": never call ``.result()`` on it, which + # would raise here and skip the cleanup below. connected = ( - self._ready is not None and self._ready.done() and self._ready.result() + self._ready is not None + and self._ready.done() + and not self._ready.cancelled() + and self._ready.result() ) if connected and self._queue is not None: # Reached the serve loop: a sentinel gives a clean, cancel-free - # teardown, with cleanup() running on the supervising task. + # teardown, with cleanup() running on the supervising task. Bound + # it, though: a hung in-flight call would otherwise leave the + # sentinel queued behind it forever, so cancel the task if the + # drain does not finish in time (wait_for cancels it on timeout). with contextlib.suppress(Exception): await self._queue.put(None) + with contextlib.suppress( + asyncio.TimeoutError, asyncio.CancelledError, Exception + ): + await asyncio.wait_for(self._task, _SHUTDOWN_TIMEOUT) else: - # Still stuck in connect() (or never connected): cancel to - # unstick it. ``_closing`` is already set, so the supervisor - # treats the cancellation as shutdown and still cleans up on - # its own task. + # Still stuck in connect(), never connected, or connect + # cancelled: cancel to unstick it. self._task.cancel() with contextlib.suppress(asyncio.CancelledError, Exception): await self._task diff --git a/tests/test_mcp_client.py b/tests/test_mcp_client.py index 6305a08a..c75efa4d 100644 --- a/tests/test_mcp_client.py +++ b/tests/test_mcp_client.py @@ -8,6 +8,7 @@ two dispatch tools ``describe_mcp`` and ``call_mcp``. from __future__ import annotations import asyncio +import contextlib import json import re from typing import TYPE_CHECKING, Any @@ -39,6 +40,7 @@ from strix.tools.mcp import ( resolve_mcp_call, ) from strix.tools.mcp import client as mcp_client +from strix.tools.mcp import session as mcp_session_mod if TYPE_CHECKING: @@ -1374,6 +1376,88 @@ async def test_flapping_idle_session_is_marked_dead_without_looping( await session.aclose() +class _HangingCallServer(FakeMCPServer): + """A connected server whose ``call_tool`` never returns, modeling a hung + in-flight call so teardown can be tested for boundedness.""" + + def __init__(self, name: str, tools: list[MCPTool]) -> None: + super().__init__(name, tools) + self.cleaned = False + + async def call_tool( + self, + tool_name: str, + arguments: dict[str, Any] | None, + meta: dict[str, Any] | None = None, + ) -> CallToolResult: + await asyncio.Event().wait() + raise AssertionError("unreachable") + + async def cleanup(self) -> None: + self.cleaned = True + + +class _HangingConnectServer(FakeMCPServer): + """A server whose ``connect`` never finishes, so ``start`` blocks on readiness + and can be cancelled mid-connect.""" + + def __init__(self, name: str, tools: list[MCPTool]) -> None: + super().__init__(name, tools) + self.cleaned = False + + async def connect(self) -> None: + await asyncio.Event().wait() + + async def cleanup(self) -> None: + self.cleaned = True + + +@pytest.mark.asyncio +async def test_aclose_is_bounded_when_an_in_flight_call_hangs( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # A hung call must not queue the shutdown sentinel behind itself forever: + # aclose falls back to cancelling the supervising task, and cleanup still runs. + monkeypatch.setattr(mcp_session_mod, "_SHUTDOWN_TIMEOUT", 0.2) + server = _HangingCallServer("fs", [_mcp_tool("read_file")]) + monkeypatch.setattr(mcp_client, "_build_server", lambda _config: server) + + session = await _started_session(_secret_config("fs")) + call = asyncio.create_task(session.dispatch("read_file", {}, label="fs_read_file")) + await asyncio.sleep(0.05) # let the serve loop pick up the request and hang + + # Must return promptly rather than block on the hung call. + await asyncio.wait_for(session.aclose(), timeout=3.0) + assert session._task is not None and session._task.done() + assert server.cleaned is True + + # The abandoned caller gets a value (dead), not a hang. + out = await asyncio.wait_for(call, timeout=3.0) + assert isinstance(out, dict) and out["success"] is False + + +@pytest.mark.asyncio +async def test_aclose_cleans_up_when_connect_is_cancelled_mid_await( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # If the scan is cancelled while start() awaits readiness, the readiness future + # is cancelled; aclose must not raise on it and must still cancel + clean up the + # partially connected supervisor. + server = _HangingConnectServer("fs", [_mcp_tool("read_file")]) + monkeypatch.setattr(mcp_client, "_build_server", lambda _config: server) + + session = SupervisedMcpSession(_secret_config("fs")) + start = asyncio.create_task(session.start()) + await asyncio.sleep(0.05) # let the supervisor reach the hanging connect() + start.cancel() + with contextlib.suppress(asyncio.CancelledError): + await start + + await asyncio.wait_for(session.aclose(), timeout=3.0) + assert session._task is not None and session._task.done() + assert server.cleaned is True + + @pytest.mark.asyncio async def test_a_session_death_is_contained_and_other_connections_survive( monkeypatch: pytest.MonkeyPatch,