fix (ai/core): race condition in mergeStreams (#2325)

Author: lgrammel

.changeset/mean-radios-flash.md

@@ -0,0 +1,5 @@
+---
+'ai': patch
+---
+
+fix (ai/core): race condition in mergeStreams

packages/core/core/util/merge-streams.test.ts

@@ -84,3 +84,43 @@ it('should return values from the 2nd stream until the 1st stream has values', a
     '2f',
   ]);
 });
+
+it('should not duplicate last value when parallel calls happen', async () => {
+  let stream1Controller: ReadableStreamDefaultController<string> | undefined;
+  const stream1 = new ReadableStream({
+    start(controller) {
+      stream1Controller = controller;
+    },
+  });
+
+  stream1Controller!.enqueue('1a');
+  stream1Controller!.close();
+
+  let stream2Controller: ReadableStreamDefaultController<string> | undefined;
+  const stream2 = new ReadableStream({
+    start(controller) {
+      stream2Controller = controller;
+    },
+  });
+
+  const mergedStream = mergeStreams(stream1, stream2);
+
+  const reader = mergedStream.getReader();
+
+  const resultsPromise = Promise.all([
+    reader.read(),
+    reader.read(),
+    reader.read(),
+    reader.read(),
+    reader.read(),
+  ]);
+
+  stream2Controller!.enqueue('2a');
+  stream2Controller!.enqueue('2b');
+  stream2Controller!.enqueue('2c');
+  stream2Controller!.close();
+
+  const values = (await resultsPromise).map(result => result.value);
+
+  expect(values).toEqual(['1a', '2a', '2b', '2c', undefined]);
+});

packages/core/core/util/merge-streams.ts

@@ -75,13 +75,13 @@ export function mergeStreams<VALUE1, VALUE2>(
       try {
         // stream 1 is done, we can only read from stream 2:
         if (stream1Done) {
-          readStream2(controller);
+          await readStream2(controller);
           return;
         }
 
         // stream 2 is done, we can only read from stream 1:
         if (stream2Done) {
-          readStream1(controller);
+          await readStream1(controller);
           return;
         }
 
@@ -109,15 +109,15 @@ export function mergeStreams<VALUE1, VALUE2>(
           lastRead1 = undefined;
           if (result.done) {
             // stream 1 is done, we can only read from stream 2:
-            readStream2(controller);
+            await readStream2(controller);
             stream1Done = true;
           }
         } else {
           lastRead2 = undefined;
           // stream 2 is done, we can only read from stream 1:
           if (result.done) {
             stream2Done = true;
-            readStream1(controller);
+            await readStream1(controller);
           }
         }
       } catch (error) {