Fixed consuming realtime runs w/streams after the run is already finished

Author: ericallam

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

@@ -116,11 +116,14 @@ export function runShapeStream<TRunTypes extends AnyRunTypes>(
     { once: true }
   );
 
+  const runStreamInstance = zodShapeStream(SubscribeRunRawShape, url, {
+    ...options,
+    signal: abortController.signal,
+  });
+
   const $options: RunSubscriptionOptions = {
-    runShapeStream: zodShapeStream(SubscribeRunRawShape, url, {
-      ...options,
-      signal: abortController.signal,
-    }),
+    runShapeStream: runStreamInstance.stream,
+    stopRunShapeStream: runStreamInstance.stop,
     streamFactory: new VersionedStreamSubscriptionFactory(version1, version2),
     abortController,
     ...options,
@@ -215,7 +218,7 @@ export class ElectricStreamSubscription implements StreamSubscription {
 
   async subscribe(): Promise<ReadableStream<unknown>> {
     return zodShapeStream(SubscribeRealtimeStreamChunkRawShape, this.url, this.options)
-      .pipeThrough(
+      .stream.pipeThrough(
         new TransformStream({
           transform(chunk, controller) {
             controller.enqueue(chunk.value);
@@ -298,12 +301,12 @@ export interface RunShapeProvider {
 
 export type RunSubscriptionOptions = RunShapeStreamOptions & {
   runShapeStream: ReadableStream<SubscribeRunRawShape>;
+  stopRunShapeStream: () => void;
   streamFactory: StreamSubscriptionFactory;
   abortController: AbortController;
 };
 
 export class RunSubscription<TRunTypes extends AnyRunTypes> {
-  private unsubscribeShape?: () => void;
   private stream: AsyncIterableStream<RunShape<TRunTypes>>;
   private packetCache = new Map<string, any>();
   private _closeOnComplete: boolean;
@@ -330,7 +333,7 @@ export class RunSubscription<TRunTypes extends AnyRunTypes> {
           ) {
             console.log("Closing stream because run is complete");
 
-            this.options.abortController.abort();
+            this.options.stopRunShapeStream();
           }
         },
       },
@@ -342,7 +345,7 @@ export class RunSubscription<TRunTypes extends AnyRunTypes> {
     if (!this.options.abortController.signal.aborted) {
       this.options.abortController.abort();
     }
-    this.unsubscribeShape?.();
+    this.options.stopRunShapeStream();
   }
 
   [Symbol.asyncIterator](): AsyncIterator<RunShape<TRunTypes>> {

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

@@ -16,25 +16,40 @@ export type ZodShapeStreamOptions = {
   signal?: AbortSignal;
 };
 
+export type ZodShapeStreamInstance<TShapeSchema extends z.ZodTypeAny> = {
+  stream: AsyncIterableStream<z.output<TShapeSchema>>;
+  stop: () => void;
+};
+
 export function zodShapeStream<TShapeSchema extends z.ZodTypeAny>(
   schema: TShapeSchema,
   url: string,
   options?: ZodShapeStreamOptions
-) {
-  const stream = new ShapeStream<z.input<TShapeSchema>>({
+): ZodShapeStreamInstance<TShapeSchema> {
+  const abortController = new AbortController();
+
+  options?.signal?.addEventListener(
+    "abort",
+    () => {
+      abortController.abort();
+    },
+    { once: true }
+  );
+
+  const shapeStream = new ShapeStream({
     url,
     headers: {
       ...options?.headers,
       "x-trigger-electric-version": "1.0.0-beta.1",
     },
     fetchClient: options?.fetchClient,
-    signal: options?.signal,
+    signal: abortController.signal,
   });
 
-  const readableShape = new ReadableShapeStream(stream);
+  const readableShape = new ReadableShapeStream(shapeStream);
 
-  return readableShape.stream.pipeThrough(
-    new TransformStream({
+  const stream = readableShape.stream.pipeThrough(
+    new TransformStream<unknown, z.output<TShapeSchema>>({
       async transform(chunk, controller) {
         const result = schema.safeParse(chunk);
 
@@ -46,6 +61,14 @@ export function zodShapeStream<TShapeSchema extends z.ZodTypeAny>(
       },
     })
   );
+
+  return {
+    stream: stream as AsyncIterableStream<z.output<TShapeSchema>>,
+    stop: () => {
+      console.log("Stopping zodShapeStream with abortController.abort()");
+      abortController.abort();
+    },
+  };
 }
 
 export type AsyncIterableStream<T> = AsyncIterable<T> & ReadableStream<T>;
@@ -104,14 +127,19 @@ class ReadableShapeStream<T extends Row<unknown> = Row> {
   readonly #currentState: Map<string, T> = new Map();
   readonly #changeStream: AsyncIterableStream<T>;
   #error: FetchError | false = false;
+  #unsubscribe?: () => void;
+
+  stop() {
+    this.#unsubscribe?.();
+  }
 
   constructor(stream: ShapeStreamInterface<T>) {
     this.#stream = stream;
 
     // Create the source stream that will receive messages
     const source = new ReadableStream<Message<T>[]>({
       start: (controller) => {
-        this.#stream.subscribe(
+        this.#unsubscribe = this.#stream.subscribe(
           (messages) => controller.enqueue(messages),
           this.#handleError.bind(this)
         );
@@ -121,41 +149,44 @@ class ReadableShapeStream<T extends Row<unknown> = Row> {
     // Create the transformed stream that processes messages and emits complete rows
     this.#changeStream = createAsyncIterableStream(source, {
       transform: (messages, controller) => {
-        messages.forEach((message) => {
+        const updatedKeys = new Set<string>();
+
+        for (const message of messages) {
           if (isChangeMessage(message)) {
+            const key = message.key;
             switch (message.headers.operation) {
               case "insert": {
-                this.#currentState.set(message.key, message.value);
-                controller.enqueue(message.value);
+                // New row entirely
+                this.#currentState.set(key, message.value);
+                updatedKeys.add(key);
                 break;
               }
               case "update": {
-                const existingRow = this.#currentState.get(message.key);
-                if (existingRow) {
-                  const updatedRow = {
-                    ...existingRow,
-                    ...message.value,
-                  };
-                  this.#currentState.set(message.key, updatedRow);
-                  controller.enqueue(updatedRow);
-                } else {
-                  this.#currentState.set(message.key, message.value);
-                  controller.enqueue(message.value);
-                }
+                // Merge updates into existing row if any, otherwise treat as new
+                const existingRow = this.#currentState.get(key);
+                const updatedRow = existingRow
+                  ? { ...existingRow, ...message.value }
+                  : message.value;
+                this.#currentState.set(key, updatedRow);
+                updatedKeys.add(key);
                 break;
               }
             }
+          } else if (isControlMessage(message)) {
+            if (message.headers.control === "must-refetch") {
+              this.#currentState.clear();
+              this.#error = false;
+            }
           }
+        }
 
-          if (isControlMessage(message)) {
-            switch (message.headers.control) {
-              case "must-refetch":
-                this.#currentState.clear();
-                this.#error = false;
-                break;
-            }
+        // Now enqueue only one updated row per key, after all messages have been processed.
+        for (const key of updatedKeys) {
+          const finalRow = this.#currentState.get(key);
+          if (finalRow) {
+            controller.enqueue(finalRow);
           }
-        });
+        }
       },
     });
   }

packages/react-hooks/src/hooks/useRealtime.ts

@@ -109,10 +109,13 @@ export function useRealtimeRun<TTask extends AnyTask>(
     }
   }, [runId, mutateRun, abortControllerRef, apiClient, setError]);
 
+  const hasCalledOnCompleteRef = useRef(false);
+
   // Effect to handle onComplete callback
   useEffect(() => {
-    if (isComplete && options?.onComplete && run) {
+    if (isComplete && run && options?.onComplete && !hasCalledOnCompleteRef.current) {
       options.onComplete(run, error);
+      hasCalledOnCompleteRef.current = true;
     }
   }, [isComplete, run, error, options?.onComplete]);
 
@@ -261,10 +264,13 @@ export function useRealtimeRunWithStreams<
     }
   }, [runId, mutateRun, mutateStreams, streamsRef, abortControllerRef, apiClient, setError]);
 
+  const hasCalledOnCompleteRef = useRef(false);
+
   // Effect to handle onComplete callback
   useEffect(() => {
-    if (isComplete && options?.onComplete && run) {
+    if (isComplete && run && options?.onComplete && !hasCalledOnCompleteRef.current) {
       options.onComplete(run, error);
+      hasCalledOnCompleteRef.current = true;
     }
   }, [isComplete, run, error, options?.onComplete]);
 
@@ -593,7 +599,7 @@ async function processRealtimeRunWithStreams<
       nextStreamData[type] = [...(existingDataRef.current[type] || []), ...chunks];
     }
 
-    await mutateStreamData(nextStreamData);
+    mutateStreamData(nextStreamData);
   }, throttleInMs);
 
   for await (const part of subscription.withStreams<TStreams>()) {