[moetemp]

What was changed

This PR makes Worker.__aexit__ wait for the run task after shutdown() instead of cancelling it. The new Worker option plugin_shutdown_timeout (default 60 s, None for no limit) bounds that wait, and a cancel of __aexit__ cancels the run task too.

Why?

After shutdown(), the run task has only plugin cleanup left, such as the finally of Plugin.run_worker. __aexit__ cancelled that cleanup while it awaited. The run task caught the CancelledError and cancelled the caller's task too. So the cleanup didn't finish, and every later await in the caller raised CancelledError, even after the async with block.

Here's a repro. A plugin awaits a client call in its finally, and the caller awaits after the block.

class CleanupPlugin(SimplePlugin):
    def __init__(self, client: Client) -> None:
        super().__init__("cleanup")
        self.client = client

    async def run_worker(self, worker, next):
        try:
            await next(worker)
        finally:
            await self.client.count_workflows()  # cancelled before this change

async with Worker(
    client, task_queue="tq", workflows=[MyWorkflow], plugins=[CleanupPlugin(client)]
):
    pass
await client.count_workflows()  # raised CancelledError before this change

Checklist

1. How was this tested: tests/worker/test_worker_exit.py runs the repro. It fails on main and passes with this change. I also ran poe lint and the full suite.

2. Any docs updates needed? No.

[CLAassistant]

All committers have signed the CLA.