Use MCP server in workflow with in-memory transport

Author: dandavison

nexusmcp/__init__.py

@@ -1,9 +1,10 @@
-from temporalio import workflow
+import temporalio.workflow
+
+with temporalio.workflow.unsafe.imports_passed_through():
+    from nexusmcp import workflow
 
-with workflow.unsafe.imports_passed_through():
     from .inbound_gateway import InboundGateway
     from .service import MCPService
     from .service_handler import MCPServiceHandler, exclude
-    from .workflow_transport import WorkflowTransport
 
-__all__ = ["MCPService", "MCPServiceHandler", "InboundGateway", "exclude", "WorkflowTransport"]
+__all__ = ["MCPService", "MCPServiceHandler", "InboundGateway", "exclude", "workflow"]

nexusmcp/workflow/__init__.py

@@ -0,0 +1,3 @@
+from .mcp_client import MCPClient
+
+__all__ = ["MCPClient"]

nexusmcp/workflow/mcp_client.py

@@ -0,0 +1,66 @@
+from contextlib import asynccontextmanager
+from typing import Any, AsyncGenerator
+
+import mcp.types as types
+from mcp.client.session import ClientSession
+from mcp.server.lowlevel import Server
+from mcp.shared.memory import create_connected_server_and_client_session
+from temporalio import workflow
+
+from nexusmcp.service import MCPService
+
+
+class MCPClient:
+    """
+    An MCP client for use in Temporal workflows.
+
+    This class provides a client that proxies MCP traffic from a Temporal Workflow to a Temporal
+    Nexus service. It works by running an MCP server in the workflow whose handlers delegate to
+    nexus operations, and connecting to it via an in-memory transport.
+
+    Example:
+        ```python
+        client = MCPClient("my-endpoint")
+        async with client.connect() as session:
+            await session.list_tools()
+            await session.call_tool("my-service_my-operation", {"arg": "value"})
+        ```
+    """
+
+    def __init__(
+        self,
+        endpoint: str,
+    ):
+        self.endpoint = endpoint
+        # Run an in-workflow MCP server whose handlers make nexus calls
+        self.mcp_server = Server("workflow-gateway-mcp-server")
+        self.mcp_server.list_tools()(self._handle_list_tools)  # type: ignore[no-untyped-call]
+        self.mcp_server.call_tool()(self._handle_call_tool)
+
+    @asynccontextmanager
+    async def connect(self) -> AsyncGenerator[ClientSession, None]:
+        """
+        Create a connected MCP ClientSession.
+
+        The session is automatically initialized before being yielded.
+        """
+        async with create_connected_server_and_client_session(
+            self.mcp_server,
+            raise_exceptions=True,
+        ) as session:
+            yield session
+
+    async def _handle_list_tools(self) -> list[types.Tool]:
+        nexus_client = workflow.create_nexus_client(
+            endpoint=self.endpoint,
+            service=MCPService,
+        )
+        return await nexus_client.execute_operation(MCPService.list_tools, None)
+
+    async def _handle_call_tool(self, name: str, arguments: dict[str, Any]) -> Any:
+        service, _, operation = name.partition("_")
+        nexus_client = workflow.create_nexus_client(
+            endpoint=self.endpoint,
+            service=service,
+        )
+        return await nexus_client.execute_operation(operation, arguments)

nexusmcp/workflow_transport.py

@@ -2,15 +2,14 @@
 from dataclasses import dataclass
 
 import pytest
-from mcp import ClientSession
 from temporalio import workflow
 from temporalio.api.nexus.v1 import EndpointSpec, EndpointTarget
 from temporalio.api.operatorservice.v1 import CreateNexusEndpointRequest
 from temporalio.contrib.pydantic import pydantic_data_converter
 from temporalio.testing import WorkflowEnvironment
 from temporalio.worker import Worker
 
-from nexusmcp import WorkflowTransport
+import nexusmcp.workflow
 
 from .service import TestServiceHandler, mcp_service
 
@@ -26,18 +25,15 @@ class MCPCallerWorkflowInput:
 class MCPCallerWorkflow:
     @workflow.run
     async def run(self, input: MCPCallerWorkflowInput) -> None:
-        transport = WorkflowTransport(input.endpoint)
-        async with transport.connect() as (read_stream, write_stream):
-            async with ClientSession(read_stream, write_stream) as session:
-                await session.initialize()
+        async with nexusmcp.workflow.MCPClient(input.endpoint).connect() as session:
+            # Session is already initialized
+            list_tools_result = await session.list_tools()
+            assert len(list_tools_result.tools) == 2
+            assert list_tools_result.tools[0].name == "modified-service-name_modified-op-name"
+            assert list_tools_result.tools[1].name == "modified-service-name_op2"
 
-                list_tools_result = await session.list_tools()
-                assert len(list_tools_result.tools) == 2
-                assert list_tools_result.tools[0].name == "modified-service-name_modified-op-name"
-                assert list_tools_result.tools[1].name == "modified-service-name_op2"
-
-                call_result = await session.call_tool("modified-service-name_modified-op-name", {"name": "World"})
-                assert call_result.structuredContent == {"message": "Hello, World"}
+            call_result = await session.call_tool("modified-service-name_modified-op-name", {"name": "World"})
+            assert call_result.structuredContent == {"message": "Hello, World"}
 
 
 @pytest.mark.asyncio