Test stateful session

Author: dandavison

tests/test_stateful_session.py

@@ -0,0 +1,127 @@
+import asyncio
+import uuid
+from dataclasses import dataclass
+
+import pytest
+from mcp import ClientSession
+from nexusrpc import Operation, service
+from nexusrpc.handler import StartOperationContext, service_handler, sync_operation
+from temporalio import nexus, workflow
+from temporalio.api.nexus.v1 import EndpointSpec, EndpointTarget
+from temporalio.api.operatorservice.v1 import CreateNexusEndpointRequest
+from temporalio.client import WithStartWorkflowOperation
+from temporalio.common import WorkflowIDConflictPolicy
+from temporalio.contrib.pydantic import pydantic_data_converter
+from temporalio.testing import WorkflowEnvironment
+from temporalio.worker import Worker
+
+from nexusmcp import MCPServiceHandler, WorkflowTransport
+
+mcp_service = MCPServiceHandler()
+
+
+@dataclass
+class AppendInput:
+    session_id: str
+    value: int
+
+
+@service
+class TestService:
+    append: Operation[AppendInput, list[int]]
+
+
+@workflow.defn(sandboxed=False)
+class AppendWorkflow:
+    def __init__(self):
+        self.data = []
+
+    @workflow.run
+    async def run(self) -> None:
+        await asyncio.Event().wait()
+
+    @workflow.update
+    async def append(self, input: int) -> list[int]:
+        self.data.append(input)
+        return self.data
+
+
+@mcp_service.register
+@service_handler(service=TestService)
+class TestServiceHandler:
+    @sync_operation
+    async def append(self, ctx: StartOperationContext, input: AppendInput) -> list[int]:
+        # TODO: start workflow in a call made during initialize()?
+        # wf = nexus.client().get_workflow_handle_for(AppendWorkflow.run, input.session_id)
+        with_start = WithStartWorkflowOperation(
+            AppendWorkflow.run,
+            id=input.session_id,
+            task_queue=nexus.info().task_queue,
+            id_conflict_policy=WorkflowIDConflictPolicy.USE_EXISTING,
+        )
+        return await nexus.client().execute_update_with_start_workflow(
+            AppendWorkflow.append,
+            input.value,
+            start_workflow_operation=with_start,
+        )
+
+
+@dataclass
+class MCPCallerWorkflowInput:
+    endpoint: str
+
+
+# sandbox disabled due to use of ThreadLocal by sniffio
+# TODO: make this unnecessary
+@workflow.defn(sandboxed=False)
+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()
+
+                list_tools_result = await session.list_tools()
+                assert len(list_tools_result.tools) == 1
+                assert list_tools_result.tools[0].name == "TestService/append"
+
+                call_result = await session.call_tool("TestService/append", {"session_id": "123", "value": 1})
+                assert call_result.content[0].text == "[1]"
+
+                call_result = await session.call_tool("TestService/append", {"session_id": "123", "value": 2})
+                assert call_result.content[0].text == "[1, 2]"
+
+
+@pytest.mark.asyncio
+async def test_workflow_caller() -> None:
+    endpoint_name = "endpoint"
+    task_queue = "handler-queue"
+
+    async with await WorkflowEnvironment.start_local(data_converter=pydantic_data_converter) as env:
+        await env.client.operator_service.create_nexus_endpoint(
+            CreateNexusEndpointRequest(
+                spec=EndpointSpec(
+                    name=endpoint_name,
+                    target=EndpointTarget(
+                        worker=EndpointTarget.Worker(
+                            namespace=env.client.namespace,
+                            task_queue=task_queue,
+                        )
+                    ),
+                )
+            )
+        )
+
+        async with Worker(
+            env.client,
+            task_queue=task_queue,
+            workflows=[MCPCallerWorkflow, AppendWorkflow],
+            nexus_service_handlers=[mcp_service, TestServiceHandler()],
+        ):
+            await env.client.execute_workflow(
+                MCPCallerWorkflow.run,
+                arg=MCPCallerWorkflowInput(endpoint=endpoint_name),
+                id=str(uuid.uuid4()),
+                task_queue=task_queue,
+            )