Fix bug in EOS for data streams

Runtime/Scripts/DataStreams/ByteDataStream.cs

@@ -145,8 +145,9 @@ internal ReadAllInstruction(ulong asyncId)
         /// YieldInstruction for <see cref="ReadIncremental"/>.
         /// </summary>
         /// <remarks>
-        /// Usage: while <see cref="IsEos"/> is false (i.e. the stream has not ended),
-        /// call <see cref="Reset"/>, yield the instruction, and then access <see cref="Bytes"/>.
+        /// Usage: loop yielding the instruction; when it resumes, break if <see cref="IsEos"/>
+        /// is true (the stream has ended and every chunk has been drained), otherwise read
+        /// <see cref="Bytes"/> and call <see cref="Reset"/> to advance to the next chunk.
         /// </remarks>
         public sealed class ReadIncrementalInstruction : ReadIncrementalInstructionBase<byte[]>
         {

Runtime/Scripts/DataStreams/DataStream.cs

@@ -72,19 +72,43 @@ internal StreamError(Proto.StreamError proto) : base(proto.Description) { }
 
     /// <summary>
     /// Shared state and helpers for incremental stream reader yield instructions.
-    /// Holds the latest chunk, end-of-stream flag, and error; subclasses own the
-    /// typed event subscription and convert raw event payloads via <see cref="OnChunk"/>
-    /// and <see cref="OnEos"/>.
+    /// Models the incoming stream as a single ordered queue whose items are either a
+    /// chunk or a the end-of-stream marker (carrying an optional <see cref="StreamError"/>).
+    /// The consumer advances through the queue one item at a time via <see cref="Reset"/>;
+    /// end-of-stream is only observed once every buffered chunk has been drained, mirroring
+    /// the JS and Python SDKs. Subclasses own the typed event subscription and convert raw
+    /// event payloads via <see cref="OnChunk"/> and <see cref="OnEos"/>.
     /// </summary>
     public abstract class ReadIncrementalInstructionBase<TContent> : StreamYieldInstruction
     {
         private readonly ulong _handleValue;
-        private readonly Queue<TContent> _pendingChunks = new();
-        private TContent _latestChunk;
 
-        // Chunk events arrive on the FFI thread; Reset() and the LatestChunk getter
-        // run on the main-thread coroutine. _gate serializes mutations of the queue,
-        // _latestChunk, IsCurrentReadDone, IsEos, and Error across both sides.
+        // A chunk or the end-of-stream marker. Chunks and the marker share one
+        // FIFO so the consumer can never observe end-of-stream while chunks remain buffered.
+        // The head of the queue is the item the consumer is currently positioned on; Reset()
+        // dequeues it to advance to the next.
+        private readonly struct Item
+        {
+            public readonly TContent Chunk;
+            public readonly bool IsEos;
+            public readonly StreamError Error;
+
+            private Item(TContent chunk, bool isEos, StreamError error)
+            {
+                Chunk = chunk;
+                IsEos = isEos;
+                Error = error;
+            }
+
+            public static Item ForChunk(TContent chunk) => new Item(chunk, false, null);
+            public static Item ForEos(StreamError error) => new Item(default, true, error);
+        }
+
+        private readonly Queue<Item> _queue = new();
+
+        // Chunk/EOS events arrive on the FFI thread; Reset() and the LatestChunk getter run on
+        // the main-thread coroutine. _gate serializes mutations of the queue, IsCurrentReadDone,
+        // IsEos, and Error across both sides.
         private readonly object _gate = new();
 
         /// <summary>
@@ -99,10 +123,11 @@ public abstract class ReadIncrementalInstructionBase<TContent> : StreamYieldInst
 
         /// <summary>
         /// The chunk from the most recent completed read. Throws the captured
-        /// <see cref="StreamError"/> if the last read errored. Internal so the optional
-        /// UniTask async-enumerable adapter (which has InternalsVisibleTo access) can read
-        /// it generically; the typed <c>Bytes</c>/<c>Text</c> accessors on the concrete
-        /// readers delegate here.
+        /// <see cref="StreamError"/> if the stream ended with an error. Returns the default
+        /// value once positioned on a normal end-of-stream marker (there is no chunk to read).
+        /// Internal so the optional UniTask async-enumerable adapter (which has
+        /// InternalsVisibleTo access) can read it generically; the typed <c>Bytes</c>/<c>Text</c>
+        /// accessors on the concrete readers delegate here.
         /// </summary>
         internal TContent LatestChunk
         {
@@ -111,7 +136,7 @@ internal TContent LatestChunk
                 lock (_gate)
                 {
                     if (Error != null) throw Error;
-                    return _latestChunk;
+                    return _queue.Count > 0 ? _queue.Peek().Chunk : default;
                 }
             }
         }
@@ -123,55 +148,57 @@ protected ReadIncrementalInstructionBase(FfiHandle readerHandle)
 
         protected bool MatchesHandle(ulong eventHandle) => eventHandle == _handleValue;
 
-        protected void OnChunk(TContent content)
+        protected void OnChunk(TContent content) => Enqueue(Item.ForChunk(content));
+
+        protected void OnEos(Proto.StreamError protoError) =>
+            Enqueue(Item.ForEos(protoError != null ? new StreamError(protoError) : null));
+
+        private void Enqueue(Item item)
         {
             lock (_gate)
             {
-                if (IsCurrentReadDone)
-                {
-                    // Consumer hasn't yielded since the last chunk; buffer until Reset().
-                    _pendingChunks.Enqueue(content);
-                }
-                else
-                {
-                    _latestChunk = content;
-                    IsCurrentReadDone = true;
-                }
+                // When the queue was empty this item becomes the head, so publish it now;
+                // otherwise it stays buffered in order behind the current head.
+                bool wasEmpty = _queue.Count == 0;
+                _queue.Enqueue(item);
+                if (wasEmpty) SettleHead();
             }
         }
 
         public override void Reset()
         {
-            // base.Reset() must run under the same lock as OnChunk, otherwise the
-            // window between IsCurrentReadDone=false (from base) and the dequeue
-            // below lets a producer race in, write _latestChunk, and have its
-            // chunk immediately overwritten by the dequeue. That race lost ~4% of
-            // chunks under stress before this fix.
+            // base.Reset() must run under the same lock as OnChunk/OnEos so a producer can't
+            // race between the flag clear and the dequeue below. It also throws when the
+            // consumer tries to advance past the Eos marker (IsEos), the intended guard.
             lock (_gate)
             {
                 base.Reset();
-                if (_pendingChunks.Count > 0)
-                {
-                    _latestChunk = _pendingChunks.Dequeue();
-                    IsCurrentReadDone = true;
-                }
+                if (_queue.Count > 0) _queue.Dequeue(); // drop the item just consumed
+                SettleHead();
             }
         }
 
-        protected void OnEos(Proto.StreamError protoError)
+        // Reflects the queue head in the base completion flags: a chunk becomes readable
+        // (IsCurrentReadDone), the Eos marker (IsEos) ends the stream. An empty queue
+        // parks the consumer (keepWaiting stays true) until the next item arrives.
+        // Caller must hold _gate.
+        private void SettleHead()
         {
-            lock (_gate)
+            if (_queue.Count == 0) return;
+            var head = _queue.Peek();
+            if (head.IsEos)
             {
                 // Assign Error before flipping IsEos. The IsEos setter fires the awaiter
-                // continuation, which inspects IsError/Error on resume; when completion runs
-                // inline on the main thread, setting IsEos first would let the continuation
-                // observe IsError == false and silently swallow the stream error.
-                if (protoError != null)
-                {
-                    Error = new StreamError(protoError);
-                }
+                // continuation, which inspects IsError/Error on resume; setting IsEos first
+                // would let the continuation observe IsError == false and silently swallow
+                // the stream error.
+                if (head.Error != null) Error = head.Error;
                 IsEos = true;
             }
+            else
+            {
+                IsCurrentReadDone = true;
+            }
         }
     }
 

