gh-159620: Fix remote asyncio stacks stopping at async generators (GH-159623)

Follow asynchronous-generator send/throw wrappers when unwinding suspended
and running asyncio tasks.

Co-authored-by: Pablo Galindo Salgado <pablogsal@gmail.com>

Lib/test/test_external_inspection.py

@@ -475,6 +475,17 @@ sock.connect(('localhost', {port}))
 
 @requires_remote_subprocess_debugging()
 class TestSelfStackTrace(RemoteInspectionTestBase):
+    def _extract_coroutine_funcnames(self, trace):
+        return {
+            task.task_name: [
+                frame.funcname.rpartition(".")[2]
+                for coro in task.coroutine_stack
+                for frame in coro.call_stack
+            ]
+            for info in trace
+            for task in info.awaited_by
+        }
+
     @skip_if_not_supported
     def test_long_task_name_is_truncated(self):
         # gh-157788
@@ -508,6 +519,101 @@ class TestSelfStackTrace(RemoteInspectionTestBase):
 
         self.assertEqual(asyncio.run(rec(3)), ["rec"] * 4)
 
+    @skip_if_not_supported
+    def test_async_generator_stack_is_not_truncated(self):
+        # gh-159620
+        async def agen():
+            await asyncio.sleep(10_000)
+            yield
+
+        async def agen_with_cleanup():
+            try:
+                yield
+            finally:
+                await asyncio.sleep(10_000)
+
+        async def outer_agen():
+            async for _ in agen():
+                yield
+
+        async def consumer(gen):
+            async for _ in gen:
+                pass
+
+        async def main():
+            g = agen_with_cleanup()
+            await anext(g)
+            async_for_task = asyncio.create_task(consumer(agen()), name="async for")
+            nested_task = asyncio.create_task(
+                consumer(outer_agen()), name="nested async for")
+            anext_task = asyncio.create_task(anext(agen()), name="anext")
+            aclose_task = asyncio.create_task(g.aclose(), name="aclose")
+            await asyncio.sleep(0)
+            return RemoteUnwinder(os.getpid()).get_all_awaited_by()
+
+        stacks = self._extract_coroutine_funcnames(asyncio.run(main()))
+        self.assertEqual(stacks["async for"], ["sleep", "agen", "consumer"])
+        self.assertEqual(stacks["nested async for"],
+                         ["sleep", "agen", "outer_agen", "consumer"])
+        self.assertEqual(stacks["anext"], ["sleep", "agen"])
+        self.assertEqual(stacks["aclose"], ["sleep", "agen_with_cleanup"])
+
+    @skip_if_not_supported
+    def test_async_generator_waiter_in_async_stack_trace(self):
+        # gh-159620
+        async def agen(task):
+            await task
+            yield
+
+        async def main():
+            waiter_task = asyncio.create_task(
+                anext(agen(asyncio.current_task())), name="anext")
+            await asyncio.sleep(0)
+            return RemoteUnwinder(os.getpid()).get_async_stack_trace()
+
+        stacks = self._extract_coroutine_funcnames(asyncio.run(main()))
+        self.assertEqual(stacks["anext"], ["agen"])
+
+    @skip_if_not_supported
+    def test_async_generator_running_task_stack(self):
+        # gh-159620
+        async def agen():
+            yield RemoteUnwinder(os.getpid()).get_async_stack_trace()
+
+        async def main():
+            return await asyncio.create_task(anext(agen()), name="anext")
+
+        stacks = self._extract_coroutine_funcnames(asyncio.run(main()))
+        self.assertEqual(stacks["anext"], ["agen"])
+
+    @skip_if_not_supported
+    def test_class_named_like_async_generator_wrapper_is_not_walked(self):
+        # gh-159620
+        async def decoy():
+            await asyncio.sleep(10_000)
+
+        class async_generator_asend:
+            __slots__ = ("gen",)
+            def __init__(self, gen):
+                self.gen = gen
+            def __await__(self):
+                return self
+            def __next__(self):
+                return None
+
+        async def consumer(aw):
+            await aw
+
+        async def main():
+            with contextlib.closing(decoy()) as gen:
+                consumer_task = asyncio.create_task(
+                    consumer(async_generator_asend(gen)), name="consumer")
+                await asyncio.sleep(0)
+                return RemoteUnwinder(os.getpid()).get_all_awaited_by()
+
+        stacks = self._extract_coroutine_funcnames(asyncio.run(main()))
+        self.assertEqual(stacks["consumer"], ["consumer"])
+
     @skip_if_not_supported
     @unittest.skipIf(
         sys.platform == "linux" and not PROCESS_VM_READV_SUPPORTED,

Misc/NEWS.d/next/Library/2026-10-17-15-22-16.gh-issue-159620.Ynsu5L.rst

@@ -0,0 +1,3 @@
+Fix ``python -m asyncio ps``, ``pstree`` and async-aware
+:mod:`profiling.sampling` stacks stopping at asynchronous generators.
+Patch by Timofei Ivankov.

Modules/_remote_debugging/asyncio.c

@@ -240,11 +240,63 @@ parse_task_name(
  * COROUTINE CHAIN PARSING
  * ============================================================================ */
 
+// gh-159620
+/* Return in *coro_addr the coroutine or async generator to walk for the
+   awaitable at addr, or 0 if it is neither (for example, a future). */
+static int
+get_coro_like_address(
+    RemoteUnwinderObject *unwinder,
+    uintptr_t addr,
+    uintptr_t *coro_addr
+) {
+    *coro_addr = 0;
+
+    // Skip user classes that share a name with the generator types below
+    uintptr_t type_addr;
+    unsigned long tp_flags;
+    if (read_ptr(unwinder, addr + (uintptr_t)unwinder->debug_offsets.pyobject.ob_type,
+                 &type_addr) < 0
+        || _Py_RemoteDebug_PagedReadRemoteMemory(&unwinder->handle,
+                type_addr + (uintptr_t)unwinder->debug_offsets.type_object.tp_flags,
+                sizeof(tp_flags), &tp_flags) < 0)
+    {
+        set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to read awaitable type");
+        return -1;
+    }
+
+    if (tp_flags & Py_TPFLAGS_HEAPTYPE) {
+        return 0;
+    }
+
+    uintptr_t name_addr;
+    char name[32] = {0};
+    if (read_ptr(unwinder, type_addr + (uintptr_t)unwinder->debug_offsets.type_object.tp_name,
+                 &name_addr) < 0
+        || _Py_RemoteDebug_PagedReadRemoteMemory(&unwinder->handle, name_addr,
+                                                 sizeof(name) - 1, name) < 0)
+    {
+        set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to read awaitable type name");
+        return -1;
+    }
+
+    if (strcmp(name, "coroutine") == 0) {
+        *coro_addr = addr;
+    }
+    else if (strcmp(name, "async_generator_asend") == 0
+             || strcmp(name, "async_generator_athrow") == 0)
+    {
+        // asend and athrow are separate wrapper objects: read the pointer
+        // to their parent async generator, stored right after PyObject_HEAD
+        return read_ptr(unwinder, addr + (uintptr_t)unwinder->debug_offsets.pyobject.size,
+                        coro_addr);
+    }
+    return 0;
+}
+
 static int
 get_awaited_coro_address(
     RemoteUnwinderObject *unwinder,
     uintptr_t gi_iframe_addr,
-    uintptr_t gen_type_addr,
     uintptr_t *next_coro
 ) {
     *next_coro = 0;
@@ -272,41 +324,14 @@ get_awaited_coro_address(
     uintptr_t stackpointer_addr = GET_MEMBER_NO_TAG(uintptr_t, iframe, unwinder->debug_offsets.interpreter_frame.stackpointer);
 
     if ((void*)stackpointer_addr != NULL) {
-        uintptr_t gi_await_addr;
         err = read_py_ptr(
             unwinder,
             stackpointer_addr - sizeof(void*) * 2,
-            &gi_await_addr);
+            next_coro);
         if (err) {
             set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to read gi_await address");
             return -1;
         }
-
-        if ((void*)gi_await_addr != NULL) {
-            uintptr_t gi_await_addr_type_addr;
-            err = read_ptr(
-                unwinder,
-                gi_await_addr + (uintptr_t)unwinder->debug_offsets.pyobject.ob_type,
-                &gi_await_addr_type_addr);
-            if (err) {
-                set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to read gi_await type address");
-                return -1;
-            }
-
-            if (gen_type_addr == gi_await_addr_type_addr) {
-                /* This needs an explanation. We always start with parsing
-                   native coroutine / generator frames. Ultimately they
-                   are awaiting on something. That something can be
-                   a native coroutine frame or... an iterator.
-                   If it's the latter -- we can't continue building
-                   our chain. So the condition to bail out of this is
-                   to do that when the type of the current coroutine
-                   doesn't match the type of whatever it points to
-                   in its cr_await.
-                */
-                *next_coro = gi_await_addr;
-            }
-        }
     }
 
     return 0;
@@ -329,6 +354,13 @@ parse_coro_chain(
             return -1;
         }
 
+        if (get_coro_like_address(unwinder, coro_address, &coro_address) < 0) {
+            return -1;
+        }
+        if ((void*)coro_address == NULL) {
+            return 0;
+        }
+
         // Read the entire generator object at once
         char gen_object[SIZEOF_GEN_OBJ];
         int err = _Py_RemoteDebug_PagedReadRemoteMemory(
@@ -346,8 +378,6 @@ parse_coro_chain(
             return 0;
         }
 
-        uintptr_t gen_type_addr = GET_MEMBER(uintptr_t, gen_object, unwinder->debug_offsets.pyobject.ob_type);
-
         PyObject* name = NULL;
 
         // Parse the previous frame using the gi_iframe from local copy
@@ -374,8 +404,7 @@ parse_coro_chain(
             return 0;
         }
 
-        if (get_awaited_coro_address(unwinder, gi_iframe_addr, gen_type_addr,
-                                     &coro_address) < 0) {
+        if (get_awaited_coro_address(unwinder, gi_iframe_addr, &coro_address) < 0) {
             return -1;
         }
     }
@@ -733,6 +762,10 @@ get_task_frame(RemoteUnwinderObject *unwinder, uintptr_t task_addr, uintptr_t *f
         return -1;
     }
 
+    if (get_coro_like_address(unwinder, running_coro_addr, &running_coro_addr) < 0) {
+        return -1;
+    }
+
     *frame_addr = running_coro_addr + (uintptr_t)unwinder->debug_offsets.gen_object.gi_iframe;
     return 0;
 }