💥 Replayer configuration from plugins (#1011)

* POC for replayer configuration from existing plugins

* Handle non-combined cases

* Fixing type checking

* Move shared configuration into plugin definition

* Moving replay configuration to worker plugin, add replay execution hook

* Readme and some cleanup

* Remove debug logging

* Make plugin interface experimental and abstract

* Update readme

* Update readme

* Readme fixes, move plugin to its own file

* Fix readme

* Update temporalio/worker/_plugin.py

Co-authored-by: Spencer Judge <[email protected]>

* Rename run_replayer, remove need to return self from init

* Docstring change

---------

Co-authored-by: Spencer Judge <[email protected]>

Author: tconley1428

README.md

@@ -1515,18 +1515,21 @@ import temporalio.service
 class AuthenticationPlugin(Plugin):
     def __init__(self, api_key: str):
         self.api_key = api_key
+        
+    def init_client_plugin(self, next: Plugin) -> None:
+      self.next_client_plugin = next
 
     def configure_client(self, config: ClientConfig) -> ClientConfig:
         # Modify client configuration
         config["namespace"] = "my-secure-namespace"
-        return super().configure_client(config)
+        return self.next_client_plugin.configure_client(config)
 
     async def connect_service_client(
         self, config: temporalio.service.ConnectConfig
     ) -> temporalio.service.ServiceClient:
         # Add authentication to the connection
         config.api_key = self.api_key
-        return await super().connect_service_client(config)
+        return await self.next_client_plugin.connect_service_client(config)
 
 # Use the plugin when connecting
 client = await Client.connect(
@@ -1538,31 +1541,55 @@ client = await Client.connect(
 #### Worker Plugins
 
 Worker plugins can modify worker configuration and intercept worker execution. They are useful for adding monitoring, 
-custom lifecycle management, or modifying worker settings.
+custom lifecycle management, or modifying worker settings. Worker plugins can also configure replay. 
+They should do this in the case that they modified the worker in a way which would also need to be present 
+for replay to function. For instance, changing the data converter or adding workflows. 
 
 Here's an example of a worker plugin that adds custom monitoring:
 
 ```python
-from temporalio.worker import Plugin, WorkerConfig, Worker
+import temporalio
+from contextlib import asynccontextmanager
+from typing import AsyncIterator
+from temporalio.worker import Plugin, WorkerConfig, ReplayerConfig, Worker, Replayer, WorkflowReplayResult
 import logging
 
 class MonitoringPlugin(Plugin):
     def __init__(self):
         self.logger = logging.getLogger(__name__)
 
+    def init_worker_plugin(self, next: Plugin) -> None:
+        self.next_worker_plugin = next
+        
     def configure_worker(self, config: WorkerConfig) -> WorkerConfig:
         # Modify worker configuration
         original_task_queue = config["task_queue"]
         config["task_queue"] = f"monitored-{original_task_queue}"
         self.logger.info(f"Worker created for task queue: {config['task_queue']}")
-        return super().configure_worker(config)
+        return self.next_worker_plugin.configure_worker(config)
 
     async def run_worker(self, worker: Worker) -> None:
         self.logger.info("Starting worker execution")
         try:
-            await super().run_worker(worker)
+            await self.next_worker_plugin.run_worker(worker)
         finally:
             self.logger.info("Worker execution completed")
+      
+  def configure_replayer(self, config: ReplayerConfig) -> ReplayerConfig:
+      return self.next_worker_plugin.configure_replayer(config)
+  
+  @asynccontextmanager
+  async def run_replayer(
+      self,
+      replayer: Replayer,
+      histories: AsyncIterator[temporalio.client.WorkflowHistory],
+  ) -> AsyncIterator[AsyncIterator[WorkflowReplayResult]]:
+        self.logger.info("Starting replay execution")
+        try:
+          async with self.next_worker_plugin.run_replayer(replayer, histories) as results:
+              yield results
+        finally:
+            self.logger.info("Replay execution completed")
 
 # Use the plugin when creating a worker
 worker = Worker(
@@ -1577,38 +1604,63 @@ worker = Worker(
 For plugins that need to work with both clients and workers, you can implement both interfaces in a single class:
 
 ```python
+import temporalio
+from contextlib import AbstractAsyncContextManager
+from typing import AsyncIterator
 from temporalio.client import Plugin as ClientPlugin, ClientConfig
-from temporalio.worker import Plugin as WorkerPlugin, WorkerConfig
+from temporalio.worker import Plugin as WorkerPlugin, WorkerConfig, ReplayerConfig, Worker, Replayer, WorkflowReplayResult
 
 
 class UnifiedPlugin(ClientPlugin, WorkerPlugin):
-  def configure_client(self, config: ClientConfig) -> ClientConfig:
-    # Client-side customization
-    config["namespace"] = "unified-namespace"
-    return super().configure_client(config)
-
-  def configure_worker(self, config: WorkerConfig) -> WorkerConfig:
-    # Worker-side customization
-    config["max_cached_workflows"] = 500
-    return super().configure_worker(config)
-
-  async def run_worker(self, worker: Worker) -> None:
-    print("Starting unified worker")
-    await super().run_worker(worker)
-
+    def init_client_plugin(self, next: ClientPlugin) -> None:
+        self.next_client_plugin = next
 
+    def init_worker_plugin(self, next: WorkerPlugin) -> None:
+        self.next_worker_plugin = next
+  
+    def configure_client(self, config: ClientConfig) -> ClientConfig:
+        # Client-side customization
+        config["data_converter"] = pydantic_data_converter
+        return self.next_client_plugin.configure_client(config)
+  
+    async def connect_service_client(
+        self, config: temporalio.service.ConnectConfig
+    ) -> temporalio.service.ServiceClient:
+        # Add authentication to the connection
+        config.api_key = self.api_key
+        return await self.next_client_plugin.connect_service_client(config)
+        
+    def configure_worker(self, config: WorkerConfig) -> WorkerConfig:
+        # Worker-side customization
+        return self.next_worker_plugin.configure_worker(config)
+  
+    async def run_worker(self, worker: Worker) -> None:
+        print("Starting unified worker")
+        await self.next_worker_plugin.run_worker(worker)
+        
+    def configure_replayer(self, config: ReplayerConfig) -> ReplayerConfig:
+        config["data_converter"] = pydantic_data_converter
+        return config
+    
+    async def run_replayer(
+        self,
+        replayer: Replayer,
+        histories: AsyncIterator[temporalio.client.WorkflowHistory],
+    ) -> AbstractAsyncContextManager[AsyncIterator[WorkflowReplayResult]]:
+        return self.next_worker_plugin.run_replayer(replayer, histories)
+              
 # Create client with the unified plugin
 client = await Client.connect(
-  "localhost:7233",
-  plugins=[UnifiedPlugin()]
+    "localhost:7233",
+    plugins=[UnifiedPlugin()]
 )
 
 # Worker will automatically inherit the plugin from the client
 worker = Worker(
-  client,
-  task_queue="my-task-queue",
-  workflows=[MyWorkflow],
-  activities=[my_activity]
+    client,
+    task_queue="my-task-queue",
+    workflows=[MyWorkflow],
+    activities=[my_activity]
 )
 ```
 
@@ -1617,7 +1669,7 @@ worker = Worker(
 - Plugins are executed in reverse order (last plugin wraps the first), forming a chain of responsibility
 - Client plugins that also implement worker plugin interfaces are automatically propagated to workers
 - Avoid providing the same plugin to both client and worker to prevent double execution
-- Plugin methods should call `super()` to maintain the plugin chain
+- Plugin methods should call the plugin provided during initialization to maintain the plugin chain
 - Each plugin's `name()` method returns a unique identifier for debugging purposes
 
 

temporalio/client.py

@@ -191,7 +191,8 @@ async def connect(
 
         root_plugin: Plugin = _RootPlugin()
         for plugin in reversed(plugins):
-            root_plugin = plugin.init_client_plugin(root_plugin)
+            plugin.init_client_plugin(root_plugin)
+            root_plugin = plugin
 
         service_client = await root_plugin.connect_service_client(connect_config)
 
@@ -235,7 +236,8 @@ def __init__(
 
         root_plugin: Plugin = _RootPlugin()
         for plugin in reversed(plugins):
-            root_plugin = plugin.init_client_plugin(root_plugin)
+            plugin.init_client_plugin(root_plugin)
+            root_plugin = plugin
 
         self._init_from_config(root_plugin.configure_client(config))
 
@@ -7398,22 +7400,21 @@ def name(self) -> str:
         """
         return type(self).__module__ + "." + type(self).__qualname__
 
-    def init_client_plugin(self, next: Plugin) -> Plugin:
+    @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 storing a reference
+        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.
-
-        Returns:
-            This plugin instance for method chaining.
         """
-        self.next_client_plugin = next
-        return self
 
+    @abstractmethod
     def configure_client(self, config: ClientConfig) -> ClientConfig:
         """Hook called when creating a client to allow modification of configuration.
 
@@ -7427,8 +7428,8 @@ def configure_client(self, config: ClientConfig) -> ClientConfig:
         Returns:
             The modified client configuration.
         """
-        return self.next_client_plugin.configure_client(config)
 
+    @abstractmethod
     async def connect_service_client(
         self, config: temporalio.service.ConnectConfig
     ) -> temporalio.service.ServiceClient:
@@ -7444,10 +7445,12 @@ async def connect_service_client(
         Returns:
             The connected service client.
         """
-        return await self.next_client_plugin.connect_service_client(config)
 
 
 class _RootPlugin(Plugin):
+    def init_client_plugin(self, next: Plugin) -> None:
+        raise NotImplementedError()
+
     def configure_client(self, config: ClientConfig) -> ClientConfig:
         return config
 

temporalio/contrib/openai_agents/_temporal_openai_agents.py

@@ -1,6 +1,6 @@
 """Initialize Temporal OpenAI Agents overrides."""
 
-from contextlib import contextmanager
+from contextlib import asynccontextmanager, contextmanager
 from datetime import timedelta
 from typing import AsyncIterator, Callable, Optional, Union
 
@@ -24,7 +24,7 @@
 
 import temporalio.client
 import temporalio.worker
-from temporalio.client import ClientConfig
+from temporalio.client import ClientConfig, Plugin
 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 TemporalOpenAIRunner
@@ -41,7 +41,13 @@
 from temporalio.converter import (
     DataConverter,
 )
-from temporalio.worker import Worker, WorkerConfig
+from temporalio.worker import (
+    Replayer,
+    ReplayerConfig,
+    Worker,
+    WorkerConfig,
+    WorkflowReplayResult,
+)
 
 
 @contextmanager
@@ -231,6 +237,20 @@ def __init__(
         self._model_params = model_params
         self._model_provider = model_provider
 
+    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
+
     def configure_client(self, config: ClientConfig) -> ClientConfig:
         """Configure the Temporal client for OpenAI agents integration.
 
@@ -246,7 +266,7 @@ def configure_client(self, config: ClientConfig) -> ClientConfig:
         config["data_converter"] = DataConverter(
             payload_converter_class=_OpenAIPayloadConverter
         )
-        return super().configure_client(config)
+        return self.next_client_plugin.configure_client(config)
 
     def configure_worker(self, config: WorkerConfig) -> WorkerConfig:
         """Configure the Temporal worker for OpenAI agents integration.
@@ -268,7 +288,7 @@ def configure_worker(self, config: WorkerConfig) -> WorkerConfig:
         config["activities"] = list(config.get("activities") or []) + [
             ModelActivity(self._model_provider).invoke_model_activity
         ]
-        return super().configure_worker(config)
+        return self.next_worker_plugin.configure_worker(config)
 
     async def run_worker(self, worker: Worker) -> None:
         """Run the worker with OpenAI agents temporal overrides.
@@ -281,4 +301,27 @@ async def run_worker(self, worker: Worker) -> None:
             worker: The worker instance to run.
         """
         with set_open_ai_agent_temporal_overrides(self._model_params):
-            await super().run_worker(worker)
+            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"] = DataConverter(
+            payload_converter_class=_OpenAIPayloadConverter
+        )
+        return 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

temporalio/worker/__init__.py

@@ -21,6 +21,7 @@
     WorkflowInterceptorClassInput,
     WorkflowOutboundInterceptor,
 )
+from ._plugin import Plugin
 from ._replayer import (
     Replayer,
     ReplayerConfig,
@@ -44,7 +45,6 @@
     WorkflowSlotInfo,
 )
 from ._worker import (
-    Plugin,
     PollerBehavior,
     PollerBehaviorAutoscaling,
     PollerBehaviorSimpleMaximum,