Return stream parts from SSE class

Author: ericallam

apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.runs.$runParam.streams.$streamKey/route.tsx

@@ -1,5 +1,5 @@
 import { LoaderFunctionArgs } from "@remix-run/server-runtime";
-import { SSEStreamSubscription } from "@trigger.dev/core/v3";
+import { SSEStreamPart, SSEStreamSubscription } from "@trigger.dev/core/v3";
 import { useEffect, useRef, useState } from "react";
 import { Paragraph } from "~/components/primitives/Paragraph";
 import { $replica } from "~/db.server";
@@ -12,6 +12,7 @@ import { cn } from "~/utils/cn";
 import { v3RunStreamParamsSchema } from "~/utils/pathBuilder";
 
 type StreamChunk = {
+  id: string;
   data: unknown;
   timestamp: number;
 };
@@ -261,7 +262,7 @@ function useRealtimeStream(resourcePath: string, startIndex?: number) {
 
   useEffect(() => {
     const abortController = new AbortController();
-    let reader: ReadableStreamDefaultReader<unknown> | null = null;
+    let reader: ReadableStreamDefaultReader<SSEStreamPart<unknown>> | null = null;
 
     async function connectAndConsume() {
       try {
@@ -288,8 +289,9 @@ function useRealtimeStream(resourcePath: string, startIndex?: number) {
             setChunks((prev) => [
               ...prev,
               {
-                data: value,
-                timestamp: Date.now(),
+                id: value.id,
+                data: value.chunk,
+                timestamp: value.timestamp,
               },
             ]);
           }

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

@@ -74,6 +74,7 @@ import {
   TaskRunShape,
   runShapeStream,
   RealtimeRunSkipColumns,
+  type SSEStreamPart,
 } from "./runStream.js";
 import {
   CreateEnvironmentVariableParams,
@@ -142,6 +143,7 @@ export type {
   RunStreamCallback,
   RunSubscription,
   TaskRunShape,
+  SSEStreamPart,
 };
 
 export * from "./getBranch.js";
@@ -1095,7 +1097,13 @@ export class ApiClient {
 
     const stream = await subscription.subscribe();
 
-    return stream as AsyncIterableStream<T>;
+    return stream.pipeThrough(
+      new TransformStream<SSEStreamPart, T>({
+        transform(chunk, controller) {
+          controller.enqueue(chunk.chunk as T);
+        },
+      })
+    );
   }
 
   async createStream(

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

@@ -1,12 +1,12 @@
-import { EventSourceParserStream } from "eventsource-parser/stream";
+import { EventSourceMessage, EventSourceParserStream } from "eventsource-parser/stream";
 import { DeserializedJson } from "../../schemas/json.js";
 import { createJsonErrorObject } from "../errors.js";
-import {
-  RunStatus,
-  SubscribeRealtimeStreamChunkRawShape,
-  SubscribeRunRawShape,
-} from "../schemas/api.js";
+import { RunStatus, SubscribeRunRawShape } from "../schemas/api.js";
 import { SerializedError } from "../schemas/common.js";
+import {
+  AsyncIterableStream,
+  createAsyncIterableReadable,
+} from "../streams/asyncIterableStream.js";
 import { AnyRunTypes, AnyTask, InferRunTypes } from "../types/tasks.js";
 import { getEnvVar } from "../utils/getEnv.js";
 import {
@@ -16,11 +16,7 @@ import {
 } from "../utils/ioSerialization.js";
 import { ApiError } from "./errors.js";
 import { ApiClient } from "./index.js";
-import { LineTransformStream, zodShapeStream } from "./stream.js";
-import {
-  AsyncIterableStream,
-  createAsyncIterableReadable,
-} from "../streams/asyncIterableStream.js";
+import { zodShapeStream } from "./stream.js";
 
 export type RunShape<TRunTypes extends AnyRunTypes> = TRunTypes extends AnyRunTypes
   ? {
@@ -157,7 +153,7 @@ export function runShapeStream<TRunTypes extends AnyRunTypes>(
 
 // First, define interfaces for the stream handling
 export interface StreamSubscription {
-  subscribe(): Promise<ReadableStream<unknown>>;
+  subscribe(): Promise<ReadableStream<SSEStreamPart<unknown>>>;
 }
 
 export type CreateStreamSubscriptionOptions = {
@@ -176,6 +172,12 @@ export interface StreamSubscriptionFactory {
   ): StreamSubscription;
 }
 
+export type SSEStreamPart<TChunk = unknown> = {
+  id: string;
+  chunk: TChunk;
+  timestamp: number;
+};
+
 // Real implementation for production
 export class SSEStreamSubscription implements StreamSubscription {
   private lastEventId: string | undefined;
@@ -197,7 +199,7 @@ export class SSEStreamSubscription implements StreamSubscription {
     this.lastEventId = options.lastEventId;
   }
 
-  async subscribe(): Promise<ReadableStream<unknown>> {
+  async subscribe(): Promise<ReadableStream<SSEStreamPart>> {
     const self = this;
 
     return new ReadableStream({
@@ -210,7 +212,9 @@ export class SSEStreamSubscription implements StreamSubscription {
     });
   }
 
-  private async connectStream(controller: ReadableStreamDefaultController): Promise<void> {
+  private async connectStream(
+    controller: ReadableStreamDefaultController<SSEStreamPart>
+  ): Promise<void> {
     try {
       const headers: Record<string, string> = {
         Accept: "text/event-stream",
@@ -259,14 +263,21 @@ export class SSEStreamSubscription implements StreamSubscription {
         .pipeThrough(new TextDecoderStream())
         .pipeThrough(new EventSourceParserStream())
         .pipeThrough(
-          new TransformStream({
+          new TransformStream<EventSourceMessage, SSEStreamPart>({
             transform: (chunk, chunkController) => {
               if (streamVersion === "v1") {
                 // Track the last event ID for resume support
                 if (chunk.id) {
                   this.lastEventId = chunk.id;
                 }
-                chunkController.enqueue(safeParseJSON(chunk.data));
+
+                const timestamp = parseRedisStreamIdTimestamp(chunk.id);
+
+                chunkController.enqueue({
+                  id: chunk.id ?? "unknown",
+                  chunk: safeParseJSON(chunk.data),
+                  timestamp,
+                });
               } else {
                 if (chunk.event === "batch") {
                   const data = safeParseJSON(chunk.data) as {
@@ -276,7 +287,11 @@ export class SSEStreamSubscription implements StreamSubscription {
                   for (const record of data.records) {
                     this.lastEventId = record.seq_num.toString();
 
-                    chunkController.enqueue(safeParseJSON(record.body));
+                    chunkController.enqueue({
+                      id: record.seq_num.toString(),
+                      chunk: safeParseJSON(record.body),
+                      timestamp: record.timestamp,
+                    });
                   }
                 }
               }
@@ -490,7 +505,7 @@ export class RunSubscription<TRunTypes extends AnyRunTypes> {
                         transform(chunk, controller) {
                           controller.enqueue({
                             type: streamKey,
-                            chunk: chunk as TStreams[typeof streamKey],
+                            chunk: chunk.chunk as TStreams[typeof streamKey],
                             run,
                           });
                         },
@@ -740,3 +755,17 @@ function getStreamsFromRunShape(run: AnyRunShape): string[] {
 
   return run.realtimeStreams;
 }
+
+// Redis stream IDs are in the format: <timestamp>-<sequence>
+function parseRedisStreamIdTimestamp(id?: string): number {
+  if (!id) {
+    return Date.now();
+  }
+
+  const timestamp = parseInt(id.split("-")[0] as string);
+  if (isNaN(timestamp)) {
+    return Date.now();
+  }
+
+  return timestamp;
+}

packages/trigger-sdk/src/v3/streams.ts

@@ -19,11 +19,16 @@ export type AppendStreamOptions = {
   requestOptions?: ApiRequestOptions;
 };
 
+export type AppendStreamResult<T> = {
+  stream: AsyncIterableStream<T>;
+  waitUntilComplete: () => Promise<void>;
+};
+
 async function append<T>(
   key: string,
   value: AsyncIterable<T> | ReadableStream<T>,
   options?: AppendStreamOptions
-): Promise<RealtimeStreamInstance<T>> {
+): Promise<AppendStreamResult<T>> {
   const runId = getRunIdForOptions(options);
 
   if (!runId) {
@@ -64,7 +69,10 @@ async function append<T>(
       span.end();
     });
 
-    return instance;
+    return {
+      stream: instance.stream,
+      waitUntilComplete: () => instance.wait(),
+    };
   } catch (error) {
     // if the error is a signal abort error, we need to end the span but not record an exception
     if (error instanceof Error && error.name === "AbortError") {

references/realtime-streams/src/trigger/streams.ts

@@ -113,7 +113,7 @@ export const streamsTask = task({
 
     const mockStream = createStreamFromGenerator(generator);
 
-    const { wait } = await streams.append("stream", mockStream);
+    const { waitUntilComplete } = await streams.append("stream", mockStream);
 
     await setTimeout(1000);
 
@@ -128,7 +128,7 @@ export const streamsTask = task({
       tokenCount++;
     }
 
-    await wait();
+    await waitUntilComplete();
 
     logger.info("Stream completed", { scenario, tokenCount });