💥 Plugin Overhaul (#1139)

* Add static plugin constructor

* Allow callables as parameters

* Convert OpenAI plugin as example

* Delay openai client creation

* Change plugin structure to remove initializers

* PR feedback - exposing SimplePlugin type

* Early check OpenAI connectivity if outside a workflow

* PR Feedback

* Don't register activities in replayer

* Delay activity construction until needed

* Linting

* Linting

* Simplify statement

Author: tconley1428

temporalio/client.py

@@ -70,11 +70,13 @@
     WorkflowSerializationContext,
 )
 from temporalio.service import (
+    ConnectConfig,
     HttpConnectProxyConfig,
     KeepAliveConfig,
     RetryConfig,
     RPCError,
     RPCStatusCode,
+    ServiceClient,
     TLSConfig,
 )
 
@@ -198,12 +200,14 @@ async def connect(
             http_connect_proxy_config=http_connect_proxy_config,
         )
 
-        root_plugin: Plugin = _RootPlugin()
+        def make_lambda(plugin, next):
+            return lambda config: plugin.connect_service_client(config, next)
+
+        next_function = ServiceClient.connect
         for plugin in reversed(plugins):
-            plugin.init_client_plugin(root_plugin)
-            root_plugin = plugin
+            next_function = make_lambda(plugin, next_function)
 
-        service_client = await root_plugin.connect_service_client(connect_config)
+        service_client = await next_function(connect_config)
 
         return Client(
             service_client,
@@ -243,12 +247,10 @@ def __init__(
             plugins=plugins,
         )
 
-        root_plugin: Plugin = _RootPlugin()
-        for plugin in reversed(plugins):
-            plugin.init_client_plugin(root_plugin)
-            root_plugin = plugin
+        for plugin in plugins:
+            config = plugin.configure_client(config)
 
-        self._init_from_config(root_plugin.configure_client(config))
+        self._init_from_config(config)
 
     def _init_from_config(self, config: ClientConfig):
         self._config = config
@@ -7541,20 +7543,6 @@ def name(self) -> str:
         """
         return type(self).__module__ + "." + type(self).__qualname__
 
-    @abstractmethod
-    def init_client_plugin(self, next: Plugin) -> None:
-        """Initialize this plugin in the plugin chain.
-
-        This method sets up the chain of responsibility pattern by providing a reference
-        to the next plugin in the chain. It is called during client creation to build
-        the plugin chain. Note, this may be called twice in the case of :py:meth:`connect`.
-        Implementations should store this reference and call the corresponding method
-        of the next plugin on method calls.
-
-        Args:
-            next: The next plugin in the chain to delegate to.
-        """
-
     @abstractmethod
     def configure_client(self, config: ClientConfig) -> ClientConfig:
         """Hook called when creating a client to allow modification of configuration.
@@ -7572,8 +7560,10 @@ def configure_client(self, config: ClientConfig) -> ClientConfig:
 
     @abstractmethod
     async def connect_service_client(
-        self, config: temporalio.service.ConnectConfig
-    ) -> temporalio.service.ServiceClient:
+        self,
+        config: ConnectConfig,
+        next: Callable[[ConnectConfig], Awaitable[ServiceClient]],
+    ) -> ServiceClient:
         """Hook called when connecting to the Temporal service.
 
         This method is called during service client connection and allows plugins
@@ -7586,16 +7576,3 @@ async def connect_service_client(
         Returns:
             The connected service client.
         """
-
-
-class _RootPlugin(Plugin):
-    def init_client_plugin(self, next: Plugin) -> None:
-        raise NotImplementedError()
-
-    def configure_client(self, config: ClientConfig) -> ClientConfig:
-        return config
-
-    async def connect_service_client(
-        self, config: temporalio.service.ConnectConfig
-    ) -> temporalio.service.ServiceClient:
-        return await temporalio.service.ServiceClient.connect(config)

temporalio/contrib/openai_agents/_invoke_model_activity.py

@@ -36,7 +36,7 @@
 from pydantic_core import to_json
 from typing_extensions import Required, TypedDict
 
-from temporalio import activity
+from temporalio import activity, workflow
 from temporalio.contrib.openai_agents._heartbeat_decorator import _auto_heartbeater
 from temporalio.exceptions import ApplicationError
 

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,13 +43,8 @@
     DataConverter,
     DefaultPayloadConverter,
 )
-from temporalio.worker import (
-    Replayer,
-    ReplayerConfig,
-    Worker,
-    WorkerConfig,
-    WorkflowReplayResult,
-)
+from temporalio.plugin import SimplePlugin
+from temporalio.worker import WorkflowRunner
 from temporalio.worker.workflow_sandbox import SandboxedWorkflowRunner
 
 # Unsupported on python 3.9
@@ -172,7 +163,21 @@ def __init__(self) -> None:
         super().__init__(ToJsonOptions(exclude_unset=True))
 
 
