From 8eb843e78c21a28dd79cd70634a8cadc6a967231 Mon Sep 17 00:00:00 2001
From: Timofey Ivankov <playboicartina@gmail.com>
Date: Fri, 28 Aug 2026 18:39:04 +0300
Subject: [PATCH] gh-156523: Fix asyncio.as_completed() not recording the
 awaiting task

---
 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 498eec3f31b292b..f432cf0afa895a2 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 928b618fe5c55b7..220c0179f27bf3b 100644
--- a/Lib/test/test_asyncio/test_graph.py
+++ b/Lib/test/test_asyncio/test_graph.py
@@ -298,6 +298,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.