Mutliple streams can be now consumed simultaneously

Author: ericallam

packages/core/src/v3/apiClient/runStream.ts

@@ -97,7 +97,7 @@ export function runShapeStream<TRunTypes extends AnyRunTypes>(
 
 // First, define interfaces for the stream handling
 export interface StreamSubscription {
-  subscribe(onChunk: (chunk: unknown) => Promise<void>): Promise<() => void>;
+  subscribe(): Promise<ReadableStream<unknown>>;
 }
 
 export interface StreamSubscriptionFactory {
@@ -111,33 +111,29 @@ export class SSEStreamSubscription implements StreamSubscription {
     private options: { headers?: Record<string, string>; signal?: AbortSignal }
   ) {}
 
-  async subscribe(onChunk: (chunk: unknown) => Promise<void>): Promise<() => void> {
-    const response = await fetch(this.url, {
+  async subscribe(): Promise<ReadableStream<unknown>> {
+    return fetch(this.url, {
       headers: {
         Accept: "text/event-stream",
         ...this.options.headers,
       },
       signal: this.options.signal,
-    });
-
-    if (!response.body) {
-      throw new Error("No response body");
-    }
-
-    const reader = response.body
-      .pipeThrough(new TextDecoderStream())
-      .pipeThrough(new EventSourceParserStream())
-      .getReader();
-
-    while (true) {
-      const { done, value } = await reader.read();
-
-      if (done) break;
-
-      await onChunk(safeParseJSON(value.data));
-    }
+    }).then((response) => {
+      if (!response.body) {
+        throw new Error("No response body");
+      }
 
-    return () => reader.cancel();
+      return response.body
+        .pipeThrough(new TextDecoderStream())
+        .pipeThrough(new EventSourceParserStream())
+        .pipeThrough(
+          new TransformStream({
+            transform(chunk, controller) {
+              controller.enqueue(safeParseJSON(chunk.data));
+            },
+          })
+        );
+    });
   }
 }
 
@@ -254,13 +250,31 @@ export class RunSubscription<TRunTypes extends AnyRunTypes> {
                 this.options.client?.baseUrl
               );
 
-              await subscription.subscribe(async (chunk) => {
-                controller.enqueue({
-                  type: streamKey,
-                  chunk: chunk as TStreams[typeof streamKey],
-                  run,
-                } as StreamPartResult<RunShape<TRunTypes>, TStreams>);
-              });
+              const stream = await subscription.subscribe();
+
+              // Create the pipeline and start it
+              stream
+                .pipeThrough(
+                  new TransformStream({
+                    transform(chunk, controller) {
+                      controller.enqueue({
+                        type: streamKey,
+                        chunk: chunk as TStreams[typeof streamKey],
+                        run,
+                      } as StreamPartResult<RunShape<TRunTypes>, TStreams>);
+                    },
+                  })
+                )
+                .pipeTo(
+                  new WritableStream({
+                    write(chunk) {
+                      controller.enqueue(chunk);
+                    },
+                  })
+                )
+                .catch((error) => {
+                  console.error(`Error in stream ${streamKey}:`, error);
+                });
             }
           }
         }

packages/core/test/runStream.test.ts

@@ -9,17 +9,23 @@ import {
 import type { SubscribeRunRawShape } from "../src/v3/schemas/api.js";
 
 // Test implementations
+// Update TestStreamSubscription to return a ReadableStream
 class TestStreamSubscription implements StreamSubscription {
   constructor(private chunks: unknown[]) {}
 
-  async subscribe(onChunk: (chunk: unknown) => Promise<void>): Promise<() => void> {
-    for (const chunk of this.chunks) {
-      await onChunk(chunk);
-    }
-    return () => {};
+  async subscribe(): Promise<ReadableStream<unknown>> {
+    return new ReadableStream({
+      start: async (controller) => {
+        for (const chunk of this.chunks) {
+          controller.enqueue(chunk);
+        }
+        controller.close();
+      },
+    });
   }
 }
 
+// TestStreamSubscriptionFactory can remain the same
 class TestStreamSubscriptionFactory implements StreamSubscriptionFactory {
   private streams = new Map<string, unknown[]>();