closing an async_generator_asend does not throw GeneratorExit into the `await` statement

agen_close.py

import types


def underline(text):
    last_line = text.splitlines()[-1]
    return text + "\n" + "^" * len(last_line)

@types.coroutine
def _async_yield(v):
    return (yield v)

import itertools
import logging

logger = logging.getLogger(__name__)


do_close = object()


class MyCancelledError(Exception):
    def __init__(self, do_raise, /, *args):
        super().__init__(do_raise, *args)
        self.do_raise = do_raise


async def agenfn(side_channel):
    for i in itertools.count():
        for j in itertools.count():
            try:
                if (await _async_yield((i, j)) is do_close):
                    break
            except MyCancelledError as e:
                logger.exception("async yield cancelled!")
                if e.do_raise:
                    raise
            except BaseException:
                logger.exception("async yield broke!")
                side_channel.append((i, j))
                raise
        try:
            if ((yield i) is do_close):
                break
        except BaseException:
            logger.exception("yield broke!")
            side_channel.append(i)
            raise



agen = agenfn(side_channel1 := [])
agen_asend1 = agen.asend(None)
print(f"{agen_asend1.send(None)=}")
print(f"{agen_asend1.send(1)=}")
print(f"{agen_asend1.send(2)=}")
try:
    print(f"{agen_asend1.throw(MyCancelledError(True))=}")
except MyCancelledError:
    print("gen raised cancelled error correctly!")
agen_asend2 = agen.asend(1)
try:
    print(f"{agen_asend2.send(None)=}")
except StopAsyncIteration:
    print("gen raised StopAsyncIteration correctly!")
print(f"{agen_asend2.close()=}")


agen2 = agenfn(side_channel2 := [])
agen2_asend1 = agen2.asend(None)
print(f"{agen2_asend1.send(None)=}")
print(f"{agen2_asend1.send(1)=}")
print(f"{agen2_asend1.send(2)=}")
print(f"{agen2_asend1.close()}")
if not side_channel2:
    print(underline(f"{side_channel2=}\ndid not raise GeneratorExit into the await!"))
agen2_asend2 = agen2.asend(1)
try:
    print(f"{agen2_asend2.send(None)=}")
except StopAsyncIteration:
    print("gen raised StopAsyncIteration correctly!")
except RuntimeError:
    logger.exception(underline("gen raised RuntimeError incorrectly!"))
print(f"{agen2_asend2.close()=}")
print(underline("atexit runs now, but asyncgen cleanup happens too late:"))

agen_asend1.send(None)=(0, 0)
agen_asend1.send(1)=(0, 1)
agen_asend1.send(2)=(0, 2)
async yield cancelled!
Traceback (most recent call last):
  File "/home/graingert/projects/cpython/demo.py", line 31, in agenfn
    if (await _async_yield((i, j)) is do_close):
  File "/home/graingert/projects/cpython/demo.py", line 10, in _async_yield
    return (yield v)
MyCancelledError: True
gen raised cancelled error correctly!
gen raised StopAsyncIteration correctly!
agen_asend2.close()=None
agen2_asend1.send(None)=(0, 0)
agen2_asend1.send(1)=(0, 1)
agen2_asend1.send(2)=(0, 2)
None
side_channel2=[]
did not raise GeneratorExit into the await!
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
gen raised RuntimeError incorrectly!
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
Traceback (most recent call last):
  File "/home/graingert/projects/cpython/demo.py", line 78, in <module>
    print(f"{agen2_asend2.send(None)=}")
RuntimeError: anext(): asynchronous generator is already running
agen2_asend2.close()=None
atexit runs now, but asyncgen cleanup happens too late:
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
async yield broke!
Traceback (most recent call last):
  File "/home/graingert/projects/cpython/demo.py", line 31, in agenfn
    if (await _async_yield((i, j)) is do_close):
GeneratorExit