Keep the generator’s Cancelled when as_safe_channel is cancelled

newsfragments/3554.bugfix.rst

@@ -0,0 +1,4 @@
+When an ``async with`` on an `@as_safe_channel <trio.as_safe_channel>` generator is cancelled
+while the generator's own code is running, the generator's `~trio.Cancelled` now propagates, as
+it would from a plain ``async for``. Previously the choice was random, so an exception the
+generator raised before its cleanup was cancelled could be missing from the traceback.

src/trio/_channel.py

@@ -22,6 +22,7 @@
     MultipleExceptionError,
     NoPublicConstructor,
     final,
+    raise_saving_context,
     raise_single_exception_from_group,
 )
 
@@ -559,6 +560,7 @@ async def context_manager(
         *args: P.args, **kwargs: P.kwargs
     ) -> AsyncGenerator[trio._channel.RecvChanWrapper[T], None]:
         send_chan, recv_chan = trio.open_memory_channel[T](0)
+        agen_cancelled: list[trio.Cancelled] = []
         try:
             async with trio.open_nursery(strict_exception_groups=True) as nursery:
                 agen = fn(*args, **kwargs)
@@ -567,7 +569,11 @@ async def context_manager(
                 # If this errors we don't close `recv_chan`, but the caller
                 # never gets access to it, so that's not a problem.
                 await nursery.start(
-                    _move_elems_to_channel, agen, send_chan, send_semaphore
+                    _move_elems_to_channel,
+                    agen,
+                    send_chan,
+                    send_semaphore,
+                    agen_cancelled,
                 )
                 # `async with recv_chan` could eat exceptions, so use sync cm
                 with RecvChanWrapper(recv_chan, send_semaphore) as wrapped_recv_chan:
@@ -578,6 +584,13 @@ async def context_manager(
                     "exited trio.as_safe_channel context manager"
                 )
         except BaseExceptionGroup as eg:
+            # If the generator's own code was cancelled and the group holds nothing but
+            # Cancelled, raise what a plain `async for` would: the generator's Cancelled,
+            # with its context. (A group raised by the body is not a Cancelled.)
+            if agen_cancelled and all(
+                isinstance(e, trio.Cancelled) for e in eg.exceptions
+            ):
+                raise_saving_context(agen_cancelled[0])
             try:
                 raise_single_exception_from_group(eg)
             except MultipleExceptionError:
@@ -590,11 +603,16 @@ async def context_manager(
                     )
 
                 raise eg from None
+        finally:
+            # Break the reference cycle: the Cancelled's traceback holds this frame,
+            # which holds the list.
+            agen_cancelled.clear()
 
     async def _move_elems_to_channel(
         agen: AsyncGenerator[T, None],
         send_chan: trio.MemorySendChannel[T],
         send_semaphore: trio.Semaphore,
+        agen_cancelled: list[trio.Cancelled],
         task_status: trio.TaskStatus,
     ) -> None:
         # `async with send_chan` will eat exceptions,
@@ -615,6 +633,9 @@ async def _move_elems_to_channel(
                         value = await agen.__anext__()
                     except StopAsyncIteration:
                         return
+                    except trio.Cancelled as cancelled:
+                        agen_cancelled.append(cancelled)
+                        raise
                     # Send the value to the channel
                     try:
                         await send_chan.send(value)

src/trio/_tests/test_channel.py

@@ -638,6 +638,32 @@ async def agen(events: list[str]) -> AsyncGenerator[None]:
     assert events == ["body cancel", "agen cancel"]
 
 
+async def test_as_safe_channel_cancelled_in_cleanup_keeps_context(
+    autojump_clock: trio.abc.Clock,
+) -> None:
+    error = ValueError("agen")
+
+    @as_safe_channel
+    async def agen() -> AsyncGenerator[None]:
+        try:
+            yield
+            raise error
+        finally:
+            # The timeout cancels the body and this task. The shield makes this task's
+            # Cancelled, which carries `error`, reach the nursery after the body's;
+            # unshielded, the order is random.
+            with trio.CancelScope(shield=True):
+                await trio.sleep(2)
+            await trio.lowlevel.checkpoint()
+
+    with trio.move_on_after(1):
+        with pytest.raises(trio.Cancelled) as excinfo:
+            async with agen() as recv_chan:
+                async for _ in recv_chan:  # pragma: no branch
+                    pass
+    assert excinfo.value.__context__ is error
+
+
 async def test_as_safe_channel_genexit_exception_group() -> None:
     @as_safe_channel
     async def agen() -> AsyncGenerator[None]:

src/trio/_util.py

@@ -358,7 +358,7 @@ def raise_single_exception_from_group(
     the cause/context of the group(s).
 
     If multiple non-cancelled exceptions are encountered, it raises
-    :exc:`AssertionError`.
+    :exc:`MultipleExceptionError`.
     """
     # immediately bail out if there's any KI or SystemExit
     for e in eg.exceptions: