Change stateless to a provider model, add caching option

Author: tconley1428

temporalio/contrib/openai_agents/__init__.py

@@ -12,7 +12,7 @@
 try:
     from temporalio.contrib.openai_agents._mcp import (
         StatefulMCPServerProvider,
-        StatelessMCPServer,
+        StatelessMCPServerProvider,
     )
 except ImportError:
     pass
@@ -36,7 +36,7 @@
     "ModelActivityParameters",
     "OpenAIAgentsPlugin",
     "OpenAIPayloadConverter",
-    "StatelessMCPServer",
+    "StatelessMCPServerProvider",
     "StatefulMCPServerProvider",
     "TestModel",
     "TestModelProvider",

temporalio/contrib/openai_agents/_mcp.py

@@ -30,11 +30,18 @@
 
 
 class _StatelessMCPServerReference(MCPServer):
-    def __init__(self, server: str, config: Optional[ActivityConfig] = None):
+    def __init__(
+        self,
+        server: str,
+        config: Optional[ActivityConfig],
+        cache_tools_list: bool,
+    ):
         self._name = server + "-stateless"
         self._config = config or ActivityConfig(
             start_to_close_timeout=timedelta(minutes=1)
         )
+        self._cache_tools_list = cache_tools_list
+        self._tools = None
         super().__init__()
 
     @property
@@ -52,12 +59,17 @@ async def list_tools(
         run_context: Optional[RunContextWrapper[Any]] = None,
         agent: Optional[AgentBase] = None,
     ) -> list[MCPTool]:
