Include stateful option

tconley1428 · tconley1428 · commit 0fd788676782 · 2025-08-11T13:03:23.000-07:00
diff --git a/temporalio/contrib/openai_agents/__init__.py b/temporalio/contrib/openai_agents/__init__.py
@@ -11,8 +11,8 @@
 # Best Effort mcp, as it is not supported on Python 3.9
 try:
     from temporalio.contrib.openai_agents._mcp import (
-        TemporalMCPServer,
-        TemporalMCPServerWorkflowShim,
+        StatefulTemporalMCPServer,
+        StatelessTemporalMCPServer,
     )
 except ImportError:
     pass
@@ -33,8 +33,8 @@
     "OpenAIAgentsPlugin",
     "ModelActivityParameters",
     "workflow",
-    "TemporalMCPServer",
-    "TemporalMCPServerWorkflowShim",
+    "StatelessTemporalMCPServer",
+    "StatefulTemporalMCPServer",
     "TestModel",
     "TestModelProvider",
 ]
diff --git a/temporalio/contrib/openai_agents/_mcp.py b/temporalio/contrib/openai_agents/_mcp.py
@@ -1,3 +1,6 @@
+import asyncio
+import logging
+import uuid
 from datetime import timedelta
 from typing import Any, Callable, Optional, Sequence, Union
 
@@ -8,13 +11,26 @@
 from mcp.types import CallToolResult  # type:ignore
 
 from temporalio import activity, workflow
-from temporalio.workflow import ActivityConfig
+from temporalio.api.enums.v1.workflow_pb2 import (
+    TIMEOUT_TYPE_HEARTBEAT,
+    TIMEOUT_TYPE_SCHEDULE_TO_START,
+)
+from temporalio.exceptions import ActivityError, ApplicationError
+from temporalio.worker import PollerBehaviorSimpleMaximum, Worker
+from temporalio.workflow import ActivityConfig, ActivityHandle
+
+logger = logging.getLogger(__name__)
+
 
 class StatelessTemporalMCPServer(MCPServer):
-    def __init__(self, server: Union[MCPServer, str], config: Optional[ActivityConfig] = None):
+    def __init__(
+        self, server: Union[MCPServer, str], config: Optional[ActivityConfig] = None
+    ):
         self.server = server if isinstance(server, MCPServer) else None
         self._name = (server if isinstance(server, str) else server.name) + "-stateless"
-        self.config = config or ActivityConfig(start_to_close_timeout=timedelta(minutes=1))
+        self.config = config or ActivityConfig(
+            start_to_close_timeout=timedelta(minutes=1)
+        )
         super().__init__()
 
     @property
@@ -68,32 +84,209 @@ async def get_prompt(
         )
 
     def get_activities(self) -> Sequence[Callable]:
-        if self.server is None:
-            raise ValueError("A full MCPServer implementation should have been provided when adding a server to the worker.")
+        server = self.server
+        if server is None:
+            raise ValueError(
+                "A full MCPServer implementation should have been provided when adding a server to the worker."
+            )
 
         @activity.defn(name=self.name + "-list-tools")
         async def list_tools() -> list[MCPTool]:
-            activity.logger.info("Listing tools in activity")
-            async with self.server:
-                return await self.server.list_tools()
+            try:
+                await server.connect()
+                return await server.list_tools()
+            finally:
+                await server.cleanup()
 
         @activity.defn(name=self.name + "-call-tool")
         async def call_tool(
             tool_name: str, arguments: Optional[dict[str, Any]]
         ) -> CallToolResult:
-            async with self.server:
-                return await self.server.call_tool(tool_name, arguments)
+            try:
+                await server.connect()
+                return await server.call_tool(tool_name, arguments)
+            finally:
+                await server.cleanup()
 
         @activity.defn(name=self.name + "-list-prompts")
         async def list_prompts() -> ListPromptsResult:
-            async with self.server:
-                return await self.server.list_prompts()
+            try:
+                await server.connect()
+                return await server.list_prompts()
+            finally:
+                await server.cleanup()
 
         @activity.defn(name=self.name + "-get-prompt")
         async def get_prompt(
             name: str, arguments: Optional[dict[str, Any]]
         ) -> GetPromptResult:
-            async with self.server:
-                return await self.server.get_prompt(name, arguments)
+            try:
+                await server.connect()
+                return await server.get_prompt(name, arguments)
+            finally:
+                await server.cleanup()
 
         return list_tools, call_tool, list_prompts, get_prompt