Runtime/Scripts/DataStreams/TextDataStream.cs

@@ -136,8 +136,9 @@ public ReadIncrementalInstruction ReadIncremental()
         /// YieldInstruction for <see cref="ReadIncremental"/>.
         /// </summary>
         /// <remarks>
-        /// Usage: while <see cref="IsEos"/> is false (i.e. the stream has not ended),
-        /// call <see cref="Reset"/>, yield the instruction, and then access <see cref="Text"/>.
+        /// Usage: loop yielding the instruction; when it resumes, break if <see cref="IsEos"/>
+        /// is true (the stream has ended and every chunk has been drained), otherwise read
+        /// <see cref="Text"/> and call <see cref="Reset"/> to advance to the next chunk.
         /// </remarks>
         public sealed class ReadIncrementalInstruction : ReadIncrementalInstructionBase<string>
         {

Runtime/Scripts/UniTask/StreamReaderUniTaskExtensions.cs

@@ -18,17 +18,13 @@ public static class StreamReaderUniTaskExtensions
         /// <see cref="TextStreamReader.ReadIncrementalInstruction"/> (<c>string</c>).
         /// </summary>
         /// <remarks>
-        /// Iteration ends when the stream reaches end-of-stream. If the stream ends with an
-        /// error, the enumerable throws that <see cref="StreamError"/> (idiomatic for
-        /// <c>await foreach</c>; this is the one place the UniTask surface throws rather than
-        /// exposing <c>IsError</c>). Cancellation (via the token or the enumerator) surfaces as
-        /// <see cref="System.OperationCanceledException"/> with abandon-awaiter semantics — the
-        /// underlying FFI read is not cancelled on the wire.
-        ///
-        /// The current chunk is delivered on the iteration where end-of-stream is also observed,
-        /// then iteration stops. Chunks buffered <em>beyond</em> the current one when
-        /// end-of-stream arrives are not drainable, because the reader disallows <c>Reset()</c>
-        /// past end-of-stream.
+        /// Iteration ends when the stream reaches end-of-stream. Every buffered chunk is
+        /// drained before iteration stops, since end-of-stream is an ordered item at the tail
+        /// of the chunk queue. If the stream ends with an error, the enumerable throws that
+        /// <see cref="StreamError"/> (idiomatic for <c>await foreach</c>; this is the one place
+        /// the UniTask surface throws rather than exposing <c>IsError</c>). Cancellation (via
+        /// the token or the enumerator) surfaces as <see cref="System.OperationCanceledException"/>
+        /// with abandon-awaiter semantics — the underlying FFI read is not cancelled on the wire.
         /// </remarks>
         public static IUniTaskAsyncEnumerable<TChunk> AsAsyncEnumerable<TChunk>(
             this ReadIncrementalInstructionBase<TChunk> instruction,
@@ -44,30 +40,20 @@ public static IUniTaskAsyncEnumerable<TChunk> AsAsyncEnumerable<TChunk>(
 
                 while (true)
                 {
-                    // Completes when a chunk is ready (IsCurrentReadDone) or the stream ends (IsEos).
+                    // Completes when a chunk is ready (IsCurrentReadDone) or the consumer has
+                    // advanced onto the terminal marker (IsEos) — the two are mutually exclusive.
                     await instruction.AsUniTask(ct);
 
                     if (instruction.IsCurrentReadDone)
                     {
                         var chunk = instruction.LatestChunk;
                         await writer.YieldAsync(chunk);
-
-                        // Re-check IsEos AFTER yielding: end-of-stream may have arrived while
-                        // the consumer was suspended. Reset() throws once IsEos is set, so this
-                        // re-check (not a value captured before the yield) is what keeps the
-                        // loop safe — mirroring the coroutine consumer's "if (IsEos) break;
-                        // else Reset()" ordering.
-                        if (instruction.IsEos)
-                        {
-                            if (instruction.IsError) throw instruction.Error;
-                            return;
-                        }
-
-                        instruction.Reset();
+                        instruction.Reset(); // advance to the next chunk or the terminal marker
                         continue;
                     }
 
-                    // Not IsCurrentReadDone => end-of-stream with nothing left to read.
+                    // Positioned on the terminal marker: the queue is fully drained. Surface an
+                    // error close by throwing, otherwise end iteration cleanly.
                     if (instruction.IsError) throw instruction.Error;
                     return;
                 }

Samples~/Agents/Assets/Runtime/Agent/TranscriptionReader.cs

@@ -69,24 +69,18 @@ IEnumerator PumpNpc(TextStreamReader reader, string identity)
         var append = _beginReply(Speaker.Npc);
         var filter = new AnnotationFilter();
         var r = reader.ReadIncremental();
-        string lastSeen = null;
         while (true)
         {
             yield return r;
+            if (r.IsEos) break;
 
-            // EOS ticks don't refresh Text — only forward when the reference actually changed.
             var chunk = r.Text;
-            if (!ReferenceEquals(chunk, lastSeen))
+            if (!string.IsNullOrEmpty(chunk))
             {
-                if (!string.IsNullOrEmpty(chunk))
-                {
-                    var filtered = filter.Filter(chunk);
-                    if (!string.IsNullOrEmpty(filtered)) append(filtered);
-                }
-                lastSeen = chunk;
+                var filtered = filter.Filter(chunk);
+                if (!string.IsNullOrEmpty(filtered)) append(filtered);
             }
 
-            if (r.IsEos) break;
             r.Reset();
         }
     }