loom-kernel 2.2.0__tar.gz → 2.2.1__tar.gz

{loom_kernel-2.2.0 → loom_kernel-2.2.1}/CHANGELOG.md

@@ -123,6 +123,21 @@
   The `etl-polars` extra now requires `polars>=1.24`, the release that
   named that join option `nulls_equal`.
 
+### ai
+
+- **ai:** an MCP session that dies after it opened no longer fails every later
+  call until the worker restarts. The worker holds one session per server for
+  its whole life, and that holder kept the dead client from ever being
+  re-entered: after a transport error under an in-flight call, a redeploy of
+  the server or a stdio child that exited, each call through `Mcp()` failed in
+  a millisecond while the health probe answered `ok`. `_ToolsetSession` now
+  treats an exception out of a call's round trip as the transport failing: that
+  call is raised as it is and never repeated, and the next call opens a new
+  session first, once even under concurrency. A renewal that cannot connect
+  raises and is tried again by the following call. An agent's own `mcp`
+  capability does not recover yet; see "A session that died after it opened is
+  replaced by the next call" in `docs/ai/mcp.md`.
+
 ## ♻️ Refactor
 
 ### etl

{loom_kernel-2.2.0 → loom_kernel-2.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: loom-kernel
-Version: 2.2.0
+Version: 2.2.1
 Summary: Loom Python project
 License-File: LICENSE
 Classifier: Typing :: Typed

{loom_kernel-2.2.0 → loom_kernel-2.2.1}/docs/ai/mcp.md

@@ -154,8 +154,10 @@ behind an address; stdio is for the servers you only consume.
 
 What stdio does not do:
 
-- **it does not reconnect.** A dead subprocess fails the call, exactly as a
-  dead HTTP server does; nothing restarts it;
+- **it does not supervise the child.** A dead subprocess fails the call that
+  finds it, exactly as a dead HTTP server does. A call made through `Mcp()`
+  after that one spawns a new child (see "A session that died after it opened
+  is replaced by the next call"); an agent's own capability does not;
 - **it spawns one process per server, for the whole worker.** Start-up spawns
   it and every agent granted that server speaks to that one child. It dies with
   the last holder of the connection — normally the runtime, briefly a straggler
@@ -585,6 +587,43 @@ surfaces only on the first business request that reaches it, as
 `TOOL_UNAVAILABLE`. Extending the probe to a use-case-only server is future
 work, not something this version does.
 
+### A session that died after it opened is replaced by the next call
+
+A server that connected at start-up can still go away later: a transport error
+under an in-flight call, a redeploy of the server, a stdio child that exits.
+The worker's one session is held open for its whole life, and that holder is
+also what stops the underlying client from ever being re-entered — so before
+this was handled, every call through `Mcp()` failed from then on, in a
+millisecond each, until the worker restarted, while the health probe kept
+answering `ok`.
+
+`_ToolsetSession` now replaces the session itself. Nothing can be asked of the
+client: once its transport is gone it still reports itself connected. The
+evidence is the call. A tool's own error comes back as a result flagged
+`is_error`; **an exception out of the round trip is the transport failing**.
+So:
+
+- **the call that found it dead is raised as it is, and never repeated.** It
+  may have run on the server, and a tool with side effects is not run twice on
+  your behalf. Retrying it is your decision;
+- **the next call opens a new session first**, against the same connection —
+  over stdio, that is a new child process. Concurrent calls that find it dead
+  renew it once;
+- **a renewal that cannot connect raises**, and the call after it tries again:
+  an outage costs one connection attempt per call, bounded by the handshake
+  deadline, and ends with the first call made after the server is back;
+- **suspecting a healthy session costs one reconnection.** A single failed call
+  is not proof the session died; opening another is cheap and harmless, so it
+  is not weighed further.
+
+This covers the `Mcp()` path and whatever else calls through the runtime's
+session. **An agent's own `mcp` capability does not recover yet**: the compiled
+agent holds the toolset it was built with, and swapping that needs the engine
+build to go through an indirection it does not have. Two other things are
+unchanged: a server that never connected at start-up under `optional` is still
+not reconnected (above), and the health probe still does not see a
+use-case-only server.
+
 ### A use case's call runs in parallel with other grant views
 
 A use case's tool call goes through `_ToolsetSession`

{loom_kernel-2.2.0 → loom_kernel-2.2.1}/src/loom/ai/engines/pydantic_ai/_mcp.py

@@ -36,7 +36,9 @@ break every deployment that declares no ``mcp`` grant.
 
 from __future__ import annotations
 
-from collections.abc import AsyncIterator, Mapping
+import asyncio
+import logging
+from collections.abc import AsyncIterator, Callable, Mapping
 from contextlib import asynccontextmanager
 from typing import TYPE_CHECKING, Any, Final, cast
 
@@ -49,6 +51,8 @@ if TYPE_CHECKING:
     from fastmcp.client.transports import ClientTransport
     from pydantic_ai.mcp import MCPToolset
 
+_logger = logging.getLogger(__name__)
+
 DEFAULT_MCP_CONNECT_TIMEOUT_SECONDS: Final[float] = 10.0
 """Fallback handshake deadline, seconds, matching
 :attr:`~loom.ai.config.AiConfig.startup_timeout_ms`'s own default. Production
@@ -74,12 +78,58 @@ class _ToolsetSession(ConcurrentMcpSession):
     ``__aexit__`` to decide when the connection opens and closes; that is a
     connection-lifecycle concern, not what makes concurrent calls safe.
 
+    A session that died after it opened is replaced by the next call. The
+    ``fastmcp`` client still reports itself connected once its transport is
+    gone, so the evidence is the call: a tool's own error is a result flagged
+    ``is_error``, and an exception out of the round trip is the transport
+    failing. That call is raised and never repeated -- it may have run on the
+    server -- and the one after it opens a fresh toolset first.
+
     Args:
         toolset: Already-entered toolset speaking to one server.
+        rebuild: Builds a new, not yet entered toolset for the same
+            connection; ``None`` never replaces the one given.
     """
 
-    def __init__(self, toolset: MCPToolset[Any]) -> None:
+    def __init__(
+        self,
+        toolset: MCPToolset[Any],
+        *,
+        rebuild: Callable[[], MCPToolset[Any]] | None = None,
+    ) -> None:
         self._toolset = toolset
+        self._rebuild = rebuild
+        self._suspect: MCPToolset[Any] | None = None
+        self._renewing = asyncio.Lock()
+        self._renewed: list[MCPToolset[Any]] = []
+
+    async def _live(self) -> MCPToolset[Any]:
+        """Return the toolset to speak through, replacing a suspect one at most once.
+
+        A renewal that cannot connect raises and leaves the suspect in place,
+        so the next call tries again.
+        """
+        toolset = self._toolset
+        if self._rebuild is None or self._suspect is not toolset:
+            return toolset
+        async with self._renewing:
+            if self._toolset is not toolset:
+                return self._toolset
+            fresh = self._rebuild()
+            await fresh.__aenter__()
+            self._renewed.append(fresh)
+            self._toolset = fresh
+            _logger.warning("mcp session renewed after a failed call")
+            return fresh
+
+    async def aclose(self) -> None:
+        """Close the toolsets this session opened itself; the first belongs to its opener."""
+        while self._renewed:
+            toolset = self._renewed.pop()
+            try:
+                await toolset.__aexit__(None, None, None)
+            except Exception:
+                _logger.warning("a renewed mcp session did not close cleanly", exc_info=True)
 
     async def list_tools(self) -> tuple[McpToolInfo, ...]:
         """Return the tools the server exposes, each with its output-schema flag.
@@ -91,7 +141,7 @@ class _ToolsetSession(ConcurrentMcpSession):
         """
         return tuple(
             McpToolInfo(name=tool.name, has_output_schema=tool.output_schema is not None)
-            for tool in await self._toolset.list_tools()
+            for tool in await (await self._live()).list_tools()
         )
 
     async def call_tool(self, name: str, arguments: Mapping[str, Any]) -> McpToolCallResult:
@@ -125,8 +175,13 @@ class _ToolsetSession(ConcurrentMcpSession):
         Returns:
             The server's own error flag and structured content, verbatim.
         """
-        async with self._toolset:
-            result = await self._toolset.client.call_tool_mcp(name, dict(arguments))
+        toolset = await self._live()
+        try:
+            async with toolset:
+                result = await toolset.client.call_tool_mcp(name, dict(arguments))
+        except Exception:
+            self._suspect = toolset
+            raise
         return McpToolCallResult(ok=not result.is_error, structured=result.structured_content)
 
 
@@ -370,7 +425,14 @@ class SharedMcpToolsets:
         self._opened.add(connection)
         toolset = self._toolset(connection)
         async with toolset:
-            yield _ToolsetSession(toolset)
+            session = _ToolsetSession(
+                toolset,
+                rebuild=lambda: build_mcp_toolset(connection, init_timeout=self._init_timeout),
+            )
+            try:
+                yield session
+            finally:
+                await session.aclose()
 
 
 @asynccontextmanager

loom_kernel-2.2.1/tests/integration/ai/test_mcp_session_recovery.py

@@ -0,0 +1,65 @@
+"""A shared MCP session that died after it opened serves the next call again.
+
+The worker opens one session per server at start-up and every run and every
+``Mcp()`` marker speaks through it. When the server behind it goes away after
+that -- a transport error under an in-flight call, a redeploy, a child that
+exits -- the holder the runtime keeps open also keeps the dead client from ever
+being re-entered, so every later call failed until the process was restarted.
+Driven over a real stdio child, which the test can kill.
+"""
+
+from __future__ import annotations
+
+import os
+import signal
+import sys
+from pathlib import Path
+
+import pytest
+
+from loom.ai.compiler import CompiledMcpCapability
+from loom.ai.engines.pydantic_ai._mcp import SharedMcpToolsets
+from tests.integration.ai.conftest import _until
+
+pytest.importorskip("fastmcp", reason="fastmcp is not installed: uv sync --group mcp-tests")
+
+_SCRIPT = str(Path(__file__).parent / "mcp_stdio_server.py")
+
+
+def _capability(pid_file: Path) -> CompiledMcpCapability:
+    return CompiledMcpCapability(
+        server="orders",
+        transport="stdio",
+        command=sys.executable,
+        args=(_SCRIPT, str(pid_file), "recovery"),
+        timeout_ms=5_000,
+    )
+
+
+def _kill(pid_file: Path) -> int:
+    pid = int(pid_file.read_text(encoding="utf-8"))
+    os.kill(pid, signal.SIGKILL)
+    return pid
+
+
+class TestASessionThatDiedAfterOpening:
+    async def test_the_call_after_the_failed_one_is_served_by_a_new_session(
+        self, tmp_path: Path
+    ) -> None:
+        pid_file = tmp_path / "server.pid"
+        shared = SharedMcpToolsets()
+
+        async with shared.open(_capability(pid_file)) as session:
+            first = await session.call_tool("read_orders", {"customer": "acme"})
+            assert first.ok
+
+            dead = _kill(pid_file)
+            pid_file.unlink()
+            with pytest.raises(Exception):  # noqa: B017, PT011 -- whatever the dead transport raises
+                await session.call_tool("read_orders", {"customer": "acme"})
+
+            again = await session.call_tool("read_orders", {"customer": "acme"})
+
+            assert again.ok
+            await _until(pid_file.exists)
+            assert int(pid_file.read_text(encoding="utf-8")) != dead

loom_kernel-2.2.1/tests/unit/ai/engines/test_pydantic_ai_mcp_session_recovery.py

@@ -0,0 +1,170 @@
+"""What ``_ToolsetSession`` promises about a session that died after it opened.
+
+Driven over toolset doubles, so each promise is one observable fact: how many
+toolsets were built, which one served a call, and how often the tool ran.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from types import SimpleNamespace
+from typing import Any
+
+import pytest
+
+from loom.ai.engines.pydantic_ai._mcp import _ToolsetSession
+
+
+class _Toolset:
+    """A toolset double: counts its calls and fails them while ``dead``."""
+
+    def __init__(self, label: str, *, dead: bool = False, connects: bool = True) -> None:
+        self.label = label
+        self.dead = dead
+        self.connects = connects
+        self.calls: list[str] = []
+        self.entered = 0
+        self.exited = 0
+        self.client = SimpleNamespace(call_tool_mcp=self._call)
+
+    async def __aenter__(self) -> _Toolset:
+        if not self.connects:
+            raise ConnectionError(f"{self.label} refused")
+        self.entered += 1
+        return self
+
+    async def __aexit__(self, *args: object) -> None:
+        self.exited += 1
+
+    async def _call(self, name: str, arguments: dict[str, Any]) -> Any:
+        self.calls.append(name)
+        await asyncio.sleep(0)
+        if self.dead:
+            raise ConnectionError("Connection closed")
+        return SimpleNamespace(is_error=name == "fails", structured_content={"by": self.label})
+
+
+class _Factory:
+    """Hands out the toolsets it was given, in order, and counts how many it built."""
+
+    def __init__(self, *toolsets: _Toolset) -> None:
+        self._toolsets = list(toolsets)
+        self.built = 0
+
+    def __call__(self) -> Any:
+        self.built += 1
+        return self._toolsets.pop(0)
+
+
+def _session(first: _Toolset, factory: _Factory | None) -> _ToolsetSession:
+    return _ToolsetSession(first, rebuild=factory)  # type: ignore[arg-type]
+
+
+class TestTheCallThatFoundTheSessionDead:
+    async def test_it_is_raised_as_it_is_and_never_run_a_second_time(self) -> None:
+        first, second = _Toolset("first", dead=True), _Toolset("second")
+        factory = _Factory(second)
+        session = _session(first, factory)
+
+        with pytest.raises(ConnectionError, match="Connection closed"):
+            await session.call_tool("write_orders", {})
+
+        assert first.calls == ["write_orders"]
+        assert second.calls == []
+        assert factory.built == 0
+
+    async def test_the_next_call_is_served_by_a_new_toolset(self) -> None:
+        first, second = _Toolset("first", dead=True), _Toolset("second")
+        session = _session(first, _Factory(second))
+        with pytest.raises(ConnectionError):
+            await session.call_tool("read_orders", {})
+
+        result = await session.call_tool("read_orders", {})
+
+        assert (result.ok, result.structured) == (True, {"by": "second"})
+        assert second.entered - second.exited == 1
+
+    async def test_listing_tools_also_goes_through_the_new_toolset(self) -> None:
+        first, second = _Toolset("first", dead=True), _Toolset("second")
+        second.list_tools = lambda: _async(  # type: ignore[attr-defined]
+            [SimpleNamespace(name="read_orders", output_schema=None)]
+        )
+        session = _session(first, _Factory(second))
+        with pytest.raises(ConnectionError):
+            await session.call_tool("read_orders", {})
+
+        (tool,) = await session.list_tools()
+
+        assert tool.name == "read_orders"
+
+
+class TestWhatDoesNotRenewASession:
+    async def test_a_tools_own_error_is_a_result_not_a_dead_session(self) -> None:
+        first = _Toolset("first")
+        factory = _Factory(_Toolset("never"))
+        session = _session(first, factory)
+
+        failed = await session.call_tool("fails", {})
+        await session.call_tool("read_orders", {})
+
+        assert failed.ok is False
+        assert factory.built == 0
+        assert first.calls == ["fails", "read_orders"]
+
+    async def test_a_session_without_a_rebuild_keeps_its_toolset_for_life(self) -> None:
+        first = _Toolset("first", dead=True)
+        session = _session(first, None)
+
+        for _ in range(2):
+            with pytest.raises(ConnectionError):
+                await session.call_tool("read_orders", {})
+
+        assert first.calls == ["read_orders", "read_orders"]
+
+
+class TestRenewal:
+    async def test_concurrent_calls_that_find_it_dead_renew_it_once(self) -> None:
+        first, second = _Toolset("first", dead=True), _Toolset("second")
+        factory = _Factory(second, _Toolset("one too many"))
+        session = _session(first, factory)
+        with pytest.raises(ConnectionError):
+            await session.call_tool("read_orders", {})
+
+        results = await asyncio.gather(*(session.call_tool("read_orders", {}) for _ in range(8)))
+
+        assert factory.built == 1
+        assert [result.structured for result in results] == [{"by": "second"}] * 8
+
+    async def test_a_renewal_that_cannot_connect_is_tried_again_by_the_next_call(self) -> None:
+        first = _Toolset("first", dead=True)
+        down, back = _Toolset("down", connects=False), _Toolset("back")
+        factory = _Factory(down, back)
+        session = _session(first, factory)
+        with pytest.raises(ConnectionError, match="Connection closed"):
+            await session.call_tool("read_orders", {})
+
+        with pytest.raises(ConnectionError, match="down refused"):
+            await session.call_tool("read_orders", {})
+        result = await session.call_tool("read_orders", {})
+
+        assert result.structured == {"by": "back"}
+        assert factory.built == 2
+
+    async def test_closing_the_session_closes_what_it_opened_and_not_its_openers(self) -> None:
+        first, second = _Toolset("first", dead=True), _Toolset("second")
+        session = _session(first, _Factory(second))
+        with pytest.raises(ConnectionError):
+            await session.call_tool("read_orders", {})
+        await session.call_tool("read_orders", {})
+
+        held_before = (first.entered - first.exited, second.entered - second.exited)
+
+        await session.aclose()
+
+        assert held_before == (0, 1)
+        assert second.entered - second.exited == 0
+        assert first.exited == 1
+
+
+async def _async(value: Any) -> Any:
+    return value