From c1c57be6769e090f4ac6e7b7bef5fc5b75f4615f Mon Sep 17 00:00:00 2001
From: Timofei Ivankov <128279579+deadlovelll@users.noreply.github.com>
Date: Sat, 29 Aug 2026 11:39:07 +0300
Subject: [PATCH] gh-156523: Fix asyncio.as_completed() not recording the
 awaiting task (GH-156527) (cherry picked from commit
 b0c9fc37dd0ed60456313600e6a01478308f195e)

Co-authored-by: Timofei Ivankov <128279579+deadlovelll@users.noreply.github.com>
---
 Lib/asyncio/tasks.py                          |  4 ++
 Lib/test/test_asyncio/test_graph.py           | 62 +++++++++++++++++++
 ...-08-28-18-10-25.gh-issue-156523.llJqo9.rst |  2 +
 3 files changed, 68 insertions(+)
 create mode 100644 Misc/NEWS.d/next/Library/2026-08-28-18-10-25.gh-issue-156523.llJqo9.rst

diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py
index 0417830467fb212..e989498d6e8632b 100644
--- a/Lib/asyncio/tasks.py
+++ b/Lib/asyncio/tasks.py
@@ -562,9 +562,11 @@ def __init__(self, aws, timeout):
         self._timeout_handle = None
 
         loop = events.get_event_loop()
+        self._cur_task = current_task()
         todo = {ensure_future(aw, loop=loop) for aw in set(aws)}
         for f in todo:
             f.add_done_callback(self._handle_completion)
+            futures.future_add_to_awaited_by(f, self._cur_task)
         if todo and timeout is not None:
             self._timeout_handle = (
                 loop.call_later(timeout, self._handle_timeout)
@@ -595,6 +597,7 @@ def __next__(self):
     def _handle_timeout(self):
         for f in self._todo:
             f.remove_done_callback(self._handle_completion)
+            futures.future_discard_from_awaited_by(f, self._cur_task)
             self._done.put_nowait(None)  # Sentinel for _wait_for_one().
         self._todo.clear()  # Can't do todo.remove(f) in the loop.
 
@@ -602,6 +605,7 @@ def _handle_completion(self, f):
         if not self._todo:
             return  # _handle_timeout() was here first.
         self._todo.remove(f)
+        futures.future_discard_from_awaited_by(f, self._cur_task)
         self._done.put_nowait(f)
         if not self._todo and self._timeout_handle is not None:
             self._timeout_handle.cancel()
diff --git a/Lib/test/test_asyncio/test_graph.py b/Lib/test/test_asyncio/test_graph.py
index 2f22fbccba42bc8..ab7576d25e231ed 100644
--- a/Lib/test/test_asyncio/test_graph.py
+++ b/Lib/test/test_asyncio/test_graph.py
@@ -271,6 +271,68 @@ async def main(t1, t2):
             ]
         ])
 
+    async def test_stack_as_completed(self):
+        # gh-156523: as_completed() must record the awaiting task
+        stack_for_inner = None
+
+        async def inner():
+            await asyncio.sleep(0)
+            nonlocal stack_for_inner
+            stack_for_inner = capture_test_stack()
+
+        async def main(t):
+            for f in asyncio.as_completed([t]):
+                await f
+
+        t = asyncio.create_task(inner(), name='inner')
+        await main(t)
+        self.assertFalse(t._asyncio_awaited_by)
+
+        self.assertEqual(stack_for_inner[0], [
+            'T<inner>',
+            ['s capture_test_stack', 'a inner'],
+            [
+                ['T<anon>',
+                    ['a get', 'a _wait_for_one', 'a main',
+                     'a test_stack_as_completed'],
+                    []
+                ]
+            ]
+        ])
+
+    async def test_stack_as_completed_timeout(self):
+        # gh-156523: the awaiting task must be dropped when as_completed() times out
+        stack_for_inner = None
+
+        async def inner():
+            nonlocal stack_for_inner
+            stack_for_inner = capture_test_stack()
+            await asyncio.sleep(3600)
+
+        async def main(t):
+            with self.assertRaises(TimeoutError):
+                for f in asyncio.as_completed([t], timeout=0.01):
+                    await f
+
+        t = asyncio.create_task(inner(), name='inner')
+        await main(t)
+        self.assertFalse(t._asyncio_awaited_by)
+        t.cancel()
+        with self.assertRaises(asyncio.CancelledError):
+            await t
+
+        self.assertEqual(stack_for_inner[0], [
+            'T<inner>',
+            ['s capture_test_stack', 'a inner'],
+            [
+                ['T<anon>',
+                    ['a get', 'a _wait_for_one', 'a main',
+                     'a test_stack_as_completed_timeout'],
+                    []
+                ]
+            ]
+        ])
+
     async def test_stack_task(self):
 
         stack_for_inner = None
diff --git a/Misc/NEWS.d/next/Library/2026-08-28-18-10-25.gh-issue-156523.llJqo9.rst b/Misc/NEWS.d/next/Library/2026-08-28-18-10-25.gh-issue-156523.llJqo9.rst
new file mode 100644
index 000000000000000..c1c95d76c0d23bb
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2026-08-28-18-10-25.gh-issue-156523.llJqo9.rst
@@ -0,0 +1,2 @@
+Fix :func:`asyncio.as_completed` not recording the awaiting task in the call
+graph.