+
+
+class StatefulTemporalMCPServer(MCPServer):
+    def __init__(
+        self,
+        server: Union[MCPServer, str],
+        config: Optional[ActivityConfig] = None,
+        connect_config: Optional[ActivityConfig] = None,
+    ):
+        self.server = server if isinstance(server, MCPServer) else None
+        self._name = (server if isinstance(server, str) else server.name) + "-stateful"
+        self.config = config or ActivityConfig(
+            start_to_close_timeout=timedelta(minutes=1),
+            schedule_to_start_timeout=timedelta(seconds=30),
+        )
+        self.connect_config = connect_config or ActivityConfig(
+            start_to_close_timeout=timedelta(hours=1),
+        )
+        self._connect_handle: Optional[ActivityHandle] = None
+        super().__init__()
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    async def connect(self) -> None:
+        self.config["task_queue"] = workflow.info().workflow_id + "-" + self.name
+        self._connect_handle = workflow.start_activity(
+            self.name + "-connect",
+            args=[],
+            **self.connect_config,
+        )
+
+    async def cleanup(self) -> None:
+        if self._connect_handle:
+            self._connect_handle.cancel()
+
+    async def __aenter__(self):
+        await self.connect()
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, traceback):
+        await self.cleanup()
+
+    async def list_tools(
+        self,
+        run_context: Optional[RunContextWrapper[Any]] = None,
+        agent: Optional[AgentBase] = None,
+    ) -> list[MCPTool]:
+        try:
+            logger.info("Executing list-tools: %s", self.config)
+            return await workflow.execute_activity(
+                self.name + "-list-tools",
+                args=[],
+                result_type=list[MCPTool],
+                **self.config,
+            )
+        except ActivityError as e:
+            failure = e.failure
+            if failure:
+                cause = failure.cause
+                if cause:
+                    if (
+                        cause.timeout_failure_info.timeout_type
+                        == TIMEOUT_TYPE_SCHEDULE_TO_START
+                    ):
+                        raise ApplicationError(
+                            "MCP Stateful Server Worker failed to schedule activity."
+                        ) from e
+                    if (
+                        cause.timeout_failure_info.timeout_type
+                        == TIMEOUT_TYPE_HEARTBEAT
+                    ):
+                        raise ApplicationError(
+                            "MCP Stateful Server Worker failed to heartbeat."
+                        ) from e
+            raise e
+
+    async def call_tool(
+        self, tool_name: str, arguments: Optional[dict[str, Any]]
+    ) -> CallToolResult:
+        return await workflow.execute_activity(
+            self.name + "-call-tool",
+            args=[tool_name, arguments],
+            result_type=CallToolResult,
+            **self.config,
+        )
+
+    async def list_prompts(self) -> ListPromptsResult:
+        return await workflow.execute_activity(
+            self.name + "-list-prompts",
+            args=[],
+            result_type=ListPromptsResult,
+            **self.config,
+        )
+
+    async def get_prompt(
+        self, name: str, arguments: Optional[dict[str, Any]] = None
+    ) -> GetPromptResult:
+        return await workflow.execute_activity(
+            self.name + "-get-prompt",
+            args=[name, arguments],
+            result_type=GetPromptResult,
+            **self.config,
+        )
+
+    def get_activities(self) -> Sequence[Callable]:
+        server = self.server
+        if server is None:
+            raise ValueError(
+                "A full MCPServer implementation should have been provided when adding a server to the worker."
+            )
+
+        @activity.defn(name=self.name + "-list-tools")
+        async def list_tools() -> list[MCPTool]:
+            return await server.list_tools()
+
+        @activity.defn(name=self.name + "-call-tool")
+        async def call_tool(
+            tool_name: str, arguments: Optional[dict[str, Any]]
+        ) -> CallToolResult:
+            return await server.call_tool(tool_name, arguments)
+
+        @activity.defn(name=self.name + "-list-prompts")
+        async def list_prompts() -> ListPromptsResult:
+            return await server.list_prompts()
+
+        @activity.defn(name=self.name + "-get-prompt")
+        async def get_prompt(
+            name: str, arguments: Optional[dict[str, Any]]
+        ) -> GetPromptResult:
+            return await server.get_prompt(name, arguments)
+
+        async def heartbeat_every(delay: float, *details: Any) -> None:
+            """Heartbeat every so often while not cancelled"""
+            while True:
+                await asyncio.sleep(delay)
+                activity.heartbeat(*details)
+
+        @activity.defn(name=self.name + "-connect")
+        async def connect() -> None:
+            logger.info("Connect activity")
+            heartbeat_task = asyncio.create_task(heartbeat_every(30))
+            try:
+                await server.connect()
+
+                worker = Worker(
+                    activity.client(),
+                    task_queue=activity.info().workflow_id + "-" + self.name,
+                    activities=[list_tools, call_tool, list_prompts, get_prompt],
+                    activity_task_poller_behavior=PollerBehaviorSimpleMaximum(1),
+                )
+
+                await worker.run()
+            finally:
+                await server.cleanup()
+                heartbeat_task.cancel()
+                try:
+                    await heartbeat_task
+                except asyncio.CancelledError:
+                    pass
+
+        return (connect,)
diff --git a/temporalio/contrib/openai_agents/_temporal_openai_agents.py b/temporalio/contrib/openai_agents/_temporal_openai_agents.py
@@ -1,8 +1,8 @@
 """Initialize Temporal OpenAI Agents overrides."""
 
 import dataclasses
