Skip to content
← All articles

Never yield inside a TaskGroup or timeout scope

The most appealing-looking abstraction in async Python — an @asynccontextmanager wrapping a TaskGroup — delivers cancellations to unrelated tasks. PEP 789 documents it precisely and has not landed, so the mitigation is a review rule.

There is one rule in this article. It fits on a sticky note:

Never yield while a TaskGroup or asyncio.timeout scope is open in the same frame.

The rest of the article is why that rule exists, why the code that breaks it looks so good, and what to write instead.

The shape that lures everyone

You have three call sites that each need a task group. You are a tidy person. You factor it out:

from contextlib import asynccontextmanager
import asyncio


@asynccontextmanager
async def worker_pool() -> AsyncIterator[asyncio.TaskGroup]:
    async with asyncio.TaskGroup() as group:
        yield group
async with worker_pool() as pool:
    pool.create_task(fetch(a))
    pool.create_task(fetch(b))
    await do_something_else()

It reads beautifully. It is a cancellation bug.

What actually happens

@asynccontextmanager is implemented as an async generator. __aenter__ advances it to the yield; __aexit__ resumes it afterwards. So the frame containing async with asyncio.TaskGroup() is suspended at the yield, with the group still open, for the entire duration of your async with worker_pool() block.

A TaskGroup handles a child failure by cancelling the task that entered the group. Not the group. Not the generator. The task, identified when __aenter__ ran.

Now consider what the interpreter can actually do when fetch(a) raises while the generator frame is parked:

  • The group wants to cancel “its” task and deliver an ExceptionGroup at the point the async with block ends.
  • That point is inside a suspended generator frame that no task currently owns.
  • The cancellation lands on whichever task next drives that generator — which may be a different task than the one that entered, and which may be sitting in completely unrelated code.

The symptom in production is a CancelledError or an ExceptionGroup surfacing in a task that has nothing to do with the code that opened the scope. You will spend a day looking at the wrong file.

asyncio.timeout fails the same way for the same reason: it schedules a cancellation against the task that entered the scope.

async def batched(source: AsyncIterator[T], size: int) -> AsyncGenerator[list[T], None]:
    async with asyncio.timeout(30):        # scope opens
        batch: list[T] = []
        async for item in source:
            batch.append(item)
            if len(batch) == size:
                yield batch                # scope is STILL OPEN across this
                batch = []

Reading that, most people see “a 30-second budget for reading the source”. What it actually is: a 30-second budget for reading the source plus every consumer’s processing time between batches, whose expiry fires into whoever resumes the generator. If the consumer spends 40 seconds writing batch one to a database, the timeout expires there.

💡A colleague argues this is fine as long as only one task ever uses the generator. Is it? click to reveal

It is much better, and it is still not fine.

With a single consumer you have dodged the “wrong task” half of the problem. What remains is the “wrong duration” half, which is not a race and shows up every time: the scope’s clock keeps running while the generator is suspended, so a timeout meant to bound your work also bounds the consumer’s. That is a real bug, just a deterministic one.

And then there is the exit path. If the consumer breaks out of the loop, aclose() throws GeneratorExit in at the yield, and the TaskGroup‘s __aexit__ now runs during finalisation — possibly from a different task, possibly during loop shutdown, in a context where its children cannot be awaited to completion. PEP 789’s analysis is that the exit is executed in a context the scope never agreed to.

“Only one consumer” is also not a property the type system can express or the reviewer can enforce. AsyncIterator[T] says nothing about how many tasks will drive it. So even when the reasoning holds today, it holds by convention, and conventions are exactly what a refactor deletes.

PEP 789, and why the fix is a habit

PEP 789, Preventing task-cancellation bugs by limiting yield in async generators, documents these failure modes precisely and proposes a sys.prevent_yields() context manager so that libraries opening a cancel scope can make yielding across it a runtime error.

It is Draft. It has not landed. There is no interpreter check, no ruff rule that catches the general case, and no type-checker diagnostic. Until there is, this is a review rule enforced by humans, which is why it is worth being able to state it from memory. trio and anyio have documented the same hazard for years; PEP 789 is the attempt to make the interpreter enforce what those communities learned by bleeding.

What to write instead

