Non-blocking await for .NET Task in Python async coroutines

Instead of blocking the thread with GetAwaiter().GetResult(),
TaskAwaitable.__next__ now yields the Task back to the runner
when it's not yet completed. The runner can then wait on the Task
and resume the coroutine, enabling true concurrency between coroutines.

Author: mikasoukhov

src/core/IronPython/Runtime/AsyncEnumerableWrapper.cs

@@ -43,12 +43,13 @@ public object __anext__() {
 
     /// <summary>
     /// The awaitable object returned by <see cref="AsyncEnumeratorWrapper{T}.__anext__"/>.
-    /// Implements both __await__ and __iter__/__next__ (the yield-from protocol) to
-    /// block on MoveNextAsync and return the current value via StopIteration.
+    /// Implements both __await__ and __iter__/__next__ (the yield-from protocol).
+    /// Non-blocking: yields the Task back to the runner if MoveNextAsync is not yet completed.
     /// </summary>
     [PythonType("async_enumerator_awaitable")]
     public sealed class AsyncEnumeratorAwaitable<T> {
         private readonly IAsyncEnumerator<T> _enumerator;
+        private Task<bool>? _moveNextTask;
 
         internal AsyncEnumeratorAwaitable(IAsyncEnumerator<T> enumerator) {
             _enumerator = enumerator;
@@ -60,7 +61,10 @@ internal AsyncEnumeratorAwaitable(IAsyncEnumerator<T> enumerator) {
 
         [LightThrowing]
         public object __next__() {
-            bool hasNext = _enumerator.MoveNextAsync().AsTask().GetAwaiter().GetResult();
+            var task = _moveNextTask ??= _enumerator.MoveNextAsync().AsTask();
+            if (!task.IsCompleted) return (Task)task;  // yield Task to runner
+            bool hasNext = task.GetAwaiter().GetResult();
+            _moveNextTask = null;  // reset for next call
             if (!hasNext) {
                 return LightExceptions.Throw(
                     new PythonExceptions._StopAsyncIteration().InitAndGetClrException());

src/core/IronPython/Runtime/TaskAwaitable.cs

@@ -15,7 +15,7 @@ namespace IronPython.Runtime {
     /// <summary>
     /// Provides an __await__ protocol wrapper for <see cref="Task"/>,
     /// enabling <c>await task</c> from Python async code.
-    /// Blocks on the task and returns None via StopIteration.
+    /// Non-blocking: yields the Task back to the runner if not yet completed.
     /// </summary>
     [PythonType("task_awaitable")]
     public sealed class TaskAwaitable {
@@ -31,15 +31,17 @@ internal TaskAwaitable(Task task) {
 
         [LightThrowing]
         public object __next__() {
-            _task.GetAwaiter().GetResult();
+            var task = _task;
+            if (!task.IsCompleted) return task;  // yield Task to runner
+            task.GetAwaiter().GetResult();       // propagate exceptions
             return LightExceptions.Throw(new PythonExceptions._StopIteration().InitAndGetClrException());
         }
     }
 
     /// <summary>
     /// Provides an __await__ protocol wrapper for <see cref="Task{T}"/>,
     /// enabling <c>result = await task</c> from Python async code.
-    /// Blocks on the task and returns the result via StopIteration.value.
+    /// Non-blocking: yields the Task back to the runner if not yet completed.
     /// </summary>
     [PythonType("task_awaitable")]
     public sealed class TaskAwaitable<T> {
@@ -55,7 +57,9 @@ internal TaskAwaitable(Task<T> task) {
 
         [LightThrowing]
         public object __next__() {
-            T result = _task.GetAwaiter().GetResult();
+            var task = _task;
+            if (!task.IsCompleted) return task;  // yield Task to runner
+            T result = task.GetAwaiter().GetResult();
             return LightExceptions.Throw(new PythonExceptions._StopIteration().InitAndGetClrException(result!));
         }
     }
@@ -64,10 +68,12 @@ public object __next__() {
     /// <summary>
     /// Provides an __await__ protocol wrapper for <see cref="ValueTask"/>,
     /// enabling <c>await valuetask</c> from Python async code.
+    /// Non-blocking: yields the Task back to the runner if not yet completed.
     /// </summary>
     [PythonType("task_awaitable")]
     public sealed class ValueTaskAwaitable {
         private readonly ValueTask _task;
+        private Task? _asTask;
 
         internal ValueTaskAwaitable(ValueTask task) {
             _task = task;
@@ -79,18 +85,22 @@ internal ValueTaskAwaitable(ValueTask task) {
 
         [LightThrowing]
         public object __next__() {
-            _task.AsTask().GetAwaiter().GetResult();
+            var task = _asTask ??= _task.AsTask();
+            if (!task.IsCompleted) return task;  // yield Task to runner
+            task.GetAwaiter().GetResult();       // propagate exceptions
             return LightExceptions.Throw(new PythonExceptions._StopIteration().InitAndGetClrException());
         }
     }
 
     /// <summary>
     /// Provides an __await__ protocol wrapper for <see cref="ValueTask{T}"/>,
     /// enabling <c>result = await valuetask</c> from Python async code.
+    /// Non-blocking: yields the Task back to the runner if not yet completed.
     /// </summary>
     [PythonType("task_awaitable")]
     public sealed class ValueTaskAwaitable<T> {
         private readonly ValueTask<T> _task;
+        private Task<T>? _asTask;
 
         internal ValueTaskAwaitable(ValueTask<T> task) {
             _task = task;
@@ -102,7 +112,9 @@ internal ValueTaskAwaitable(ValueTask<T> task) {
 
         [LightThrowing]
         public object __next__() {
-            T result = _task.AsTask().GetAwaiter().GetResult();
+            var task = _asTask ??= _task.AsTask();
+            if (!task.IsCompleted) return (Task)task;  // yield Task to runner
+            T result = task.GetAwaiter().GetResult();
             return LightExceptions.Throw(new PythonExceptions._StopIteration().InitAndGetClrException(result!));
         }
     }

tests/IronPython.Tests/AsyncInteropHelpers.cs

@@ -47,6 +47,64 @@ private static async IAsyncEnumerable<string> YieldStrings(
             }
         }
 
+        /// <summary>
+        /// Returns a real async Task&lt;int&gt; with a delay.
+        /// The runtime type will be AsyncStateMachineBox, not Task&lt;int&gt; directly.
+        /// </summary>
+        public static async Task<int> GetAsyncInt(int value, int delayMs = 50) {
+            await Task.Delay(delayMs);
+            return value;
+        }
+
+        /// <summary>
+        /// Returns a real async Task&lt;string&gt; with a delay.
+        /// </summary>
+        public static async Task<string> GetAsyncString(string value, int delayMs = 50) {
+            await Task.Delay(delayMs);
+            return value;
+        }
+
+        /// <summary>
+        /// Returns a real async Task (void result) with a delay.
+        /// </summary>
+        public static async Task DoAsync(int delayMs = 50) {
+            await Task.Delay(delayMs);
+        }
+
+        /// <summary>
+        /// Async Task&lt;int&gt; that respects a CancellationToken.
+        /// Throws OperationCanceledException if token is cancelled during the delay.
+        /// </summary>
+        public static async Task<int> GetAsyncIntWithCancellation(int value, CancellationToken token, int delayMs = 5000) {
+            await Task.Delay(delayMs, token);
+            return value;
+        }
+
+        /// <summary>
+        /// Async Task that respects a CancellationToken.
+        /// </summary>
+        public static async Task DoAsyncWithCancellation(CancellationToken token, int delayMs = 5000) {
+            await Task.Delay(delayMs, token);
+        }
+
+        /// <summary>
+        /// IAsyncEnumerable&lt;int&gt; that yields values with delay and respects cancellation.
+        /// </summary>
+        public static IAsyncEnumerable<int> GetAsyncIntsWithCancellation(CancellationToken token, params int[] values) {
+            return YieldIntsWithCancellation(values, token);
+        }
+
+        private static async IAsyncEnumerable<int> YieldIntsWithCancellation(
+            int[] values,
+            CancellationToken token,
+            [EnumeratorCancellation] CancellationToken ct = default) {
+            using var linked = CancellationTokenSource.CreateLinkedTokenSource(token, ct);
+            foreach (var v in values) {
+                await Task.Delay(50, linked.Token);
+                yield return v;
+            }
+        }
+
         /// <summary>
         /// Returns an empty IAsyncEnumerable&lt;int&gt;.
         /// </summary>

tests/suite/test_async.py

@@ -10,12 +10,16 @@
 
 
 def run_coro(coro):
-    """Run a coroutine to completion and return its result."""
-    try:
-        coro.send(None)
-        raise AssertionError("Coroutine did not raise StopIteration")
-    except StopIteration as e:
-        return e.value
+    """Run a coroutine to completion, blocking on yielded .NET Tasks."""
+    value = None
+    while True:
+        try:
+            task = coro.send(value)
+            # .NET Task yielded — block on it (test runner is synchronous)
+            task.Wait()
+            value = None
+        except StopIteration as e:
+            return e.value
 
 
 class AsyncIter:
@@ -518,4 +522,117 @@ def test_async_for_has_aiter(self):
         self.assertTrue(hasattr(stream, '__aiter__'))
 
 
+@unittest.skipUnless(_has_async_helpers, "requires IronPythonTest with AsyncInteropHelpers")
+class DotNetRealAsyncTaskTest(unittest.TestCase):
+    """Tests for await on real async Task<T> methods.
+
+    These test real .NET async methods where the runtime type is
+    AsyncStateMachineBox<TResult, TStateMachine>, not Task<T> directly.
+    All methods include real delays (Task.Delay) to ensure truly async behavior.
+    """
+
+    def test_await_real_async_int(self):
+        """await a real async Task<int> with delay."""
+        async def test():
+            return await AsyncInteropHelpers.GetAsyncInt(42)
+        self.assertEqual(run_coro(test()), 42)
+
+    def test_await_real_async_string(self):
+        """await a real async Task<string> with delay."""
+        async def test():
+            return await AsyncInteropHelpers.GetAsyncString("hello")
+        self.assertEqual(run_coro(test()), "hello")
+
+    def test_await_real_async_void(self):
+        """await a real async Task (void) with delay."""
+        async def test():
+            await AsyncInteropHelpers.DoAsync()
+            return 'done'
+        self.assertEqual(run_coro(test()), 'done')
+
+    def test_await_real_async_multiple(self):
+        """Multiple awaits on real async Task<T> in sequence."""
+        async def test():
+            a = await AsyncInteropHelpers.GetAsyncInt(10)
+            b = await AsyncInteropHelpers.GetAsyncInt(20)
+            s = await AsyncInteropHelpers.GetAsyncString("!")
+            return str(a + b) + s
+        self.assertEqual(run_coro(test()), "30!")
+
+    def test_await_real_async_mixed_with_python(self):
+        """Mix real .NET async with Python coroutines."""
+        async def py_double(x):
+            return x * 2
+
+        async def test():
+            val = await AsyncInteropHelpers.GetAsyncInt(5)
+            doubled = await py_double(val)
+            return doubled
+        self.assertEqual(run_coro(test()), 10)
+
+
+@unittest.skipUnless(_has_async_helpers, "requires IronPythonTest with AsyncInteropHelpers")
+class DotNetCancellationTest(unittest.TestCase):
+    """Tests for CancellationToken and CancelledError with .NET async methods."""
+
+    def test_cancel_async_task_int(self):
+        """Cancelling a Task<int> should raise CancelledError."""
+        from System.Threading import CancellationTokenSource
+        cts = CancellationTokenSource()
+        cts.Cancel()
+        async def test():
+            return await AsyncInteropHelpers.GetAsyncIntWithCancellation(42, cts.Token)
+        with self.assertRaises(CancelledError):
+            run_coro(test())
+
+    def test_cancel_async_task_void(self):
+        """Cancelling a Task (void) should raise CancelledError."""
+        from System.Threading import CancellationTokenSource
+        cts = CancellationTokenSource()
+        cts.Cancel()
+        async def test():
+            await AsyncInteropHelpers.DoAsyncWithCancellation(cts.Token)
+        with self.assertRaises(CancelledError):
+            run_coro(test())
+
+    def test_cancel_async_enumerable(self):
+        """Cancelling during async for over IAsyncEnumerable should raise CancelledError."""
+        from System.Threading import CancellationTokenSource
+        cts = CancellationTokenSource()
+        cts.Cancel()
+        async def test():
+            result = []
+            async for x in AsyncInteropHelpers.GetAsyncIntsWithCancellation(cts.Token, 1, 2, 3):
+                result.append(x)
+            return result
+        with self.assertRaises(CancelledError):
+            run_coro(test())
+
+    def test_cancelled_error_is_exception_subclass(self):
+        """CancelledError should be a subclass of Exception."""
+        self.assertTrue(issubclass(CancelledError, Exception))
+
+    def test_cancelled_error_catch_as_exception(self):
+        """CancelledError should be catchable as Exception."""
+        from System.Threading import CancellationTokenSource
+        cts = CancellationTokenSource()
+        cts.Cancel()
+        async def test():
+            try:
+                await AsyncInteropHelpers.GetAsyncIntWithCancellation(99, cts.Token)
+            except Exception:
+                return 'caught'
+        self.assertEqual(run_coro(test()), 'caught')
+
+    def test_operation_cancelled_maps_to_cancelled_error(self):
+        """System.OperationCanceledException raised directly should be catchable as CancelledError."""
+        from System import OperationCanceledException
+        caught = False
+        try:
+            raise OperationCanceledException("test")
+        except CancelledError:
+            caught = True
+        self.assertTrue(caught)
+
+
 run_test(__name__)