[Repo Assist] Add AsyncSeq.withCancellation (design parity with TaskSeq, #277) (#278)

* Add AsyncSeq.withCancellation for design parity with TaskSeq (#277)

Adds AsyncSeq.withCancellation, which returns a new AsyncSeq<'T> that
passes the given CancellationToken to GetAsyncEnumerator, overriding
whatever token would otherwise be supplied during iteration.

This mirrors TaskSeq.withCancellation and is useful when consuming
sequences from libraries (e.g. Entity Framework) that accept a
CancellationToken through GetAsyncEnumerator.

Closes #277 (partial - withCancellation addressed; further design parity
investigation tracked in the issue).

Co-authored-by: Copilot <[email protected]>

* ci: trigger checks

---------

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: Copilot <[email protected]>

Author: github-actions[bot]

RELEASE_NOTES.md

@@ -1,3 +1,7 @@
+### 4.10.0
+
+* Added `AsyncSeq.withCancellation` — returns a new `AsyncSeq` that passes the given `CancellationToken` to `GetAsyncEnumerator`, overriding whatever token would otherwise be supplied. Mirrors `TaskSeq.withCancellation` and is useful when consuming sequences from libraries (e.g. Entity Framework) that accept a cancellation token through `GetAsyncEnumerator`. Part of ongoing design-parity work with FSharp.Control.TaskSeq (see #277).
+
 ### 4.9.0
 
 * Performance: `filterAsync` — replaced `asyncSeq`-builder implementation with a direct optimised enumerator, reducing allocation and generator overhead.

src/FSharp.Control.AsyncSeq/AsyncSeq.fs

@@ -2522,6 +2522,13 @@ module AsyncSeq =
             (emptyAsync fillChannelTask)
       }
 
+  /// Returns a new AsyncSeq that passes the given CancellationToken to GetAsyncEnumerator,
+  /// overriding whatever token would otherwise be used. Useful when consuming sequences from
+  /// libraries (such as Entity Framework) that accept a CancellationToken through GetAsyncEnumerator.
+  let withCancellation (cancellationToken: CancellationToken) (source: AsyncSeq<'T>) : AsyncSeq<'T> =
+    { new IAsyncEnumerable<'T> with
+        member _.GetAsyncEnumerator(_ct) = source.GetAsyncEnumerator(cancellationToken) }
+
   #endif
 
 

src/FSharp.Control.AsyncSeq/AsyncSeq.fsi

@@ -824,6 +824,14 @@ module AsyncSeq =
     /// Transforms an async seq to a new one that fetches values ahead of time to improve throughput.
     val prefetch<'T> : numberToPrefetch: int -> source: AsyncSeq<'T> -> AsyncSeq<'T>
 
+    /// <summary>
+    /// Returns a new <c>AsyncSeq</c> that passes the given <c>CancellationToken</c> to
+    /// <c>GetAsyncEnumerator</c>, overriding whatever token would otherwise be used when iterating.
+    /// This is useful when consuming sequences from libraries such as Entity Framework that
+    /// accept a <c>CancellationToken</c> through <c>GetAsyncEnumerator</c>.
+    /// </summary>
+    val withCancellation<'T> : cancellationToken: System.Threading.CancellationToken -> source: AsyncSeq<'T> -> AsyncSeq<'T>
+
     #endif
 
 

tests/FSharp.Control.AsyncSeq.Tests/AsyncSeqTests.fs

@@ -3662,3 +3662,60 @@ let ``AsyncSeq.insertAt raises ArgumentException when index exceeds length`` ()
     |> AsyncSeq.toArrayAsync
     |> Async.RunSynchronously |> ignore)
   |> ignore
+
+// ===== withCancellation =====
+
+[<Test>]
+let ``AsyncSeq.withCancellation passes token to enumerator`` () =
+  use cts = new System.Threading.CancellationTokenSource()
+  let receivedToken = ref System.Threading.CancellationToken.None
+  let source =
+    { new System.Collections.Generic.IAsyncEnumerable<int> with
+        member _.GetAsyncEnumerator(ct) =
+          receivedToken.Value <- ct
+          (AsyncSeq.ofSeq [1; 2; 3]).GetAsyncEnumerator(ct) }
+  source
+  |> AsyncSeq.withCancellation cts.Token
+  |> AsyncSeq.toArrayAsync
+  |> Async.RunSynchronously
+  |> ignore
+  Assert.AreEqual(cts.Token, receivedToken.Value)
+
+[<Test>]
+let ``AsyncSeq.withCancellation overrides incoming token`` () =
+  use cts1 = new System.Threading.CancellationTokenSource()
+  use cts2 = new System.Threading.CancellationTokenSource()
+  let receivedToken = ref System.Threading.CancellationToken.None
+  let source : System.Collections.Generic.IAsyncEnumerable<int> =
+    { new System.Collections.Generic.IAsyncEnumerable<int> with
+        member _.GetAsyncEnumerator(ct) =
+          receivedToken.Value <- ct
+          (AsyncSeq.ofSeq [1; 2; 3]).GetAsyncEnumerator(ct) }
+  let wrapped = source |> AsyncSeq.withCancellation cts1.Token
+  // Enumerate with cts2's token - withCancellation should still pass cts1's token
+  let e = wrapped.GetAsyncEnumerator(cts2.Token)
+  e.MoveNextAsync().AsTask() |> Async.AwaitTask |> Async.RunSynchronously |> ignore
+  e.DisposeAsync() |> ignore
+  Assert.AreEqual(cts1.Token, receivedToken.Value)
+
+[<Test>]
+let ``AsyncSeq.withCancellation preserves sequence values`` () =
+  use cts = new System.Threading.CancellationTokenSource()
+  let result =
+    AsyncSeq.ofSeq [1; 2; 3; 4; 5]
+    |> AsyncSeq.withCancellation cts.Token
+    |> AsyncSeq.toArrayAsync
+    |> Async.RunSynchronously
+  Assert.AreEqual([| 1; 2; 3; 4; 5 |], result)
+
+[<Test>]
+let ``AsyncSeq.withCancellation with cancelled token raises OperationCanceledException`` () =
+  use cts = new System.Threading.CancellationTokenSource()
+  cts.Cancel()
+  Assert.Catch<System.OperationCanceledException>(fun () ->
+    AsyncSeq.ofSeq [1; 2; 3]
+    |> AsyncSeq.withCancellation cts.Token
+    |> AsyncSeq.toArrayAsync
+    |> Async.RunSynchronously
+    |> ignore)
+  |> ignore

version.props

@@ -1,5 +1,5 @@
 <Project>
   <PropertyGroup>
-    <Version>4.8.0</Version>
+    <Version>4.10.0</Version>
   </PropertyGroup>
 </Project>