Test stateful session

.github/workflows/ci.yml

@@ -0,0 +1,34 @@
+name: CI
+
+on:
+  pull_request:
+  push:
+    branches: [ main ]
+
+jobs:
+  lint-test-docs:
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        python-version: ['3.13']
+        os: [ubuntu-latest]
+
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v6
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install dependencies
+        run: |
+          uv sync
+
+      - name: Lint and type check
+        run: uv run poe lint
+
+      - name: Run tests
+        run: |
+          uv run pytest

nexusmcp/__init__.py

@@ -4,6 +4,5 @@
     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"]

nexusmcp/service_handler.py

@@ -9,11 +9,11 @@
 - Decorators for controlling which operations are exposed as MCP tools
 """
 
+import logging
+import re
 from collections.abc import Callable
 from dataclasses import dataclass
 from typing import Any
-import re
-import logging
 
 import mcp.types
 import nexusrpc

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)

pyproject.toml

@@ -14,10 +14,12 @@ dependencies = [
 
 [dependency-groups]
 dev = [
-    "pytest>=8.4.1",
+    "mypy>=1.13.0",
+    "poethepoet>=0.35.0",
+    "pyright>=1.1",
     "pytest-asyncio>=1.1.0",
+    "pytest>=8.4.1",
     "ruff>=0.8.0",
-    "mypy>=1.13.0",
 ]
 
 [tool.pytest.ini_options]
@@ -41,5 +43,17 @@ indent-style = "space"
 python_version = "3.13"
 strict = true
 
+[tool.poe.tasks]
+lint = [
+  {cmd = "uv run pyright"},
+  {cmd = "uv run mypy --check-untyped-defs ."},
+  {cmd = "uv run ruff check --select I"},
+  {cmd = "uv run ruff format --check"},
+]
+format = [
+  {cmd = "uv run ruff check --select I --fix"},
+  {cmd = "uv run ruff format"},
+]
+
 [tool.uv.sources]
 temporalio = { git = "https://github.com/temporalio/sdk-python" }

tests/test_inbound_gateway.py

@@ -66,8 +66,8 @@ async def test_inbound_gateway() -> None:
 
 
 async def call_tool(
-    read_stream: anyio.streams.memory.MemoryObjectReceiveStream[SessionMessage],
-    write_stream: anyio.streams.memory.MemoryObjectSendStream[SessionMessage],
+    read_stream: anyio.streams.memory.MemoryObjectReceiveStream[SessionMessage],  # pyright: ignore[reportAttributeAccessIssue]
+    write_stream: anyio.streams.memory.MemoryObjectSendStream[SessionMessage],  # pyright: ignore[reportAttributeAccessIssue]
 ) -> None:
     """Test MCP client connecting via memory streams and calling tools."""
     async with ClientSession(read_stream, write_stream) as session:
@@ -79,8 +79,8 @@ async def call_tool(
         print(f"Available tools: {[tool.name for tool in list_tools_result.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"
+        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"})
+        call_result = await session.call_tool("modified-service-name_modified-op-name", {"name": "World"})
         assert call_result.structuredContent == {"message": "Hello, World"}

tests/test_service.py

@@ -1,6 +1,6 @@
 import pytest
-
 from nexusrpc.handler import StartOperationContext
+
 from .service import MyInput, TestServiceHandler, mcp_service
 
 
@@ -27,8 +27,8 @@ async def test_list_tools() -> None:
         None,
     )
     assert len(tools) == 2
-    assert tools[0].name == "modified-service-name/modified-op-name"
-    assert tools[1].name == "modified-service-name/op2"
+    assert tools[0].name == "modified-service-name_modified-op-name"
+    assert tools[1].name == "modified-service-name_op2"
     assert tools[0].description == "This is a test operation."
     assert tools[1].description == "This is also a test operation."
     assert tools[0].inputSchema == MyInput.model_json_schema()