Change server names, overhaul tests to use a custom MCPServer which tracks calls

Author: tconley1428

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 (
-        StatefulTemporalMCPServer,
-        StatelessTemporalMCPServer,
+        StatefulMCPServer,
+        StatelessMCPServer,
     )
 except ImportError:
     pass
@@ -33,8 +33,8 @@
     "OpenAIAgentsPlugin",
     "ModelActivityParameters",
     "workflow",
-    "StatelessTemporalMCPServer",
-    "StatefulTemporalMCPServer",
+    "StatelessMCPServer",
+    "StatefulMCPServer",
     "TestModel",
     "TestModelProvider",
 ]

temporalio/contrib/openai_agents/_mcp.py

@@ -24,7 +24,7 @@
 logger = logging.getLogger(__name__)
 
 
-class _StatelessTemporalMCPServerReference(MCPServer):
+class _StatelessMCPServerReference(MCPServer):
     def __init__(self, server: str, config: Optional[ActivityConfig] = None):
         self._name = server + "-stateless"
         self._config = config or ActivityConfig(
@@ -83,7 +83,7 @@ async def get_prompt(
         )
 
 
-class StatelessTemporalMCPServer:
+class StatelessMCPServer:
     """A stateless MCP server implementation for Temporal workflows.
 
     This class wraps an MCP server to make it stateless by executing each MCP operation
@@ -180,7 +180,7 @@ async def wrapper(*args, **kwargs):
     return wrapper
 
 
-class _StatefulTemporalMCPServerReference(MCPServer, AbstractAsyncContextManager):
+class _StatefulMCPServerReference(MCPServer, AbstractAsyncContextManager):
     def __init__(
         self,
         server: str,
@@ -266,7 +266,7 @@ async def get_prompt(
         )
 
 
-class StatefulTemporalMCPServer:
+class StatefulMCPServer:
     """A stateful MCP server implementation for Temporal workflows.
 
     This class wraps an MCP server to maintain a persistent connection throughout
@@ -343,6 +343,7 @@ async def connect() -> None:
 
                 await worker.run()
             finally:
+                print("Cleanup")
                 await self._server.cleanup()
                 heartbeat_task.cancel()
                 try:

temporalio/contrib/openai_agents/_openai_runner.py

@@ -57,16 +57,16 @@ async def run(
 
         if starting_agent.mcp_servers:
             from temporalio.contrib.openai_agents._mcp import (
-                _StatefulTemporalMCPServerReference,
-                _StatelessTemporalMCPServerReference,
+                _StatefulMCPServerReference,
+                _StatelessMCPServerReference,
             )
 
             for s in starting_agent.mcp_servers:
                 if not isinstance(
                     s,
                     (
-                        _StatelessTemporalMCPServerReference,
-                        _StatefulTemporalMCPServerReference,
+                        _StatelessMCPServerReference,
+                        _StatefulMCPServerReference,
                     ),
                 ):
                     raise ValueError(

temporalio/contrib/openai_agents/_temporal_openai_agents.py

@@ -60,8 +60,8 @@
 
 if typing.TYPE_CHECKING:
     from temporalio.contrib.openai_agents import (
-        StatefulTemporalMCPServer,
-        StatelessTemporalMCPServer,
+        StatefulMCPServer,
+        StatelessMCPServer,
     )
 
 
@@ -201,7 +201,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, StatelessTemporalMCPServer
+        >>> from temporalio.contrib.openai_agents import OpenAIAgentsPlugin, ModelActivityParameters, StatelessMCPServer
         >>> from agents.mcp import MCPServerStdio
         >>> from datetime import timedelta
         >>>
@@ -212,7 +212,7 @@ class OpenAIAgentsPlugin(temporalio.client.Plugin, temporalio.worker.Plugin):
         ... )
         >>>
         >>> # Create MCP servers
-        >>> filesystem_server = StatelessTemporalMCPServer(MCPServerStdio(
+        >>> filesystem_server = StatelessMCPServer(MCPServerStdio(
         ...     name="Filesystem Server",
         ...     params={"command": "npx", "args": ["-y", "@modelcontextprotocol/server-filesystem", "."]}
         ... ))
@@ -239,9 +239,7 @@ def __init__(
         self,
         model_params: Optional[ModelActivityParameters] = None,
         model_provider: Optional[ModelProvider] = None,
-        mcp_servers: Sequence[
-            Union["StatelessTemporalMCPServer", "StatefulTemporalMCPServer"]
-        ] = (),
+        mcp_servers: Sequence[Union["StatelessMCPServer", "StatefulMCPServer"]] = (),
     ) -> None:
         """Initialize the OpenAI agents plugin.
 

temporalio/contrib/openai_agents/workflow.py

@@ -258,10 +258,10 @@ def stateless_mcp_server(
     superior durability guarantees.
     """
     from temporalio.contrib.openai_agents._mcp import (
-        _StatelessTemporalMCPServerReference,
+        _StatelessMCPServerReference,
     )
 
-    return _StatelessTemporalMCPServerReference(name, config)
+    return _StatelessMCPServerReference(name, config)
 
 
 def stateful_mcp_server(
@@ -294,10 +294,10 @@ def stateful_mcp_server(
                        Defaults to 1-hour start-to-close timeout.
     """
     from temporalio.contrib.openai_agents._mcp import (
-        _StatefulTemporalMCPServerReference,
+        _StatefulMCPServerReference,
     )
 
-    return _StatefulTemporalMCPServerReference(name, config, server_session_config)
+    return _StatefulMCPServerReference(name, config, server_session_config)
 
 
 class ToolSerializationError(TemporalError):