ai 6.0.259 → 6.0.261

package/src/util/create-stitchable-stream.ts

@@ -8,11 +8,19 @@ import { createResolvablePromise } from './create-resolvable-promise';
  */
 export function createStitchableStream<T>(): {
   stream: ReadableStream<T>;
-  addStream: (innerStream: ReadableStream<T>) => void;
+  addStream: (
+    innerStream: ReadableStream<T>,
+    callbacks?: {
+      onError?: (error: unknown) => void;
+    },
+  ) => void;
   close: () => void;
   terminate: () => void;
 } {
-  let innerStreamReaders: ReadableStreamDefaultReader<T>[] = [];
+  let innerStreams: Array<{
+    reader: ReadableStreamDefaultReader<T>;
+    onError?: (error: unknown) => void;
+  }> = [];
   let controller: ReadableStreamDefaultController<T> | null = null;
   let isClosed = false;
   let waitForNewStream = createResolvablePromise<void>();
@@ -21,34 +29,34 @@ export function createStitchableStream<T>(): {
     isClosed = true;
     waitForNewStream.resolve();
 
-    innerStreamReaders.forEach(reader => reader.cancel());
-    innerStreamReaders = [];
+    innerStreams.forEach(({ reader }) => reader.cancel());
+    innerStreams = [];
     controller?.close();
   };
 
   const processPull = async () => {
     // Case 1: Outer stream is closed and no more inner streams
-    if (isClosed && innerStreamReaders.length === 0) {
+    if (isClosed && innerStreams.length === 0) {
       controller?.close();
       return;
     }
 
     // Case 2: No inner streams available, but outer stream is open
     // wait for a new inner stream to be added or the outer stream to close
-    if (innerStreamReaders.length === 0) {
+    if (innerStreams.length === 0) {
       waitForNewStream = createResolvablePromise<void>();
       await waitForNewStream.promise;
       return processPull();
     }
 
     try {
-      const { value, done } = await innerStreamReaders[0].read();
+      const { value, done } = await innerStreams[0].reader.read();
 
       if (done) {
         // Case 3: Current inner stream is done
-        innerStreamReaders.shift(); // Remove the finished stream
+        innerStreams.shift(); // Remove the finished stream
 
-        if (innerStreamReaders.length === 0 && isClosed) {
+        if (innerStreams.length === 0 && isClosed) {
           // when closed and no more inner streams, stop pulling
           controller?.close();
         } else {
@@ -61,8 +69,9 @@ export function createStitchableStream<T>(): {
       }
     } catch (error) {
       // Case 5: Current inner stream throws an error
+      innerStreams[0].onError?.(error);
       controller?.error(error);
-      innerStreamReaders.shift(); // Remove the errored stream
+      innerStreams.shift(); // Remove the errored stream
       terminate(); // we have errored, terminate all streams
     }
   };
@@ -74,19 +83,27 @@ export function createStitchableStream<T>(): {
       },
       pull: processPull,
       async cancel() {
-        for (const reader of innerStreamReaders) {
+        for (const { reader } of innerStreams) {
           await reader.cancel();
         }
-        innerStreamReaders = [];
+        innerStreams = [];
         isClosed = true;
       },
     }),
-    addStream: (innerStream: ReadableStream<T>) => {
+    addStream: (
+      innerStream: ReadableStream<T>,
+      callbacks?: {
+        onError?: (error: unknown) => void;
+      },
+    ) => {
       if (isClosed) {
         throw new Error('Cannot add inner stream: outer stream is closed');
       }
 
-      innerStreamReaders.push(innerStream.getReader());
+      innerStreams.push({
+        reader: innerStream.getReader(),
+        ...callbacks,
+      });
       waitForNewStream.resolve();
     },
 
@@ -98,7 +115,7 @@ export function createStitchableStream<T>(): {
       isClosed = true;
       waitForNewStream.resolve();
 
-      if (innerStreamReaders.length === 0) {
+      if (innerStreams.length === 0) {
         controller?.close();
       }
     },