contrib/strands: cache MCP connections across tool calls

The per-server {server}-call-tool activity opened a fresh MCP session on
every invocation (open transport + initialize + call + teardown), so an
agent making several successive MCP calls paid that handshake per call --
and for stdio servers, a subprocess spawn per call.

Hold a lazily-opened MCP session per server in the activity worker
process so successive call-tool activities reuse one connection. A
dedicated owner task enters and exits the anyio transport/ClientSession
context managers in the same task (the cancel-scope rule); call-tool
activities on the same event loop invoke session.call_tool directly,
which MCP multiplexes by request id. Evict on idle timeout, on a call
error (so a broken session reconnects), and on worker shutdown -- scoped
to the servers the plugin registered rather than every cached connection.

A reused session now carries server-side state across workflows sharing a
worker, a behavior change from the previous per-call isolation.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: brianstrauch

temporalio/contrib/strands/_plugin.py

@@ -14,6 +14,7 @@
 from ._failure_converter import StrandsFailureConverter
 from ._model_activity import ModelActivity
 from ._temporal_mcp_client import (
+    _evict_connection,
     build_call_tool_activity,
     clear_cache,
     populate_cache,
@@ -67,6 +68,7 @@ async def run_context() -> AsyncGenerator[None, None]:
                 yield
             finally:
                 for server in mcp_clients:
+                    await _evict_connection(server)
                     clear_cache(server)
 
         super().__init__(

temporalio/contrib/strands/_temporal_mcp_client.py

@@ -1,3 +1,6 @@
+from __future__ import annotations
+
+import asyncio
 from collections.abc import Callable, Sequence
 from dataclasses import dataclass, field
 from datetime import timedelta
@@ -158,13 +161,36 @@ def clear_cache(server: str) -> None:
     _TOOL_CACHE.pop(server, None)
 
 
-def build_call_tool_activity(
-    server: str, client_factory: Callable[[], MCPClient]
-) -> Callable:
-    """Return the per-server ``{server}-call-tool`` activity for registration."""
+# How long an idle MCP connection stays open before it is disconnected.
+_MCP_CONNECTION_IDLE = timedelta(minutes=5)
 
-    @activity.defn(name=f"{server}-call-tool")
-    async def call_tool(args: _CallToolArgs) -> MCPToolResult:
+# Server name -> live connection held open in the activity worker process.
+# Activities run in the worker process , so this module state is shared across activity invocations on the worker
+_CONNECTIONS: dict[str, _ConnectionRecord] = {}
+
+
+class _ConnectionRecord:
+    """A single MCP session held open by a dedicated owner task.
+
+    The MCP transport and ``ClientSession`` are anyio context managers whose
+    cancel scope is bound to the task that enters them, so they must be entered
+    and exited in the same task. ``_run`` owns that task for the connection's
+    whole lifetime; ``call_tool`` activities on the same event loop invoke
+    ``session.call_tool`` directly (MCP multiplexes concurrent requests by id).
+    """
+
+    def __init__(self, server: str, client_factory: Callable[[], MCPClient]) -> None:
+        loop = asyncio.get_running_loop()
+        self._server = server
+        self._stop = asyncio.Event()
+        self._ready: asyncio.Future[tuple[MCPClient, ClientSession]] = (
+            loop.create_future()
+        )
+        self._idle_handle: asyncio.TimerHandle | None = None
+        self._idle_task: asyncio.Task[None] | None = None
+        self._owner = asyncio.create_task(self._run(client_factory))
+
+    async def _run(self, client_factory: Callable[[], MCPClient]) -> None:
         client = client_factory()
         try:
             async with client._transport_callable() as (read_stream, write_stream, *_):
@@ -174,9 +200,87 @@ async def call_tool(args: _CallToolArgs) -> MCPToolResult:
                     elicitation_callback=client._elicitation_callback,
                 ) as session:
                     await session.initialize()
-                    result = await session.call_tool(args.tool_name, args.arguments)
-                    return client._handle_tool_result(args.tool_use_id, result)
+                    self._ready.set_result((client, session))
+                    await self._stop.wait()
+        except BaseException as err:
+            # A failed connect should not be cached; drop it so the next call
+            # retries instead of awaiting a permanently rejected future.
+            if not self._ready.done():
+                self._ready.set_exception(err)
+            _CONNECTIONS.pop(self._server, None)
+            raise
+
+    def touch(self) -> None:
+        """Restart the idle-eviction timer because the connection was used."""
+        if self._idle_handle is not None:
+            self._idle_handle.cancel()
+        loop = asyncio.get_running_loop()
+        self._idle_handle = loop.call_later(
+            _MCP_CONNECTION_IDLE.total_seconds(), self._on_idle
+        )
+
+    def _on_idle(self) -> None:
+        self._idle_task = asyncio.ensure_future(_evict_connection(self._server))
+
+    async def aclose(self) -> None:
+        """Signal the owner task to exit its context managers and wait for it."""
+        if self._idle_handle is not None:
+            self._idle_handle.cancel()
+            self._idle_handle = None
+        self._stop.set()
+        try:
+            await self._owner
+        except BaseException:
+            pass
+
+    async def session(self) -> tuple[MCPClient, ClientSession]:
+        """Return the live client and session, or raise the connect failure."""
+        return await self._ready
+
+
+async def get_connection(
+    server: str, client_factory: Callable[[], MCPClient]
+) -> tuple[MCPClient, ClientSession]:
+    """Return the cached session for ``server``, opening one lazily if needed.
+
+    Concurrent first-callers dedupe onto a single connect handshake by awaiting
+    the same record.
+    """
+    record = _CONNECTIONS.get(server)
+    if record is None:
+        record = _ConnectionRecord(server, client_factory)
+        _CONNECTIONS[server] = record
+    record.touch()
+    return await record.session()
+
+
+async def _evict_connection(server: str) -> None:
+    record = _CONNECTIONS.pop(server, None)
+    if record is not None:
+        await record.aclose()
+
+
+def build_call_tool_activity(
+    server: str, client_factory: Callable[[], MCPClient]
+) -> Callable:
+    """Return the per-server ``{server}-call-tool`` activity for registration.
+
+    Reuses a worker-process MCP session opened lazily through ``client_factory``.
+    """
+
+    @activity.defn(name=f"{server}-call-tool")
+    async def call_tool(args: _CallToolArgs) -> MCPToolResult:
+        try:
+            client, session = await get_connection(server, client_factory)
+        except Exception as err:
+            # Connecting failed; map to a tool error result like a call would.
+            return client_factory()._handle_tool_execution_error(args.tool_use_id, err)
+        try:
+            result = await session.call_tool(args.tool_name, args.arguments)
+            return client._handle_tool_result(args.tool_use_id, result)
         except Exception as err:
+            # The session may be broken; drop it so the next call reconnects.
+            await _evict_connection(server)
             return client._handle_tool_execution_error(args.tool_use_id, err)
 
     return call_tool

tests/contrib/strands/test_mcp.py

@@ -12,12 +12,24 @@
     StrandsPlugin,
     TemporalAgent,
     TemporalMCPClient,
+    _temporal_mcp_client,
 )
 from temporalio.worker import Replayer, Worker
 from tests.contrib.strands.common import get_activities
 from tests.contrib.strands.mock_model import MockModel
 
 
+def _echo_client_factory() -> MCPClient:
+    return MCPClient(
+        lambda: stdio_client(
+            StdioServerParameters(
+                command=sys.executable,
+                args=[str(Path(__file__).parent / "echo_mcp_server.py")],
+            )
+        )
+    )
+
+
 @workflow.defn
 class MCPWorkflow:
     def __init__(self) -> None:
@@ -86,3 +98,82 @@ async def test_mcp(client: Client):
         workflows=[MCPWorkflow],
         plugins=[plugin],
     ).replay_workflow(history)
+
+
+@workflow.defn
+class MCPReuseWorkflow:
+    def __init__(self) -> None:
+        echo = TemporalMCPClient(
+            server="echo_cached",
+            start_to_close_timeout=timedelta(seconds=30),
+        )
+        self.agent = TemporalAgent(
+            model="mock",
+            start_to_close_timeout=timedelta(seconds=30),
+            tools=[echo],
+        )
+
+    @workflow.run
+    async def run(self, prompt: str) -> str:
+        result = await self.agent.invoke_async(prompt)
+        return str(result)
+
+
+async def test_mcp_reuses_connection(client: Client):
+    """Successive MCP tool calls reuse one cached worker-side connection."""
+    task_queue = "test_mcp_reuses_connection"
+    # Count how often the worker opens a connection. With caching this is one
+    # startup-discovery connection plus one cached call connection serving both
+    # tool calls (2); reconnecting per call would make it 3.
+    factory_calls = [0]
+
+    def counting_factory() -> MCPClient:
+        factory_calls[0] += 1
+        return _echo_client_factory()
+
+    plugin = StrandsPlugin(
+        models={
+            "mock": lambda: MockModel(
+                [
+                    {"name": "echo", "input": {"message": "one"}},
+                    {"name": "echo", "input": {"message": "two"}},
+                    "Done!",
+                ]
+            )
+        },
+        mcp_clients={"echo_cached": counting_factory},
+    )
+
+    async with Worker(
+        client,
+        task_queue=task_queue,
+        workflows=[MCPReuseWorkflow],
+        plugins=[plugin],
+        max_cached_workflows=0,
+    ):
+        handle = await client.start_workflow(
+            MCPReuseWorkflow.run,
+            "echo twice",
+            id=f"test_mcp_reuses_connection_{uuid4()}",
+            task_queue=task_queue,
+        )
+        assert await handle.result() == "Done!\n"
+
+    # The worker context has exited, so its run_context finally evicted the
+    # cached connection.
+    assert "echo_cached" not in _temporal_mcp_client._CONNECTIONS
+    assert factory_calls[0] == 2
+
+    history = await handle.fetch_history()
+    assert get_activities(history) == [
+        "invoke_model",
+        "echo_cached-call-tool",
+        "invoke_model",
+        "echo_cached-call-tool",
+        "invoke_model",
+    ]
+
+    await Replayer(
+        workflows=[MCPReuseWorkflow],
+        plugins=[plugin],
+    ).replay_workflow(history)