-class OpenAIAgentsPlugin(temporalio.client.Plugin, temporalio.worker.Plugin):
+def _data_converter(converter: Optional[DataConverter]) -> DataConverter:
+    if converter is None:
+        return DataConverter(payload_converter_class=OpenAIPayloadConverter)
+    elif converter.payload_converter_class is DefaultPayloadConverter:
+        return dataclasses.replace(
+            converter, payload_converter_class=OpenAIPayloadConverter
+        )
+    elif not isinstance(converter.payload_converter, OpenAIPayloadConverter):
+        raise ValueError(
+            "The payload converter must be of type OpenAIPayloadConverter."
+        )
+    return converter
+
+
+class OpenAIAgentsPlugin(SimplePlugin):
     """Temporal plugin for integrating OpenAI agents with Temporal workflows.
 
     .. warning::
@@ -278,127 +283,48 @@ def __init__(
                     "When configuring a custom provider, the model activity must have start_to_close_timeout or schedule_to_close_timeout"
                 )
 
-        self._model_params = model_params
-        self._model_provider = model_provider
-        self._mcp_server_providers = mcp_server_providers
-        self._register_activities = register_activities
-
-    def init_client_plugin(self, next: temporalio.client.Plugin) -> None:
-        """Set the next client plugin"""
-        self.next_client_plugin = next
-
-    async def connect_service_client(
-        self, config: temporalio.service.ConnectConfig
-    ) -> temporalio.service.ServiceClient:
-        """No modifications to service client"""
-        return await self.next_client_plugin.connect_service_client(config)
-
-    def init_worker_plugin(self, next: temporalio.worker.Plugin) -> None:
-        """Set the next worker plugin"""
-        self.next_worker_plugin = next
-
-    @staticmethod
-    def _data_converter(converter: Optional[DataConverter]) -> DataConverter:
-        if converter is None:
-            return DataConverter(payload_converter_class=OpenAIPayloadConverter)
-        elif converter.payload_converter_class is DefaultPayloadConverter:
-            return dataclasses.replace(
-                converter, payload_converter_class=OpenAIPayloadConverter
-            )
-        elif not isinstance(converter.payload_converter, OpenAIPayloadConverter):
-            raise ValueError(
-                "The payload converter must be of type OpenAIPayloadConverter."
-            )
-        return converter
-
-    def configure_client(self, config: ClientConfig) -> ClientConfig:
-        """Configure the Temporal client for OpenAI agents integration.
-
-        This method sets up the Pydantic data converter to enable proper
-        serialization of OpenAI agent objects and responses.
-
-        Args:
-            config: The client configuration to modify.
-
-        Returns:
-            The modified client configuration.
-        """
-        config["data_converter"] = self._data_converter(config["data_converter"])
-        return self.next_client_plugin.configure_client(config)
-
-    def configure_worker(self, config: WorkerConfig) -> WorkerConfig:
-        """Configure the Temporal worker for OpenAI agents integration.
-
-        This method adds the necessary interceptors and activities for OpenAI
-        agent execution:
-        - Adds tracing interceptors for OpenAI agent interactions
-        - Registers model execution activities
-
-        Args:
-            config: The worker configuration to modify.
-
-        Returns:
-            The modified worker configuration.
-        """
-        config["interceptors"] = list(config.get("interceptors") or []) + [
-            OpenAIAgentsTracingInterceptor()
-        ]
+        # Delay activity construction until they are actually needed
+        def add_activities(
+            activities: Optional[Sequence[Callable]],
+        ) -> Sequence[Callable]:
+            if not register_activities:
+                return activities or []
 
-        if self._register_activities:
-            new_activities = [ModelActivity(self._model_provider).invoke_model_activity]
+            new_activities = [ModelActivity(model_provider).invoke_model_activity]
 
-            server_names = [server.name for server in self._mcp_server_providers]
+            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 self._mcp_server_providers:
+            for mcp_server in mcp_server_providers:
                 new_activities.extend(mcp_server._get_activities())
-            config["activities"] = list(config.get("activities") or []) + new_activities
-
-        runner = config.get("workflow_runner")
-        if isinstance(runner, SandboxedWorkflowRunner):
-            config["workflow_runner"] = dataclasses.replace(
-                runner,
-                restrictions=runner.restrictions.with_passthrough_modules("mcp"),
-            )
-
-        config["workflow_failure_exception_types"] = list(
-            config.get("workflow_failure_exception_types") or []
-        ) + [AgentsWorkflowError]
-        return self.next_worker_plugin.configure_worker(config)
+            return list(activities or []) + new_activities
 
-    async def run_worker(self, worker: Worker) -> None:
-        """Run the worker with OpenAI agents temporal overrides.
+        def workflow_runner(runner: Optional[WorkflowRunner]) -> WorkflowRunner:
+            if not runner:
+                raise ValueError("No WorkflowRunner provided to the OpenAI plugin.")
 
-        This method sets up the necessary runtime overrides for OpenAI agents
-        to work within the Temporal worker context, including custom runners
-        and trace providers.
-
-        Args:
-            worker: The worker instance to run.
-        """
-        with set_open_ai_agent_temporal_overrides(self._model_params):
-            await self.next_worker_plugin.run_worker(worker)
-
-    def configure_replayer(self, config: ReplayerConfig) -> ReplayerConfig:
-        """Configure the replayer for OpenAI Agents."""
-        config["interceptors"] = list(config.get("interceptors") or []) + [
-            OpenAIAgentsTracingInterceptor()
-        ]
-        config["data_converter"] = self._data_converter(config.get("data_converter"))
-        return self.next_worker_plugin.configure_replayer(config)
-
-    @asynccontextmanager
-    async def run_replayer(
-        self,
-        replayer: Replayer,
-        histories: AsyncIterator[temporalio.client.WorkflowHistory],
-    ) -> AsyncIterator[AsyncIterator[WorkflowReplayResult]]:
-        """Set the OpenAI Overrides during replay"""
-        with set_open_ai_agent_temporal_overrides(self._model_params):
-            async with self.next_worker_plugin.run_replayer(
-                replayer, histories
-            ) as results:
-                yield results
+            # 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=add_activities,
+            workflow_runner=workflow_runner,
+            workflow_failure_exception_types=[AgentsWorkflowError],
+            run_context=lambda: run_context(),
+        )