-        return await workflow.execute_activity(
+        if self._tools:
+            return self._tools
+        tools = await workflow.execute_activity(
             self.name + "-list-tools",
             args=[],
             result_type=list[MCPTool],
             **self._config,
         )
+        if self._cache_tools_list:
+            self._tools = tools
+        return tools
 
     async def call_tool(
         self, tool_name: str, arguments: Optional[dict[str, Any]]
@@ -88,25 +100,26 @@ async def get_prompt(
         )
 
 
-class StatelessMCPServer:
+class StatelessMCPServerProvider:
     """A stateless MCP server implementation for Temporal workflows.
 
-    This class wraps an MCP server to make it stateless by executing each MCP operation
+    This class wraps a function to create MCP servers to make them stateless by executing each MCP operation
     as a separate Temporal activity. Each operation (list_tools, call_tool, etc.) will
     connect to the underlying server, execute the operation, and then clean up the connection.
 
     This approach will not maintain state across calls. If the desired MCPServer needs persistent state in order to
     function, this cannot be used.
     """
 
-    def __init__(self, server: MCPServer):
+    def __init__(self, server_factory: Callable[[], MCPServer]):
         """Initialize the stateless temporal MCP server.
 
         Args:
-            server: An MCPServer instance
+            server_factory: A function which will produce MCPServer instances. It should return a new server each time
+                so that state is not shared between workflow runs
         """
-        self._server = server
-        self._name = server.name + "-stateless"
+        self._server_factory = server_factory
+        self._name = server_factory().name + "-stateless"
         super().__init__()
 
     @property
@@ -117,39 +130,43 @@ def name(self) -> str:
     def _get_activities(self) -> Sequence[Callable]:
         @activity.defn(name=self.name + "-list-tools")
         async def list_tools() -> list[MCPTool]:
+            server = self._server_factory()
             try:
-                await self._server.connect()
-                return await self._server.list_tools()
+                await server.connect()
+                return await server.list_tools()
             finally:
-                await self._server.cleanup()
+                await server.cleanup()
 
         @activity.defn(name=self.name + "-call-tool")
         async def call_tool(
             tool_name: str, arguments: Optional[dict[str, Any]]
         ) -> CallToolResult:
+            server = self._server_factory()
             try:
-                await self._server.connect()
-                return await self._server.call_tool(tool_name, arguments)
+                await server.connect()
+                return await server.call_tool(tool_name, arguments)
             finally:
-                await self._server.cleanup()
+                await server.cleanup()
 
         @activity.defn(name=self.name + "-list-prompts")
         async def list_prompts() -> ListPromptsResult:
+            server = self._server_factory()
             try:
-                await self._server.connect()
-                return await self._server.list_prompts()
+                await server.connect()
+                return await server.list_prompts()
             finally:
-                await self._server.cleanup()
+                await server.cleanup()
 
         @activity.defn(name=self.name + "-get-prompt")
         async def get_prompt(
             name: str, arguments: Optional[dict[str, Any]]
         ) -> GetPromptResult:
+            server = self._server_factory()
             try:
-                await self._server.connect()
-                return await self._server.get_prompt(name, arguments)
+                await server.connect()
+                return await server.get_prompt(name, arguments)
             finally:
-                await self._server.cleanup()
+                await server.cleanup()
 
         return list_tools, call_tool, list_prompts, get_prompt
 
@@ -189,8 +206,9 @@ class _StatefulMCPServerReference(MCPServer, AbstractAsyncContextManager):
     def __init__(
         self,
         server: str,
-        config: Optional[ActivityConfig] = None,
-        server_session_config: Optional[ActivityConfig] = None,
+        config: Optional[ActivityConfig],
+        server_session_config: Optional[ActivityConfig],
+        cache_tools_list: bool,
     ):
         self._name = server + "-stateful"
         self._config = config or ActivityConfig(
@@ -201,6 +219,8 @@ def __init__(
             start_to_close_timeout=timedelta(hours=1),
         )
         self._connect_handle: Optional[ActivityHandle] = None
+        self._cache_tools_list = cache_tools_list
+        self._tools = None
         super().__init__()
 
     @property
@@ -239,16 +259,22 @@ async def list_tools(
         run_context: Optional[RunContextWrapper[Any]] = None,
         agent: Optional[AgentBase] = None,
     ) -> list[MCPTool]:
+        if self._tools:
+            return self._tools
+
         if not self._connect_handle:
             raise ApplicationError(
                 "Stateful MCP Server not connected. Call connect first."
             )
-        return await workflow.execute_activity(
+        tools = await workflow.execute_activity(
             self.name + "-list-tools",
             args=[],
             result_type=list[MCPTool],
             **self._config,
         )
+        if self._cache_tools_list:
+            self._tools = tools
+        return tools
 
     @_handle_worker_failure
     async def call_tool(
@@ -361,7 +387,7 @@ async def heartbeat_every(delay: float, *details: Any) -> None:
                 await asyncio.sleep(delay)
                 activity.heartbeat(*details)
 
-        @activity.defn(name=self._name + "-server-session")
+        @activity.defn(name=self.name + "-server-session")
         async def connect() -> None:
             heartbeat_task = asyncio.create_task(heartbeat_every(30))
 

temporalio/contrib/openai_agents/_temporal_openai_agents.py

@@ -65,7 +65,7 @@
 if typing.TYPE_CHECKING:
     from temporalio.contrib.openai_agents import (
         StatefulMCPServerProvider,
-        StatelessMCPServer,
+        StatelessMCPServerProvider,
     )
 
 
@@ -204,7 +204,7 @@ class OpenAIAgentsPlugin(temporalio.client.Plugin, temporalio.worker.Plugin):
     Example:
         >>> from temporalio.client import Client
         >>> from temporalio.worker import Worker
-        >>> from temporalio.contrib.openai_agents import OpenAIAgentsPlugin, ModelActivityParameters, StatelessMCPServer
+        >>> from temporalio.contrib.openai_agents import OpenAIAgentsPlugin, ModelActivityParameters, StatelessMCPServerProvider
         >>> from agents.mcp import MCPServerStdio
         >>> from datetime import timedelta
         >>>
@@ -215,7 +215,7 @@ class OpenAIAgentsPlugin(temporalio.client.Plugin, temporalio.worker.Plugin):
         ... )
         >>>
         >>> # Create MCP servers
-        >>> filesystem_server = StatelessMCPServer(MCPServerStdio(
+        >>> filesystem_server = StatelessMCPServerProvider(MCPServerStdio(
         ...     name="Filesystem Server",
         ...     params={"command": "npx", "args": ["-y", "@modelcontextprotocol/server-filesystem", "."]}
         ... ))
@@ -243,7 +243,7 @@ def __init__(
         model_params: Optional[ModelActivityParameters] = None,
         model_provider: Optional[ModelProvider] = None,
         mcp_servers: Sequence[
-            Union["StatelessMCPServer", "StatefulMCPServerProvider"]
+            Union["StatelessMCPServerProvider", "StatefulMCPServerProvider"]
         ] = (),
     ) -> None:
         """Initialize the OpenAI agents plugin.

temporalio/contrib/openai_agents/workflow.py

@@ -242,7 +242,9 @@ async def run_operation(ctx: RunContextWrapper[Any], input: str) -> Any:
 
 
 def stateless_mcp_server(
-    name: str, config: Optional[ActivityConfig] = None
+    name: str,
+    config: Optional[ActivityConfig] = None,
+    cache_tools_list: bool = False,
 ) -> "MCPServer":
     """A stateless MCP server implementation for Temporal workflows.
 
@@ -256,18 +258,25 @@ def stateless_mcp_server(
     This approach is suitable for simple use cases where connection overhead is acceptable
     and you don't need to maintain state between operations. It should be preferred to stateful when possible due to its
     superior durability guarantees.
+
+    Args:
+        name: A string name for the server. Should match that provided in the plugin.
+        config: Optional activity configuration for MCP operation activities.
+               Defaults to 1-minute start-to-close timeout.
+        cache_tools_list: If true, the list of tools will be cached for the duration of the server
     """
     from temporalio.contrib.openai_agents._mcp import (
         _StatelessMCPServerReference,
     )
 
-    return _StatelessMCPServerReference(name, config)
+    return _StatelessMCPServerReference(name, config, cache_tools_list)
 
 
 def stateful_mcp_server(
     name: str,
     config: Optional[ActivityConfig] = None,
     server_session_config: Optional[ActivityConfig] = None,
+    cache_tools_list: bool = False,
 ) -> AbstractAsyncContextManager["MCPServer"]:
     """A stateful MCP server implementation for Temporal workflows.
 
@@ -292,12 +301,15 @@ def stateful_mcp_server(
                Defaults to 1-minute start-to-close and 30-second schedule-to-start timeouts.
         server_session_config: Optional activity configuration for the connection activity.
                        Defaults to 1-hour start-to-close timeout.
+        cache_tools_list: If true, the list of tools will be cached for the duration of the server
     """
     from temporalio.contrib.openai_agents._mcp import (
         _StatefulMCPServerReference,
     )
 
-    return _StatefulMCPServerReference(name, config, server_session_config)
+    return _StatefulMCPServerReference(
+        name, config, server_session_config, cache_tools_list
+    )
 
 
 class ToolSerializationError(TemporalError):