POC for moving plugin run_worker to a context

Author: tconley1428

temporalio/contrib/openai_agents/_temporal_openai_agents.py

@@ -276,7 +276,8 @@ def configure_worker(self, config: WorkerConfig) -> WorkerConfig:
         ]
         return super().configure_worker(config)
 
-    async def run_worker(self, worker: Worker) -> None:
+    @asynccontextmanager
+    async def run_worker(self) -> AsyncIterator[None]:
         """Run the worker with OpenAI agents temporal overrides.
 
         This method sets up the necessary runtime overrides for OpenAI agents
@@ -287,7 +288,8 @@ 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)
+            async with super().run_worker():
+                yield
 
     def configure_replayer(self, config: ReplayerConfig) -> ReplayerConfig:
         """Configure the replayer for OpenAI Agents."""

temporalio/worker/_worker.py

@@ -140,17 +140,14 @@ def configure_worker(self, config: WorkerConfig) -> WorkerConfig:
         """
         return self.next_worker_plugin.configure_worker(config)
 
-    async def run_worker(self, worker: Worker) -> None:
+    def run_worker(self) -> AbstractAsyncContextManager[None]:
         """Hook called when running a worker to allow interception of execution.
 
         This method is called when the worker is started and allows plugins to
         intercept or wrap the worker execution. Plugins can add monitoring,
         custom lifecycle management, or other execution-time behavior.
-
-        Args:
-            worker: The worker instance to run.
         """
-        await self.next_worker_plugin.run_worker(worker)
+        return self.next_worker_plugin.run_worker()
 
     def configure_replayer(self, config: ReplayerConfig) -> ReplayerConfig:
         """Hook called when creating a replayer to allow modification of configuration.
@@ -176,8 +173,9 @@ class _RootPlugin(Plugin):
     def configure_worker(self, config: WorkerConfig) -> WorkerConfig:
         return config
 
-    async def run_worker(self, worker: Worker) -> None:
-        await worker._run()
+    @asynccontextmanager
+    async def run_worker(self) -> AsyncIterator[None]:
+        yield
 
     def workflow_replay(
         self,
@@ -794,7 +792,8 @@ async def run(self) -> None:
         also cancel the shutdown process. Therefore users are encouraged to use
         explicit shutdown instead.
         """
-        await self._plugin.run_worker(self)
+        async with self._plugin.run_worker():
+            await self._run()
 
     async def _run(self):
         # Eagerly validate which will do a namespace check in Core

temporalio/worker/workflow_sandbox/_importer.py

@@ -285,12 +285,6 @@ def module_configured_passthrough(self, name: str) -> bool:
     def _maybe_passthrough_module(self, name: str) -> Optional[types.ModuleType]:
         # If imports not passed through and all modules are not passed through
         # and name not in passthrough modules, check parents
-        logger.debug(
-            "Check passthrough module: %s - %s",
-            name,
-            temporalio.workflow.unsafe.is_imports_passed_through()
-            or self.module_configured_passthrough(name),
-        )
         if (
             not temporalio.workflow.unsafe.is_imports_passed_through()
             and not self.module_configured_passthrough(name)

tests/test_plugins.py

@@ -1,7 +1,9 @@
 import dataclasses
 import uuid
 import warnings
-from typing import cast
+from contextlib import asynccontextmanager
+from datetime import timedelta
+from typing import AsyncIterator, cast
 
 import pytest
 
@@ -68,6 +70,9 @@ def configure_worker(self, config: WorkerConfig) -> WorkerConfig:
         return super().configure_worker(config)
 
 
+IN_CONTEXT: bool = False
+
+
 class MyWorkerPlugin(temporalio.worker.Plugin):
     def configure_worker(self, config: WorkerConfig) -> WorkerConfig:
         config["task_queue"] = "replaced_queue"
@@ -79,8 +84,15 @@ def configure_worker(self, config: WorkerConfig) -> WorkerConfig:
             )
         return super().configure_worker(config)
 
-    async def run_worker(self, worker: Worker) -> None:
-        await super().run_worker(worker)
+    @asynccontextmanager
+    async def run_worker(self) -> AsyncIterator[None]:
+        global IN_CONTEXT
+        try:
+            IN_CONTEXT = True
+            async with super().run_worker():
+                yield
+        finally:
+            IN_CONTEXT = False
 
 
 async def test_worker_plugin_basic_config(client: Client) -> None:
@@ -109,6 +121,30 @@ async def test_worker_plugin_basic_config(client: Client) -> None:
     assert worker.config().get("task_queue") == "replaced_queue"
 
 
+@workflow.defn(sandboxed=False)
+class CheckContextWorkflow:
+    @workflow.run
+    async def run(self) -> bool:
+        return IN_CONTEXT
+
+
+async def test_worker_plugin_run_context(client: Client) -> None:
+    async with Worker(
+        client,
+        task_queue=str(uuid.uuid4()),
+        workflows=[CheckContextWorkflow],
+        activities=[never_run_activity],
+        plugins=[MyWorkerPlugin()],
+    ) as worker:
+        result = await client.execute_workflow(
+            CheckContextWorkflow.run,
+            task_queue=worker.task_queue,
+            id=f"workflow-{uuid.uuid4()}",
+            execution_timeout=timedelta(seconds=1),
+        )
+        assert result
+
+
 async def test_worker_duplicated_plugin(client: Client) -> None:
     new_config = client.config()
     new_config["plugins"] = [MyCombinedPlugin()]