TaskSeq.take 2 is returning 3 items, duplicating the last item yielded from the IAsyncEnumerable

[showbiz2k]

I'm experimenting with NATS for pub/sub messaging. The .NET client provides a stream consumer API that exposes a stream of messages as an IAsyncEnumerable. My sample publishes two messages to a stream, and then I attempt to take 2 messages from the stream via an IAsyncEnumerable. I receive both published messages, but the last message is duplicated. If I take 1 message, I receive two messages with the first message being duplicated.

If I implement a custom take with logging, I only see two values being yielded, but still 3 values result.

This behavior only occurs with the IAsyncEnumerable from the NATS .NET client. If I create other IAsyncEnumerables the behavior of the take is as expected.

I'm not sure what the issue could be that would cause a TaskSeq pipeline to behave differently based on the IAsyncEnumerable. Any thoughts on what could be happening here or how to debug?

Publish messages to stream

let createAndPopulateTestStream (js: INatsJSContext) =
    task {
        let streamConfig = new Models.StreamConfig("TEST", [| "test.>" |])
        let! _ = js.CreateStreamAsync streamConfig
        let! ack = js.PublishAsync("test.1", "First message on test.1 subject")
        ack.EnsureSuccess()
        let! ack = js.PublishAsync("test.1", "Second message on test.1 subject")
        ack.EnsureSuccess()
    }

Consume messages from stream

let loadEvents (js: INatsJSContext) (logger: ILogger) stream subject =
    task {
        let consumerOpts = NatsJSOrderedConsumerOpts(FilterSubjects = [| subject |])
        let! streamConsumer = js.CreateOrderedConsumerAsync(stream, consumerOpts)
        let asyncEnumerable = streamConsumer.ConsumeAsync<string>()

        do!
            asyncEnumerable
            |> TaskSeq.take 2
            |> TaskSeq.iter (fun msg -> logger.LogInformation msg.Data)

        let oneToThree = taskSeq { yield! [ 1..3 ] }

        do!
            oneToThree
            |> TaskSeq.take 2
            |> TaskSeq.iter (fun num -> logger.LogInformation $"{num}")

        let infiniteSeq =
            taskSeq {
                while true do
                    yield "infinite"
            }

        do!
            infiniteSeq
            |> TaskSeq.take 2
            |> TaskSeq.iter (fun msg -> logger.LogInformation msg)

    }

Custom take with logging

module TaskSeq =
    let empty<'T> =
        { new IAsyncEnumerable<'T> with
            member _.GetAsyncEnumerator _ =
                { new IAsyncEnumerator<'T> with
                    member _.MoveNextAsync() = ValueTask.False
                    member _.Current = Unchecked.defaultof<'T>
                    member _.DisposeAsync() = ValueTask.CompletedTask } }

    let takeWithLog count (logger: ILogger) (source: TaskSeq<_>) =
        if count = 0 then
            empty
        else
            taskSeq {
                use e = source.GetAsyncEnumerator CancellationToken.None

                for _ in count .. - 1 .. 1 do
                    let! step = e.MoveNextAsync()

                    if not step then
                        failwith "insufficient"

                    logger.LogInformation "yield"
                    yield e.Current
            }

let loadEvents (js: INatsJSContext) (logger: ILogger) stream subject =
    task {
        let consumerOpts = NatsJSOrderedConsumerOpts(FilterSubjects = [| subject |])
        let! streamConsumer = js.CreateOrderedConsumerAsync(stream, consumerOpts)
        let asyncEnumerable = streamConsumer.ConsumeAsync<string>()

        do!
            asyncEnumerable
            |> TaskSeq.takeWithLog 2 logger
            |> TaskSeq.iter (fun msg -> logger.LogInformation msg.Data)

        let oneToThree = taskSeq { yield! [ 1..3 ] }

        do!
            oneToThree
            |> TaskSeq.takeWithLog 2 logger
            |> TaskSeq.iter (fun num -> logger.LogInformation $"{num}")

        let infiniteSeq =
            taskSeq {
                while true do
                    yield "infinite"
            }

        do!
            infiniteSeq
            |> TaskSeq.takeWithLog 2 logger
            |> TaskSeq.iter (fun msg -> logger.LogInformation msg)

    }

Any help is appreciated. Thanks!

[showbiz2k]

In the custom takeWithLog above, if I change the use e = source.GetAsyncEnumerator CancellationToken.None to let e = source.GetAsyncEnumerator I no longer get the duplicated final message. So it appears it has something to do with disposing the enumerator. I'm not sure if there are specific do's or don'ts when implementing the enumerator that would prevent the duplication behavior as I've not seen this type of behavior before.

[showbiz2k]

If I add a reference to System.Linq.Async and use its implementation of Take I get the expected non-duplicated messages from the stream.

let loadEvents (js: INatsJSContext) (logger: ILogger) stream subject =
    task {
        let! lastSeq = getLastSeqNum js stream subject
        let consumerOpts = NatsJSOrderedConsumerOpts(FilterSubjects = [| subject |])
        let! streamConsumer = js.CreateOrderedConsumerAsync(stream, consumerOpts)

        do!
            streamConsumer.ConsumeAsync<string>().Take(2)
            |> TaskSeq.iter (fun msg -> logger.LogInformation msg.Data)

    }