PR feedback - exposing SimplePlugin type

Author: tconley1428

temporalio/client.py

@@ -203,7 +203,7 @@ async def connect(
         def make_lambda(plugin, next):
             return lambda config: plugin.connect_service_client(config, next)
 
-        next_function = lambda config: ServiceClient.connect(config)
+        next_function = ServiceClient.connect
         for plugin in reversed(plugins):
             next_function = make_lambda(plugin, next_function)
 

temporalio/contrib/openai_agents/_temporal_openai_agents.py

@@ -22,11 +22,7 @@
 from agents.run import get_default_agent_runner, set_default_agent_runner
 from agents.tracing import get_trace_provider
 from agents.tracing.provider import DefaultTraceProvider
-from openai.types.responses import ResponsePromptParam
 
-import temporalio.client
-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 (
@@ -47,15 +43,8 @@
     DataConverter,
     DefaultPayloadConverter,
 )
-from temporalio.plugin import Plugin, create_plugin
-from temporalio.worker import (
-    Replayer,
-    ReplayerConfig,
-    Worker,
-    WorkerConfig,
-    WorkflowReplayResult,
-    WorkflowRunner,
-)
+from temporalio.plugin import SimplePlugin
+from temporalio.worker import WorkflowRunner
 from temporalio.worker.workflow_sandbox import SandboxedWorkflowRunner
 
 # Unsupported on python 3.9
@@ -188,75 +177,142 @@ def _data_converter(converter: Optional[DataConverter]) -> DataConverter:
     return converter
 
 
-def OpenAIAgentsPlugin(
-    model_params: Optional[ModelActivityParameters] = None,
-    model_provider: Optional[ModelProvider] = None,
-    mcp_server_providers: Sequence[
-        Union["StatelessMCPServerProvider", "StatefulMCPServerProvider"]
-    ] = (),
-) -> Plugin:
-    """Create an OpenAI agents plugin.
+class OpenAIAgentsPlugin(SimplePlugin):
+    """Temporal plugin for integrating OpenAI agents with Temporal workflows.
+
+    .. warning::
+        This class is experimental and may change in future versions.
+        Use with caution in production environments.
+
+    This plugin provides seamless integration between the OpenAI Agents SDK and
+    Temporal workflows. It automatically configures the necessary interceptors,
+    activities, and data converters to enable OpenAI agents to run within
+    Temporal workflows with proper tracing and model execution.
+
+    The plugin:
+    1. Configures the Pydantic data converter for type-safe serialization
+    2. Sets up tracing interceptors for OpenAI agent interactions
+    3. Registers model execution activities
+    4. Automatically registers MCP server activities and manages their lifecycles
+    5. Manages the OpenAI agent runtime overrides during worker execution
 
     Args:
         model_params: Configuration parameters for Temporal activity execution
             of model calls. If None, default parameters will be used.
         model_provider: Optional model provider for custom model implementations.
             Useful for testing or custom model integrations.
         mcp_server_providers: Sequence of MCP servers to automatically register with the worker.
-            Each server will be wrapped in a TemporalMCPServer if not already wrapped,
-            and their activities will be automatically registered with the worker.
-            The plugin manages the connection lifecycle of these servers.
+            The plugin will wrap each server in a TemporalMCPServer if needed and
+            manage their connection lifecycles tied to the worker lifetime. This is
+            the recommended way to use MCP servers with Temporal workflows.
+
+    Example:
+        >>> from temporalio.client import Client
+        >>> from temporalio.worker import Worker
+        >>> from temporalio.contrib.openai_agents import OpenAIAgentsPlugin, ModelActivityParameters, StatelessMCPServerProvider
+        >>> from agents.mcp import MCPServerStdio
+        >>> from datetime import timedelta
+        >>>
+        >>> # Configure model parameters
+        >>> model_params = ModelActivityParameters(
+        ...     start_to_close_timeout=timedelta(seconds=30),
+        ...     retry_policy=RetryPolicy(maximum_attempts=3)
+        ... )
+        >>>
+        >>> # Create MCP servers
+        >>> filesystem_server = StatelessMCPServerProvider(MCPServerStdio(
+        ...     name="Filesystem Server",
+        ...     params={"command": "npx", "args": ["-y", "@modelcontextprotocol/server-filesystem", "."]}
+        ... ))
+        >>>
+        >>> # Create plugin with MCP servers
+        >>> plugin = OpenAIAgentsPlugin(
+        ...     model_params=model_params,
+        ...     mcp_server_providers=[filesystem_server]
+        ... )
+        >>>
+        >>> # Use with client and worker
+        >>> client = await Client.connect(
+        ...     "localhost:7233",
+        ...     plugins=[plugin]
+        ... )
+        >>> worker = Worker(
+        ...     client,
+        ...     task_queue="my-task-queue",
+        ...     workflows=[MyWorkflow],
+        ... )
     """
-    if model_params is None:
-        model_params = ModelActivityParameters()
-
-    # For the default provider, we provide a default start_to_close_timeout of 60 seconds.
-    # Other providers will need to define their own.
-    if (
-        model_params.start_to_close_timeout is None
-        and model_params.schedule_to_close_timeout is None
+
+    def __init__(
+        self,
+        model_params: Optional[ModelActivityParameters] = None,
+        model_provider: Optional[ModelProvider] = None,
+        mcp_server_providers: Sequence[
+            Union["StatelessMCPServerProvider", "StatefulMCPServerProvider"]
+        ] = (),
     ):
-        if model_provider is None:
-            model_params.start_to_close_timeout = timedelta(seconds=60)
-        else:
+        """Create an OpenAI agents plugin.
+
+        Args:
+            model_params: Configuration parameters for Temporal activity execution
+                of model calls. If None, default parameters will be used.
+            model_provider: Optional model provider for custom model implementations.
+                Useful for testing or custom model integrations.
+            mcp_server_providers: Sequence of MCP servers to automatically register with the worker.
+                Each server will be wrapped in a TemporalMCPServer if not already wrapped,
+                and their activities will be automatically registered with the worker.
+                The plugin manages the connection lifecycle of these servers.
+        """
+        if model_params is None:
+            model_params = ModelActivityParameters()
+
+        # For the default provider, we provide a default start_to_close_timeout of 60 seconds.
+        # Other providers will need to define their own.
+        if (
+            model_params.start_to_close_timeout is None
+            and model_params.schedule_to_close_timeout is None
+        ):
+            if model_provider is None:
+                model_params.start_to_close_timeout = timedelta(seconds=60)
+            else:
+                raise ValueError(
+                    "When configuring a custom provider, the model activity must have start_to_close_timeout or schedule_to_close_timeout"
+                )
+
+        new_activities = [ModelActivity(model_provider).invoke_model_activity]
+
+        server_names = [server.name for server in mcp_server_providers]
+        if len(server_names) != len(set(server_names)):
             raise ValueError(
-                "When configuring a custom provider, the model activity must have start_to_close_timeout or schedule_to_close_timeout"
+                f"More than one mcp server registered with the same name. Please provide unique names."
             )
 
-    new_activities = [ModelActivity(model_provider).invoke_model_activity]
-
-    server_names = [server.name for server in mcp_server_providers]
-    if len(server_names) != len(set(server_names)):
-        raise ValueError(
-            f"More than one mcp server registered with the same name. Please provide unique names."
+        for mcp_server in mcp_server_providers:
+            new_activities.extend(mcp_server._get_activities())
+
+        def workflow_runner(runner: Optional[WorkflowRunner]) -> WorkflowRunner:
+            if not runner:
+                raise ValueError("No WorkflowRunner provided to the OpenAI plugin.")
+
+            # If in sandbox, add additional passthrough
+            if isinstance(runner, SandboxedWorkflowRunner):
+                return dataclasses.replace(
+                    runner,
+                    restrictions=runner.restrictions.with_passthrough_modules("mcp"),
+                )
+            return runner
+
+        @asynccontextmanager
+        async def run_context() -> AsyncIterator[None]:
+            with set_open_ai_agent_temporal_overrides(model_params):
+                yield
+
+        super().__init__(
+            name="OpenAIAgentsPlugin",
+            data_converter=_data_converter,
+            worker_interceptors=[OpenAIAgentsTracingInterceptor()],
+            activities=new_activities,
+            workflow_runner=workflow_runner,
+            workflow_failure_exception_types=[AgentsWorkflowError],
+            run_context=lambda: run_context(),
         )
-
-    for mcp_server in mcp_server_providers:
-        new_activities.extend(mcp_server._get_activities())
-
-    def workflow_runner(runner: Optional[WorkflowRunner]) -> WorkflowRunner:
-        if not runner:
-            raise ValueError("No WorkflowRunner provided to the OpenAI plugin.")
-
-        # If in sandbox, add additional passthrough
-        if isinstance(runner, SandboxedWorkflowRunner):
-            return dataclasses.replace(
-                runner,
-                restrictions=runner.restrictions.with_passthrough_modules("mcp"),
-            )
-        return runner
-
-    @asynccontextmanager
-    async def run_context() -> AsyncIterator[None]:
-        with set_open_ai_agent_temporal_overrides(model_params):
-            yield
-
-    return create_plugin(
-        name="OpenAIAgentsPlugin",
-        data_converter=_data_converter,
-        worker_interceptors=[OpenAIAgentsTracingInterceptor()],
-        activities=new_activities,
-        workflow_runner=workflow_runner,
-        workflow_failure_exception_types=[AgentsWorkflowError],
-        run_context=lambda: run_context(),
-    )