From eb458af0208adaf3c5be35ddde1b221e3059e926 Mon Sep 17 00:00:00 2001
From: Timofey Ivankov <playboicartina@gmail.com>
Date: Mon, 7 Sep 2026 12:24:49 +0300
Subject: [PATCH 1/2] gh-157058: Fix missing awaited-by edge in wait_for(fut,
 0)

---
 Lib/asyncio/tasks.py                          |  7 +++++
 Lib/test/test_asyncio/test_graph.py           | 27 +++++++++++++++++++
 ...-09-07-12-24-11.gh-issue-157058.2SoQU7.rst |  2 ++
 3 files changed, 36 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 f432cf0afa895a2..adaa11a7554bfb0 100644
--- a/Lib/asyncio/tasks.py
+++ b/Lib/asyncio/tasks.py
@@ -541,6 +541,11 @@ 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()
+    if cur_task is not None:
+        futures.future_add_to_awaited_by(fut, cur_task)
+
     try:
         fut.cancel()
         # We cannot wait on *fut* directly to make
@@ -548,6 +553,8 @@ async def _cancel_and_wait(fut):
         await waiter
     finally:
         fut.remove_done_callback(cb)
+        if cur_task is not None:
+            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 36841672e1f0f65..51f93e91bfbd6fd 100644
--- a/Lib/test/test_asyncio/test_graph.py
+++ b/Lib/test/test_asyncio/test_graph.py
@@ -173,6 +173,33 @@ class FakeCoro:
 
         self.assertEqual(len(result.call_stack), 2)
 
+    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.

From 97add69b3d488713d6e1a4f62981a5fa12ebd928 Mon Sep 17 00:00:00 2001
From: Timofey Ivankov <playboicartina@gmail.com>
Date: Mon, 7 Sep 2026 15:20:01 +0300
Subject: [PATCH 2/2] remove None-checks

---
 Lib/asyncio/tasks.py | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py
index adaa11a7554bfb0..553415e0eef034b 100644
--- a/Lib/asyncio/tasks.py
+++ b/Lib/asyncio/tasks.py
@@ -543,8 +543,7 @@ async def _cancel_and_wait(fut):
 
     # gh-157058: awaiting the waiter leaves no edge on fut, add it here
     cur_task = current_task()
-    if cur_task is not None:
-        futures.future_add_to_awaited_by(fut, cur_task)
+    futures.future_add_to_awaited_by(fut, cur_task)
 
     try:
         fut.cancel()
@@ -553,8 +552,7 @@ async def _cancel_and_wait(fut):
         await waiter
     finally:
         fut.remove_done_callback(cb)
-        if cur_task is not None:
-            futures.future_discard_from_awaited_by(fut, cur_task)
+        futures.future_discard_from_awaited_by(fut, cur_task)
 
 
 class _AsCompletedIterator: