fix(python): settle cancelled tasks awaited by async workflows

src/Fable.Transforms/Python/Replacements.fs

@@ -782,6 +782,8 @@ let tryEntityIdent (com: Compiler) entFullName =
     // | BuiltinDefinition FSharpMap _ -> fail "Map" // TODO:
     | Types.matchFail -> makeImportLib com Any "MatchFailureException" "Types" |> Some
     | Types.exception_ -> makeIdentExpr "Exception" |> Some
+    | "System.Threading.Tasks.TaskCanceledException" ->
+        makeImportLib com Any "TaskCanceledError" "async_builder" |> Some
     // | "System.Collections.Generic.KeyNotFoundException" ->
     //     makeImportLib com Any "KeyNotFoundException" "System.Collections.Generic" |> Some
     | BuiltinSystemException entName -> makeImportLib com Any entName "System" |> Some
@@ -2844,6 +2846,9 @@ let exceptions (com: ICompiler) (ctx: Context) r t (i: CallInfo) (thisArg: Expr
     match i.CompiledName, thisArg with
     | ".ctor", _ ->
         match i.DeclaringEntityFullName with
+        | "System.Threading.Tasks.TaskCanceledException" ->
+            Helper.LibCall(com, "async_builder", "TaskCanceledError", t, args, isConstructor = true, ?loc = r)
+            |> Some
         // | "System.Collections.Generic.KeyNotFoundException"
         | BuiltinSystemException _ -> bclType com ctx r t i thisArg args
         | _ ->
@@ -3675,6 +3680,9 @@ let asyncBuilder (com: ICompiler) (ctx: Context) r t (i: CallInfo) (thisArg: Exp
 
 let asyncs com (ctx: Context) r t (i: CallInfo) (_: Expr option) (args: Expr list) =
     match i.CompiledName with
+    | "AwaitTask" ->
+        Helper.LibCall(com, "async_", "await_task_with_cancellation", t, args, i.SignatureArgTypes, ?loc = r)
+        |> Some
     | "FromContinuations" when
         (match t with
          | Fable.DeclaredType(_, [ Fable.Unit ]) -> true

src/fable-library-py/fable_library/async_.py

@@ -16,6 +16,7 @@
     Continuations,
     IAsyncContext,
     OperationCanceledError,
+    TaskCanceledError,
     Trampoline,
     empty_continuation,
     protected_bind,
@@ -281,6 +282,17 @@ def callback(conts: Continuations[T]) -> None:
     return from_continuations(callback)
 
 
+def await_task_with_cancellation[T](task: Awaitable[T]) -> Async[T]:
+    async def observe() -> T:
+        try:
+            return await task
+        except asyncio.CancelledError as error:
+            # decision: cancelled tasks are ordinary AwaitTask failures, matching .NET TaskCanceledException
+            raise TaskCanceledError() from error
+
+    return await_task(observe())
+
+
 def run_in_loop(computation: Callable[..., None]) -> Any:
     """Run a computation on the event loop.
 
@@ -461,6 +473,7 @@ def exception_handler(loop: asyncio.AbstractEventLoop, context: dict[str, Any])
 __all__ = [
     "await_event",
     "await_task",
+    "await_task_with_cancellation",
     "cancel",
     "cancel_after",
     "cancellation_token",

src/fable-library-py/fable_library/async_builder.py

@@ -20,6 +20,12 @@ def __init__(self, msg: str | None = None) -> None:
         super().__init__(msg or "The operation was canceled")
 
 
+class TaskCanceledError(OperationCanceledError):
+    def __init__(self, msg: str | None = None, inner_exception: Exception | None = None) -> None:
+        Exception.__init__(self, "A task was canceled." if msg is None else msg)
+        self.inner_exception: Exception | None = inner_exception
+
+
 type Continuations[T] = tuple[
     Callable[[T], None],
     Callable[[Exception], None],

tests/Python/TestAsync.fs

@@ -345,26 +345,34 @@ let ``test Interaction between Async and Task works`` () =
         equal true res
     } |> Async.RunSynchronously
 
-#if FABLE_COMPILER
 [<Fact>]
 let ``test Tasks can be cancelled`` () =
     async {
-        let mutable res = 0
-        let tcs = new System.Threading.CancellationTokenSource(50)
-        let work =
-            let work = async {
-                do! Async.Sleep 75
-                res <- -1
-            }
-            Async.StartAsTask(work, cancellationToken=tcs.Token) |> Async.AwaitTask
-        // behavior change: a cancelled task is now triggering the exception continuation instead of
-        // the cancellation continuation, see: https://github.com/Microsoft/visualfsharp/issues/1416
-        // also, System.OperationCanceledException will be changed to TaskCanceledException (not yet)
-        Async.StartWithContinuations(work, ignore, ignore, (fun _ -> res <- 1))
-        do! Async.Sleep 100
-        equal 1 res
-    } |> Async.StartImmediate
-#endif
+        use source = new System.Threading.CancellationTokenSource()
+        let entered = System.Threading.Tasks.TaskCompletionSource<int>()
+        let work = async {
+            entered.SetResult 1
+            do! Async.Sleep Int32.MaxValue
+            failwith "Cancelled task resumed"
+        }
+        let task = Async.StartAsTask(work, cancellationToken=source.Token)
+        let! _ = Async.AwaitTask entered.Task
+        source.Cancel()
+        let! result = Async.Catch(Async.AwaitTask task)
+        match result with
+        | Choice2Of2 (:? System.Threading.Tasks.TaskCanceledException) -> ()
+        | _ -> failwith "Expected TaskCanceledException through the error continuation"
+    } |> Async.RunSynchronously
+
+[<Fact>]
+let ``test TaskCanceledException preserves message and inner exception`` () =
+    let inner = Exception("inner")
+    let error: exn = System.Threading.Tasks.TaskCanceledException("task failure", inner)
+    match error with
+    | :? System.Threading.Tasks.TaskCanceledException ->
+        equal "task failure" error.Message
+        equal true (obj.ReferenceEquals(inner, error.InnerException))
+    | _ -> failwith "Expected TaskCanceledException"
 
 [<Fact>]
 let ``test Deep recursion with async doesn't cause stack overflow`` () =

tests/Python/test_async_runtime.py

@@ -2,8 +2,15 @@
 from threading import Barrier, Condition, Event, Thread
 
 import pytest
-from fable_library.async_ import await_event, sleep, start_with_continuations
-from fable_library.async_builder import CancellationCallbackError, CancellationToken, IAsyncContext, Trampoline
+from fable_library.async_ import await_event, await_task_with_cancellation, sleep, start_with_continuations
+from fable_library.async_builder import (
+    CancellationCallbackError,
+    CancellationToken,
+    IAsyncContext,
+    OperationCanceledError,
+    TaskCanceledError,
+    Trampoline,
+)
 from fable_library.event import Event as FableEvent
 
 
@@ -165,6 +172,46 @@ def test_sleep_on_cancelled_token_settles_once():
     assert calls == ["cancel"]
 
 
+@pytest.mark.parametrize("already_cancelled", [False, True])
+def test_await_task_cancellation_uses_error_continuation(already_cancelled: bool) -> None:
+    async def run() -> None:
+        loop = asyncio.get_running_loop()
+        unhandled: list[object] = []
+        loop.set_exception_handler(lambda _, context: unhandled.append(context.get("exception")))
+        task: asyncio.Future[int] = loop.create_future()
+        token = CancellationToken()
+        calls: list[str] = []
+        errors: list[Exception] = []
+        if already_cancelled:
+            task.cancel()
+
+        def on_error(error: Exception) -> None:
+            calls.append("error")
+            errors.append(error)
+
+        start_with_continuations(
+            await_task_with_cancellation(task),
+            lambda _: calls.append("success"),
+            on_error,
+            lambda _: calls.append("cancel"),
+            token,
+        )
+        await asyncio.sleep(0)
+        if not already_cancelled:
+            task.cancel()
+        await asyncio.sleep(0)
+        await asyncio.sleep(0)
+        assert calls == ["error"]
+        assert len(errors) == 1
+        assert isinstance(errors[0], TaskCanceledError)
+        assert isinstance(errors[0], OperationCanceledError)
+        assert errors[0].inner_exception is None
+        assert not token.is_cancelled
+        assert not unhandled
+
+    asyncio.run(run())
+
+
 class _DeferringTrampoline(Trampoline):
     def increment_and_check(self) -> bool:
         return True