Bound MCP session teardown and handle a cancelled connect

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

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,