From 85e88425f2e29383a727b44919f66e3abd851d91 Mon Sep 17 00:00:00 2001
From: Timofei Ivankov <128279579+deadlovelll@users.noreply.github.com>
Date: Fri, 11 Sep 2026 16:54:27 +0300
Subject: [PATCH] gh-157058: Fix missing awaited-by edge in wait_for(fut, 0)

---
 Lib/asyncio/tasks.py                          |  5 ++++
 Lib/test/test_asyncio/test_graph.py           | 27 +++++++++++++++++++
 ...-09-07-12-24-11.gh-issue-157058.2SoQU7.rst |  2 ++
 3 files changed, 34 insertions(+)
 create mode 100644 Misc/NEWS.d/next/Library/2026-09-07-12-24-11.gh-issue-157058.2SoQU7.rst

diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py
index 3cf9cb59fd6a0e9..bd6a03f530afd2e 100644
--- a/Lib/asyncio/tasks.py
+++ b/Lib/asyncio/tasks.py
@@ -541,6 +541,10 @@ async def _cancel_and_wait(fut):
     cb = functools.partial(_release_waiter, waiter)
     fut.add_done_callback(cb)
 
+    # gh-157058: awaiting the waiter leaves no edge on fut, add it here
+    cur_task = current_task()
+    futures.future_add_to_awaited_by(fut, cur_task)
+
     try:
         fut.cancel()
         # We cannot wait on *fut* directly to make
@@ -548,6 +552,7 @@ async def _cancel_and_wait(fut):
         await waiter
     finally:
         fut.remove_done_callback(cb)
+        futures.future_discard_from_awaited_by(fut, cur_task)
 
 
 class _AsCompletedIterator:
diff --git a/Lib/test/test_asyncio/test_graph.py b/Lib/test/test_asyncio/test_graph.py
index b555982beef61c6..ba02263dff0b241 100644
--- a/Lib/test/test_asyncio/test_graph.py
+++ b/Lib/test/test_asyncio/test_graph.py
@@ -146,6 +146,33 @@ async def main():
             'async generator CallStackTestBase.test_stack_async_gen.<locals>.gen()',
             stack_for_gen_nested_call[1])
 
+    async def test_stack_wait_for_non_positive_timeout(self):
+        # gh-157058: wait_for(fut, 0) must still record the waiter
+        cleanup = asyncio.Future()
+
+        async def worker():
+            try:
+                await asyncio.Future()
+            finally:
+                await cleanup
+
+        async def probe(t):
+            await asyncio.wait_for(t, 0)
+
+        t = asyncio.ensure_future(worker())
+        p = asyncio.create_task(probe(t), name='probe')
+        for _ in range(5):
+            await asyncio.sleep(0)
+
+        stack = capture_test_stack(fut=t)
+
+        cleanup.set_result(None)
+        await asyncio.gather(p, t, return_exceptions=True)
+
+        self.assertEqual(stack[0][2], [
+            ['T<probe>', ['a _cancel_and_wait', 'a wait_for', 'a probe'], []],
+        ])
+
     async def test_stack_gather(self):
 
         stack_for_deep = None
diff --git a/Misc/NEWS.d/next/Library/2026-09-07-12-24-11.gh-issue-157058.2SoQU7.rst b/Misc/NEWS.d/next/Library/2026-09-07-12-24-11.gh-issue-157058.2SoQU7.rst
new file mode 100644
index 000000000000000..9122ae0919dcb11
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2026-09-07-12-24-11.gh-issue-157058.2SoQU7.rst
@@ -0,0 +1,2 @@
+:func:`asyncio.wait_for` with a non-positive timeout now records the waiter
+in the call graph.