-import typing
-from contextlib import AsyncExitStack, asynccontextmanager, contextmanager
+import warnings
+from contextlib import asynccontextmanager, contextmanager
 from datetime import timedelta
 from typing import AsyncIterator, Callable, Optional, Sequence, Union
 
@@ -28,6 +28,7 @@
 import temporalio.worker
 from temporalio.client import ClientConfig
 from temporalio.contrib.openai_agents._invoke_model_activity import ModelActivity
+
 from temporalio.contrib.openai_agents._model_parameters import ModelActivityParameters
 from temporalio.contrib.openai_agents._openai_runner import TemporalOpenAIRunner
 from temporalio.contrib.openai_agents._temporal_trace_provider import (
@@ -52,10 +53,15 @@
 )
 from temporalio.worker.workflow_sandbox import SandboxedWorkflowRunner
 
-from agents.mcp import MCPServer
-
-from temporalio.contrib.openai_agents._mcp import StatelessTemporalMCPServer
-
+# Unsupported on python 3.9
+try:
+    from agents.mcp import MCPServer
+    from temporalio.contrib.openai_agents._mcp import (
+        StatefulTemporalMCPServer,
+        StatelessTemporalMCPServer,
+    )
+except ImportError:
+    pass
 
 @contextmanager
 def set_open_ai_agent_temporal_overrides(
@@ -165,18 +171,6 @@ def __init__(self) -> None:
         super().__init__(ToJsonOptions(exclude_unset=True))
 
 
-def _transform_mcp_servers(mcp_servers: Sequence[MCPServer]) -> list[MCPServer]:
-    def _transform_mcp_server(server: MCPServer) -> MCPServer:
-        if isinstance(server, StatelessTemporalMCPServer):
-            return server
-        else:
-            raise TypeError(f"Unsupported mcp server type {type(server)}")
-    return [
-        _transform_mcp_server(server)
-        for server in mcp_servers
-    ]
-
-
 class OpenAIAgentsPlugin(temporalio.client.Plugin, temporalio.worker.Plugin):
     """Temporal plugin for integrating OpenAI agents with Temporal workflows.
 
@@ -281,7 +275,18 @@ def __init__(
         self._model_provider = model_provider
 
         if mcp_servers:
-            self._mcp_servers = _transform_mcp_servers(mcp_servers)
+            def _transform_mcp_server(server: "MCPServer") -> "MCPServer":
+                if not (
+                        isinstance(server, StatelessTemporalMCPServer)
+                        or isinstance(server, StatefulTemporalMCPServer)
+                ):
+                    warnings.warn(
+                        f"Unsupported mcp server type {type(server)} is not guaranteed to behave reasonably."
+                    )
+
+                return server
+
+            self._mcp_servers = [_transform_mcp_server(server) for server in mcp_servers]
         else:
             self._mcp_servers = []
 
@@ -335,7 +340,11 @@ def configure_worker(self, config: WorkerConfig) -> WorkerConfig:
         ]
         new_activities = [ModelActivity(self._model_provider).invoke_model_activity]
         for mcp_server in self._mcp_servers:
-            new_activities.extend(mcp_server.get_activities())
+            if hasattr(mcp_server, "get_activities"):
+                get_activities: Callable[[], Sequence[Callable]] = getattr(
+                    mcp_server, "get_activities"
+                )
+                new_activities.extend(get_activities())
         config["activities"] = list(config.get("activities") or []) + new_activities
 
         runner = config.get("workflow_runner")
diff --git a/tests/contrib/openai_agents/test_openai.py b/tests/contrib/openai_agents/test_openai.py
