fix: await the Worker run task on exit instead of cancelling it

changelog/fixed/patient-porcupine.md

@@ -0,0 +1 @@
+Exiting `async with Worker(...)` now waits for the Worker to finish its shutdown, including the `run_worker` cleanup of plugins, instead of cancelling the task that entered the block.

temporalio/worker/_worker.py

@@ -488,6 +488,7 @@ def _init_from_config(self, client: temporalio.client.Client, config: WorkerConf
         self._async_context_inner_task: asyncio.Task | None = None
         self._async_context_run_task: asyncio.Task | None = None
         self._async_context_run_exception: BaseException | None = None
+        self._async_context_exiting = False
 
         self._activity_worker: _ActivityWorker | None = None
         self._runtime = (
@@ -960,7 +961,9 @@ async def run():
                 await self.run()
             except BaseException as err:
                 self._async_context_run_exception = err
-                self._async_context_inner_task.cancel()  # type: ignore[union-attr]
+                # After __aexit__ starts, a cancel would stop the caller's code after the block.
+                if not self._async_context_exiting:
+                    self._async_context_inner_task.cancel()  # type: ignore[union-attr]
 
         self._async_context_run_task = asyncio.create_task(run())
         return self
@@ -975,13 +978,21 @@ async def __aexit__(self, exc_type: type[BaseException] | None, *args: Any) -> N
         if not self._async_context_run_task:
             raise RuntimeError("Never started")
         await self.shutdown()
-        # Cancel our run task
-        self._async_context_run_task.cancel()
+        # Plugins can run cleanup code after shutdown. Wait for that code instead of cancelling it.
+        self._async_context_exiting = True
+        run_task = self._async_context_run_task
+        await asyncio.wait([run_task], timeout=_RUN_TASK_EXIT_TIMEOUT.total_seconds())
+        if not run_task.done():
+            run_task.cancel()
         # Only re-raise our exception if present and exc_type is cancel
         if exc_type is asyncio.CancelledError and self._async_context_run_exception:
             raise self._async_context_run_exception
 
 
+# How long __aexit__ waits for plugin cleanup code after shutdown.
+_RUN_TASK_EXIT_TIMEOUT = timedelta(seconds=60)
+
+
 class WorkerConfig(TypedDict, total=False):
     """TypedDict of config originally passed to :py:class:`Worker`."""
 

tests/worker/test_worker_exit.py

@@ -0,0 +1,44 @@
+import uuid
+from collections.abc import Awaitable, Callable
+
+from temporalio import workflow
+from temporalio.client import Client
+from temporalio.plugin import SimplePlugin
+from temporalio.worker import Worker
+
+
+@workflow.defn
+class NothingWorkflow:
+    @workflow.run
+    async def run(self) -> None:
+        pass
+
+
+class CleanupPlugin(SimplePlugin):
+    def __init__(self, client: Client) -> None:
+        super().__init__("test-cleanup")
+        self.client = client
+        self.cleanup_result: int | None = None
+
+    async def run_worker(
+        self, worker: Worker, next: Callable[[Worker], Awaitable[None]]
+    ) -> None:
+        try:
+            await super().run_worker(worker, next)
+        finally:
+            # Plugin cleanup code calls the server after the Worker stops.
+            self.cleanup_result = (await self.client.count_workflows()).count
+
+
+async def test_plugin_cleanup_awaits_after_worker_exit(client: Client):
+    plugin = CleanupPlugin(client)
+    async with Worker(
+        client,
+        task_queue=f"tq-{uuid.uuid4()}",
+        workflows=[NothingWorkflow],
+        plugins=[plugin],
+    ):
+        pass
+    assert plugin.cleanup_result is not None
+    # The caller can still await after the block.
+    await client.count_workflows()