Align Async.Await with fslang-suggestions #840 (raise TaskCanceledException on cancel)

src/FSharpPlus/Extensions/Extensions.fs

@@ -163,13 +163,13 @@ module Extensions =
         /// at the point where the overall async is started.
         /// </remarks>
         static member Await (task: Task<'T>) : Async<'T> =
-            Async.FromContinuations (fun (sc, ec, cc) ->
+            Async.FromContinuations (fun (sc, ec, _) ->
                 task.ContinueWith (fun (task: Task<'T>) ->
                     if task.IsFaulted then
                         let e = Unchecked.nonNull task.Exception
                         if e.InnerExceptions.Count = 1 then ec e.InnerExceptions[0]
                         else ec e
-                    elif task.IsCanceled then cc (TaskCanceledException ())
+                    elif task.IsCanceled then ec (TaskCanceledException ())
                     else sc task.Result)
                 |> ignore)
 
@@ -190,13 +190,13 @@ module Extensions =
         /// at the point where the overall async is started.
         /// </remarks>
         static member Await (task: Task) : Async<unit> =
-            Async.FromContinuations (fun (sc, ec, cc) ->
+            Async.FromContinuations (fun (sc, ec, _) ->
                 task.ContinueWith (fun (task: Task) ->
                     if task.IsFaulted then
                         let e = Unchecked.nonNull task.Exception
                         if e.InnerExceptions.Count = 1 then ec e.InnerExceptions[0]
                         else ec e
-                    elif task.IsCanceled then cc (TaskCanceledException ())
+                    elif task.IsCanceled then ec (TaskCanceledException ())
                     else sc ())
                 |> ignore)
 

tests/FSharpPlus.Tests/Asyncs.fs

@@ -41,8 +41,9 @@ module Async =
             let t2 = createAsync true 0 2
             let t3 = createAsync true 0 3
 
-            let c = new CancellationToken true
-            let t4 = Task.FromCanceled<int> c |> Async.Await
+            // A genuinely cancelled async: Async.Await no longer surfaces a cancelled Task as a
+            // cancellation, it raises TaskCanceledException (fslang-suggestions #840).
+            let t4 : Async<int> = Async.FromContinuations (fun (_, _, cc) -> cc (OperationCanceledException ()))
 
             let t5 = createAsync false 0 5
             let t6 = createAsync false 0 6
@@ -71,8 +72,9 @@ module Async =
             let t2 = createAsync true 10 2
             let t3 = createAsync true 30 3
 
-            let c = new CancellationToken true
-            let t4 = Task.FromCanceled<int> c |> Async.Await
+            // A genuinely cancelled async: Async.Await no longer surfaces a cancelled Task as a
+            // cancellation, it raises TaskCanceledException (fslang-suggestions #840).
+            let t4 : Async<int> = Async.FromContinuations (fun (_, _, cc) -> cc (OperationCanceledException ()))
 
             let t5 = createAsync false 20 5
             let t6 = createAsync false 10 6

tests/FSharpPlus.Tests/Task.fs

@@ -273,6 +273,30 @@ module Task =
             Assert.AreEqual (e0, e1, "Original exception is not the same as that extracted from the Async")
             Assert.AreEqual (e1, e2, "The exception extracted from the Async is not the same as that extracted from the roundtripped Task")
 
+        [<Test>]
+        let awaitOfCancelledTaskRaisesTaskCanceledException () =
+            // A cancelled Task is surfaced through the exception continuation as a TaskCanceledException, so it can be
+            // caught with an ordinary try/with around the Await, matching C# await (fslang-suggestions #840).
+            let ct = CancellationToken true
+
+            let caughtGeneric =
+                async {
+                    try
+                        let! _ = Async.Await (Task.FromCanceled<int> ct)
+                        return false
+                    with :? TaskCanceledException -> return true }
+                |> Async.RunSynchronously
+            Assert.IsTrue (caughtGeneric, "Await of a cancelled Task<'T> should raise a TaskCanceledException catchable by try/with")
+
+            let caughtNonGeneric =
+                async {
+                    try
+                        do! Async.Await (Task.FromCanceled ct)
+                        return false
+                    with :? TaskCanceledException -> return true }
+                |> Async.RunSynchronously
+            Assert.IsTrue (caughtNonGeneric, "Await of a cancelled Task should raise a TaskCanceledException catchable by try/with")
+
 
     // This module contains tests for ComputationExpression not covered by the below TaskBuilderTests module
     module ComputationExpressionTests =