[moetemp]

What was changed

This PR makes Worker.__aexit__ wait for the run task after shutdown(), for up to 60 s, instead of cancelling it.

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]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}