Move the scope to the consumer. The clean version of the worker-pool example is not a context manager at all — it is the caller opening the group, because the caller is the one whose lifetime the group should follow:

async with asyncio.TaskGroup() as group:
    group.create_task(fetch(a))
    group.create_task(fetch(b))
    await do_something_else()

Three extra characters at each call site, and the scope’s lifetime is now visibly the block’s lifetime. Most “I factored out the TaskGroup” abstractions are deleting the one piece of information the reader needed.

If you need an object, make the lifecycle explicit. A class with async def start() and async def aclose() can hold a group without a generator frame in the middle:

class Pool:
    async def __aenter__(self) -> Self:
        self._group = asyncio.TaskGroup()
        await self._group.__aenter__()
        return self

    async def __aexit__(self, *exc: object) -> bool | None:
        return await self._group.__aexit__(*exc)

This is not prettier, and it is honest: __aenter__ and __aexit__ are ordinary coroutines, so there is no suspended generator frame and no yield across the scope. Whether the ergonomics are worth it is a judgement call — the point is that @asynccontextmanager is not a free abstraction here.

Keep the scope entirely inside one segment. If the generator needs a timeout on its own work, put the scope around only that work, not around the yield:

async def batched(source: AsyncGenerator[T, None], size: int) -> AsyncGenerator[list[T], None]:
    async with aclosing(source):
        batch: list[T] = []
        async for item in source:
            batch.append(item)
            if len(batch) == size:
                yield batch               # no cancel scope is open here
                batch = []
        if batch:
            yield batch

aclosing is not a cancel scope

Notice the async with aclosing(source) in that last example, wrapped around a yield. That is fine, and knowing why is the actual skill.

contextlib.aclosing is an ordinary async context manager: its __aexit__ awaits source.aclose() and that is all. It registers no cancellation, it holds no clock, and it has no opinion about which task it is running in. Suspending inside it costs nothing.

A cancel scope is the special case: asyncio.TaskGroup, asyncio.timeout, asyncio.timeout_at, anyio’s CancelScope, trio’s nursery. What they share is that they can inject a CancelledError into a specific task at a moment they choose. That injection needs a task to aim at, and a suspended generator frame does not have one.

So the rule generalises to: do not suspend a frame across anything that can cancel you. Everything else is fine.

💡How would you catch a violation of this rule in review, or in CI, given there is no linter rule for it? click to reveal

In review, the tell is short and grep-able: an async def that contains both a yield and one of TaskGroup, asyncio.timeout, asyncio.timeout_at, move_on_after, fail_after, or CancelScope. Because @asynccontextmanager compiles to an async generator, the decorator counts as a yield for this purpose — so @asynccontextmanager plus TaskGroup in the same function is the highest-value single pattern to look for.

In CI, an AST check is about twenty lines and is worth writing once: walk each AsyncFunctionDef, collect the names used in AsyncWith items, and flag the function if it contains a Yield node and any name from your cancel-scope list. It over-approximates — a yield lexically after the with block has closed is a false positive — but sequencing that precisely needs control-flow analysis, and in practice the over-approximation finds real bugs and costs one comment to dismiss.

What cannot catch it: mypy, which has no notion of cancel scopes; and tests, because the failure needs a child task to fail while the parent frame happens to be suspended. That is exactly the interleaving your test suite does not produce and production does.

The type flip, while you are here

One more thing that trips people up in this exact code. Adding a single yield changes what your function is:

async def stream(src: AsyncIterator[T]) -> AsyncIterator[T]:
    return _wrap(src)          # a coroutine returning an async iterator
async def stream(src: AsyncIterator[T]) -> AsyncIterator[T]:
    async for item in src:
        yield item             # an async GENERATOR function

Same signature, same annotation, two different calling conventions: the first must be awaited, the second must not. mypy accepts both readings without comment, so reveal_type(stream(src)) is how you find out which one you wrote.

The consequence that matters here: an async generator runs none of its body until the first __anext__, so it cannot validate its arguments at call time. If your streaming function must reject size <= 0 before the caller starts iterating, split it: a plain def that validates and returns, plus the async def ... yield worker underneath.

Laziness, Iteration and Pipelines · step 4 of 14

Next: zip(strict=True) and the silent-truncation bug → ← Back to Laziness, Iteration and Pipelines