Add docstrings

Author: tconley1428

temporalio/contrib/openai_agents/_mcp.py

@@ -23,9 +23,26 @@
 
 
 class StatelessTemporalMCPServer(MCPServer):
+    """A stateless MCP server implementation for Temporal workflows.
+
+    This class wraps an MCP server to make it 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 is suitable for simple use cases where connection overhead is acceptable
+    and you don't need to maintain state between operations.
+    """
+
     def __init__(
         self, server: Union[MCPServer, str], config: Optional[ActivityConfig] = None
     ):
+        """Initialize the stateless temporal MCP server.
+
+        Args:
+            server: Either an MCPServer instance or a string name for the server.
+            config: Optional activity configuration for Temporal activities. Defaults to
+                   1-minute start-to-close timeout if not provided.
+        """
         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(
@@ -35,19 +52,49 @@ def __init__(
 
     @property
     def name(self) -> str:
+        """Get the server name with '-stateless' suffix.
+
+        Returns:
+            The server name with '-stateless' appended.
+        """
         return self._name
 
     async def connect(self) -> None:
+        """Connect to the MCP server.
+
+        For stateless servers, this is a no-op since connections are made
+        on a per-operation basis.
+        """
         pass
 
     async def cleanup(self) -> None:
+        """Clean up the MCP server connection.
+
+        For stateless servers, this is a no-op since connections are cleaned
+        up after each operation.
+        """
         pass
 
     async def list_tools(
         self,
         run_context: Optional[RunContextWrapper[Any]] = None,
         agent: Optional[AgentBase] = None,
     ) -> list[MCPTool]:
+        """List available tools from the MCP server.
+
+        This method executes a Temporal activity to connect to the MCP server,
+        retrieve the list of available tools, and clean up the connection.
+
+        Args:
+            run_context: Optional run context wrapper (unused in stateless mode).
+            agent: Optional agent base (unused in stateless mode).
+
+        Returns:
+            A list of available MCP tools.
+
+        Raises:
+            ActivityError: If the underlying Temporal activity fails.
+        """
         return await workflow.execute_activity(
             self.name + "-list-tools",
             args=[],
@@ -58,6 +105,21 @@ async def list_tools(
     async def call_tool(
         self, tool_name: str, arguments: Optional[dict[str, Any]]
     ) -> CallToolResult:
+        """Call a specific tool on the MCP server.
+
+        This method executes a Temporal activity to connect to the MCP server,
+        call the specified tool with the given arguments, and clean up the connection.
+
+        Args:
+            tool_name: The name of the tool to call.
+            arguments: Optional dictionary of arguments to pass to the tool.
+
+        Returns:
+            The result of the tool call.
+
+        Raises:
+            ActivityError: If the underlying Temporal activity fails.
+        """
         return await workflow.execute_activity(
             self.name + "-call-tool",
             args=[tool_name, arguments],
@@ -66,6 +128,17 @@ async def call_tool(
         )
 
     async def list_prompts(self) -> ListPromptsResult:
+        """List available prompts from the MCP server.
+
+        This method executes a Temporal activity to connect to the MCP server,
+        retrieve the list of available prompts, and clean up the connection.
+
+        Returns:
+            A list of available prompts.
+
+        Raises:
+            ActivityError: If the underlying Temporal activity fails.
+        """
         return await workflow.execute_activity(
             self.name + "-list-prompts",
             args=[],
@@ -76,6 +149,21 @@ async def list_prompts(self) -> ListPromptsResult:
     async def get_prompt(
         self, name: str, arguments: Optional[dict[str, Any]] = None
     ) -> GetPromptResult:
+        """Get a specific prompt from the MCP server.
+
+        This method executes a Temporal activity to connect to the MCP server,
+        retrieve the specified prompt with optional arguments, and clean up the connection.
+
+        Args:
+            name: The name of the prompt to retrieve.
+            arguments: Optional dictionary of arguments for the prompt.
+
+        Returns:
+            The prompt result.
+
+        Raises:
+            ActivityError: If the underlying Temporal activity fails.
+        """
         return await workflow.execute_activity(
             self.name + "-get-prompt",
             args=[name, arguments],
@@ -84,6 +172,17 @@ async def get_prompt(
         )
 
     def get_activities(self) -> Sequence[Callable]:
+        """Get the Temporal activities for this MCP server.
+
+        Creates and returns the Temporal activity functions that handle MCP operations.
+        Each activity manages its own connection lifecycle (connect -> operate -> cleanup).
+
+        Returns:
+            A sequence of Temporal activity functions.
+
+        Raises:
+            ValueError: If no MCP server instance was provided during initialization.
+        """
         server = self.server
         if server is None:
             raise ValueError(
@@ -130,12 +229,35 @@ async def get_prompt(
 
 
 class StatefulTemporalMCPServer(MCPServer):
+    """A stateful MCP server implementation for Temporal workflows.
+
+    This class wraps an MCP server to maintain a persistent connection throughout
+    the workflow execution. It creates a dedicated worker that stays connected to
+    the MCP server and processes operations on a dedicated task queue.
+
+    This approach is more efficient for workflows that make multiple MCP calls,
+    as it avoids connection overhead, but requires more resources to maintain
+    the persistent connection and worker.
+
+    The caller will have to handle cases where the dedicated worker fails, as Temporal is
+    unable to seamlessly recreate any lost state in that case.
+    """
+
     def __init__(
         self,
         server: Union[MCPServer, str],
         config: Optional[ActivityConfig] = None,
         connect_config: Optional[ActivityConfig] = None,
     ):
+        """Initialize the stateful temporal MCP server.
+
+        Args:
+            server: Either an MCPServer instance or a string name for the server.
+            config: Optional activity configuration for MCP operation activities.
+                   Defaults to 1-minute start-to-close and 30-second schedule-to-start timeouts.
+            connect_config: Optional activity configuration for the connection activity.
+                           Defaults to 1-hour start-to-close timeout.
+        """
         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(
@@ -150,9 +272,20 @@ def __init__(
 
     @property
     def name(self) -> str:
+        """Get the server name with '-stateful' suffix.
+
+        Returns:
+            The server name with '-stateful' appended.
+        """
         return self._name
 
     async def connect(self) -> None:
+        """Connect to the MCP server and start the dedicated worker.
+
+        This method creates a dedicated task queue for this workflow and starts
+        a long-running activity that maintains the connection and runs a worker
+        to handle MCP operations.
+        """
         self.config["task_queue"] = workflow.info().workflow_id + "-" + self.name
         self._connect_handle = workflow.start_activity(
             self.name + "-connect",
@@ -161,21 +294,55 @@ async def connect(self) -> None:
         )
 
     async def cleanup(self) -> None:
+        """Clean up the MCP server connection.
+
+        This method cancels the long-running connection activity, which will
+        cause the dedicated worker to shut down and the MCP server connection
+        to be closed.
+        """
         if self._connect_handle:
             self._connect_handle.cancel()
 
     async def __aenter__(self):
+        """Async context manager entry point.
+
+        Returns:
+            This server instance after connecting.
+        """
         await self.connect()
         return self
 
     async def __aexit__(self, exc_type, exc_value, traceback):
+        """Async context manager exit point.
+
+        Args:
+            exc_type: Exception type if an exception occurred.
+            exc_value: Exception value if an exception occurred.
+            traceback: Exception traceback if an exception occurred.
+        """
         await self.cleanup()
 
     async def list_tools(
         self,
         run_context: Optional[RunContextWrapper[Any]] = None,
         agent: Optional[AgentBase] = None,
     ) -> list[MCPTool]:
+        """List available tools from the MCP server.
+
+        This method executes a Temporal activity on the dedicated task queue
+        to retrieve the list of available tools from the persistent MCP connection.
+
+        Args:
+            run_context: Optional run context wrapper (unused in stateful mode).
+            agent: Optional agent base (unused in stateful mode).
+
+        Returns:
+            A list of available MCP tools.
+
+        Raises:
+            ApplicationError: If the MCP worker fails to schedule or heartbeat.
+            ActivityError: If the underlying Temporal activity fails.
+        """
         try:
             logger.info("Executing list-tools: %s", self.config)
             return await workflow.execute_activity(
@@ -208,6 +375,21 @@ async def list_tools(
     async def call_tool(
         self, tool_name: str, arguments: Optional[dict[str, Any]]
     ) -> CallToolResult:
+        """Call a specific tool on the MCP server.
+
+        This method executes a Temporal activity on the dedicated task queue
+        to call the specified tool using the persistent MCP connection.
+
+        Args:
+            tool_name: The name of the tool to call.
+            arguments: Optional dictionary of arguments to pass to the tool.
+
+        Returns:
+            The result of the tool call.
+
+        Raises:
+            ActivityError: If the underlying Temporal activity fails.
+        """
         return await workflow.execute_activity(
             self.name + "-call-tool",
             args=[tool_name, arguments],
@@ -216,6 +398,17 @@ async def call_tool(
         )
 
     async def list_prompts(self) -> ListPromptsResult:
+        """List available prompts from the MCP server.
+
+        This method executes a Temporal activity on the dedicated task queue
+        to retrieve the list of available prompts from the persistent MCP connection.
+
+        Returns:
+            A list of available prompts.
+
+        Raises:
+            ActivityError: If the underlying Temporal activity fails.
+        """
         return await workflow.execute_activity(
             self.name + "-list-prompts",
             args=[],
@@ -226,6 +419,21 @@ async def list_prompts(self) -> ListPromptsResult:
     async def get_prompt(
         self, name: str, arguments: Optional[dict[str, Any]] = None
     ) -> GetPromptResult:
+        """Get a specific prompt from the MCP server.
+
+        This method executes a Temporal activity on the dedicated task queue
+        to retrieve the specified prompt using the persistent MCP connection.
+
+        Args:
+            name: The name of the prompt to retrieve.
+            arguments: Optional dictionary of arguments for the prompt.
+
+        Returns:
+            The prompt result.
+
+        Raises:
+            ActivityError: If the underlying Temporal activity fails.
+        """
         return await workflow.execute_activity(
             self.name + "-get-prompt",
             args=[name, arguments],
@@ -234,6 +442,18 @@ async def get_prompt(
         )
 
     def get_activities(self) -> Sequence[Callable]:
+        """Get the Temporal activities for this stateful MCP server.
+
+        Creates and returns the Temporal activity functions that handle MCP operations
+        and connection management. This includes a long-running connect activity that
+        maintains the MCP connection and runs a dedicated worker.
+
+        Returns:
+            A sequence containing the connect activity function.
+
+        Raises:
+            ValueError: If no MCP server instance was provided during initialization.
+        """
         server = self.server
         if server is None:
             raise ValueError(