Fix streams.append for s2, and a bunch of other little things, plus tests

ericallam · ericallam · commit 186846bbb779 · 2025-11-10T17:16:34.000Z
diff --git a/apps/webapp/app/routes/realtime.v1.streams.$runId.$target.$streamId.append.ts b/apps/webapp/app/routes/realtime.v1.streams.$runId.$target.$streamId.append.ts
@@ -1,5 +1,6 @@
 import { json } from "@remix-run/server-runtime";
 import { tryCatch } from "@trigger.dev/core/utils";
+import { nanoid } from "nanoid";
 import { z } from "zod";
 import { $replica, prisma } from "~/db.server";
 import { getRealtimeStreamInstance } from "~/services/realtime/v1StreamsGlobal.server";
@@ -96,8 +97,10 @@ const { action } = createActionApiRoute(
       targetRun.realtimeStreamsVersion
     );
 
+    const partId = request.headers.get("X-Part-Id") ?? nanoid(7);
+
     const [appendError] = await tryCatch(
-      realtimeStream.appendPart(part, targetId, params.streamId)
+      realtimeStream.appendPart(part, partId, targetId, params.streamId)
     );
 
     if (appendError) {
diff --git a/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.runs.$runParam.streams.$streamKey/route.tsx b/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.runs.$runParam.streams.$streamKey/route.tsx
@@ -118,8 +118,6 @@ export function RealtimeStreamViewer({
   const [mouseOver, setMouseOver] = useState(false);
   const [copied, setCopied] = useState(false);
 
-  console.log("chunks.length", chunks.length);
-
   const getCompactText = useCallback(() => {
     return chunks
       .map((chunk) => {
diff --git a/apps/webapp/app/services/realtime/redisRealtimeStreams.server.ts b/apps/webapp/app/services/realtime/redisRealtimeStreams.server.ts
@@ -334,7 +334,7 @@ export class RedisRealtimeStreams implements StreamIngestor, StreamResponder {
     }
   }
 
-  async appendPart(part: string, runId: string, streamId: string): Promise<void> {
+  async appendPart(part: string, partId: string, runId: string, streamId: string): Promise<void> {
     const redis = new Redis(this.options.redis ?? {});
     const streamKey = `stream:${runId}:${streamId}`;
 
diff --git a/apps/webapp/app/services/realtime/s2realtimeStreams.server.ts b/apps/webapp/app/services/realtime/s2realtimeStreams.server.ts
@@ -92,12 +92,16 @@ export class S2RealtimeStreams implements StreamResponder, StreamIngestor {
     throw new Error("S2 streams are written to S2 via the client, not from the server");
   }
 
-  async appendPart(part: string, runId: string, streamId: string): Promise<void> {
+  async appendPart(part: string, partId: string, runId: string, streamId: string): Promise<void> {
     const s2Stream = this.toStreamName(runId, streamId);
 
-    await this.s2Append(s2Stream, {
-      records: [{ body: part }],
+    this.logger.info(`S2 appending to stream`, { part, stream: s2Stream });
+
+    const result = await this.s2Append(s2Stream, {
+      records: [{ body: JSON.stringify({ data: part, id: partId }) }],
     });
+
+    this.logger.info(`S2 append result`, { result });
   }
 
   getLastChunkIndex(runId: string, streamId: string, clientId: string): Promise<number> {
@@ -116,6 +120,8 @@ export class S2RealtimeStreams implements StreamResponder, StreamIngestor {
     const s2Stream = this.toStreamName(runId, streamId);
     const startSeq = this.parseLastEventId(options?.lastEventId);
 
+    this.logger.info(`S2 streaming records from stream`, { stream: s2Stream, startSeq });
+
     // Request SSE stream from S2 and return it directly
     const s2Response = await this.s2StreamRecords(s2Stream, {
       seq_num: startSeq ?? 0,
diff --git a/apps/webapp/app/services/realtime/types.ts b/apps/webapp/app/services/realtime/types.ts
@@ -13,7 +13,7 @@ export interface StreamIngestor {
     resumeFromChunk?: number
   ): Promise<Response>;
 
-  appendPart(part: string, runId: string, streamId: string): Promise<void>;
+  appendPart(part: string, partId: string, runId: string, streamId: string): Promise<void>;
 
   getLastChunkIndex(runId: string, streamId: string, clientId: string): Promise<number>;
 }
diff --git a/packages/core/src/v3/apiClient/index.ts b/packages/core/src/v3/apiClient/index.ts
@@ -132,7 +132,7 @@ const DEFAULT_ZOD_FETCH_OPTIONS: ZodFetchOptions = {
 };
 
 export type ApiClientFutureFlags = {
-  unstable_v2RealtimeStreams?: boolean;
+  v2RealtimeStreams?: boolean;
 };
 
 export { isRequestOptions, SSEStreamSubscription };
@@ -1214,7 +1214,7 @@ export class ApiClient {
     headers[API_VERSION_HEADER_NAME] = API_VERSION;
 
     if (
-      this.futureFlags.unstable_v2RealtimeStreams ||
+      this.futureFlags.v2RealtimeStreams ||
       getEnvVar("TRIGGER_V2_REALTIME_STREAMS") === "1" ||
       getEnvVar("TRIGGER_V2_REALTIME_STREAMS") === "true" ||
       getEnvVar("TRIGGER_REALTIME_STREAMS_V2") === "1" ||
diff --git a/packages/core/src/v3/realtimeStreams/manager.ts b/packages/core/src/v3/realtimeStreams/manager.ts
@@ -87,7 +87,7 @@ export class StandardRealtimeStreamsManager implements RealtimeStreamsManager {
 
     const result = await this.apiClient.appendToStream(
       runId,
-      options?.target ?? "self",
+      "self",
       key,
       part,
       options?.requestOptions
@@ -141,11 +141,11 @@ export class StandardRealtimeStreamsManager implements RealtimeStreamsManager {
 function getRunIdForOptions(options?: RealtimeStreamOperationOptions): string | undefined {
   if (options?.target) {
     if (options.target === "parent") {
-      return taskContext.ctx?.run?.parentTaskRunId;
+      return taskContext.ctx?.run?.parentTaskRunId ?? taskContext.ctx?.run?.id;
     }
 
     if (options.target === "root") {
-      return taskContext.ctx?.run?.rootTaskRunId;
+      return taskContext.ctx?.run?.rootTaskRunId ?? taskContext.ctx?.run?.id;
     }
 
     if (options.target === "self") {
diff --git a/packages/core/src/v3/realtimeStreams/types.ts b/packages/core/src/v3/realtimeStreams/types.ts
@@ -39,7 +39,7 @@ export type RealtimeDefinedStream<TPart> = {
   ) => PipeStreamResult<TPart>;
   read: (runId: string, options?: ReadStreamOptions) => Promise<AsyncIterableStream<TPart>>;
   append: (value: TPart, options?: AppendStreamOptions) => Promise<void>;
-  writer: (options: WriterStreamOptions<TPart>) => void;
+  writer: (options: WriterStreamOptions<TPart>) => PipeStreamResult<TPart>;
 };
 
 export type InferStreamType<T> = T extends RealtimeDefinedStream<infer TPart> ? TPart : unknown;
diff --git a/packages/core/src/v3/types/tasks.ts b/packages/core/src/v3/types/tasks.ts
@@ -593,7 +593,8 @@ export interface Task<TIdentifier extends string, TInput = void, TOutput = any>
    */
   triggerAndWait: (
     payload: TInput,
-    options?: TriggerAndWaitOptions
+    options?: TriggerAndWaitOptions,
+    requestOptions?: TriggerApiRequestOptions
   ) => TaskRunPromise<TIdentifier, TOutput>;
 
   /**
diff --git a/packages/trigger-sdk/src/v3/shared.ts b/packages/trigger-sdk/src/v3/shared.ts
@@ -185,7 +185,7 @@ export function createTask<
         params.queue?.name
       );
     },
-    triggerAndWait: (payload, options) => {
+    triggerAndWait: (payload, options, requestOptions) => {
       return new TaskRunPromise<TIdentifier, TOutput>((resolve, reject) => {
         triggerAndWait_internal<TIdentifier, TInput, TOutput>(
           "triggerAndWait()",
@@ -195,7 +195,8 @@ export function createTask<
           {
             queue: params.queue?.name,
             ...options,
-          }
+          },
+          requestOptions
         )
           .then((result) => {
             resolve(result);
diff --git a/packages/trigger-sdk/src/v3/streams.ts b/packages/trigger-sdk/src/v3/streams.ts
@@ -358,28 +358,83 @@ function append<TPart extends BodyInit>(
   options?: AppendStreamOptions
 ): Promise<void> {
   if (typeof keyOrValue === "string" && typeof valueOrOptions === "string") {
-    return realtimeStreams.append(keyOrValue, valueOrOptions, options);
+    return appendInternal(keyOrValue, valueOrOptions, options);
   }
 
   if (typeof keyOrValue === "string") {
     if (isAppendStreamOptions(valueOrOptions)) {
-      return realtimeStreams.append(DEFAULT_STREAM_KEY, keyOrValue, valueOrOptions);
+      return appendInternal(DEFAULT_STREAM_KEY, keyOrValue, valueOrOptions);
     } else {
       if (!valueOrOptions) {
-        return realtimeStreams.append(DEFAULT_STREAM_KEY, keyOrValue, options);
+        return appendInternal(DEFAULT_STREAM_KEY, keyOrValue, options);
       }
 
-      return realtimeStreams.append(keyOrValue, valueOrOptions, options);
+      return appendInternal(keyOrValue, valueOrOptions, options);
     }
   } else {
     if (isAppendStreamOptions(valueOrOptions)) {
-      return realtimeStreams.append(DEFAULT_STREAM_KEY, keyOrValue, valueOrOptions);
+      return appendInternal(DEFAULT_STREAM_KEY, keyOrValue, valueOrOptions);
     } else {
-      return realtimeStreams.append(DEFAULT_STREAM_KEY, keyOrValue, options);
+      return appendInternal(DEFAULT_STREAM_KEY, keyOrValue, options);
     }
   }
 }
 
+async function appendInternal<TPart extends BodyInit>(
+  key: string,
+  part: TPart,
+  options?: AppendStreamOptions
+): Promise<void> {
+  const runId = getRunIdForOptions(options);
+
+  if (!runId) {
+    throw new Error(
+      "Could not determine the target run ID for the realtime stream. Please specify a target run ID using the `target` option or use this function from inside a task."
+    );
+  }
+
+  const span = tracer.startSpan("streams.append()", {
+    attributes: {
+      key,
+      runId,
+      [SemanticInternalAttributes.ENTITY_TYPE]: "realtime-stream",
+      [SemanticInternalAttributes.ENTITY_ID]: `${runId}:${key}`,
+      [SemanticInternalAttributes.STYLE_ICON]: "streams",
+      ...accessoryAttributes({
+        items: [
+          {
+            text: key,
+            variant: "normal",
+          },
+        ],
+        style: "codepath",
+      }),
+    },
+  });
+
+  try {
+    await realtimeStreams.append(key, part, options);
+    span.end();
+  } 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") {
+      span.end();
+      throw error;
+    }
+
+    if (error instanceof Error || typeof error === "string") {
+      span.recordException(error);
+    } else {
+      span.recordException(String(error));
+    }
+
+    span.setStatus({ code: SpanStatusCode.ERROR });
+    span.end();
+
+    throw error;
+  }
+}
+
 function isAppendStreamOptions(val: unknown): val is AppendStreamOptions {
   return (
     typeof val === "object" &&
@@ -390,7 +445,121 @@ function isAppendStreamOptions(val: unknown): val is AppendStreamOptions {
   );
 }
 
-function writer<TPart>(key: string, options: WriterStreamOptions<TPart>) {
+/**
+ * Writes data to a realtime stream using the default stream key (`"default"`).
+ *
+ * This is a convenience overload that allows you to write to the default stream without
+ * specifying a stream key. The stream will be created/accessed with the key `"default"`.
+ *
+ * @template TPart - The type of data chunks in the stream
+ * @param options - The options for writing to the stream
+ * @returns A promise that resolves to an object containing:
+ *   - `stream`: The original stream (can be consumed in your task)
+ *   - `waitUntilComplete`: A function that returns a promise resolving when the stream is fully sent
+ *
+ * @example
+ * ```ts
+ * import { streams } from "@trigger.dev/sdk";
+ *
+ * // Write to the default stream
+ * const { waitUntilComplete } = await streams.writer({
+ *   execute: ({ write, merge }) => {
+ *     write("chunk 1");
+ *     write("chunk 2");
+ *     write("chunk 3");
+ *   },
+ * });
+ *
+ * // Wait for all chunks to be written
+ * await waitUntilComplete();
+ * ```
+ *
+ * @example
+ * ```ts
+ * // Write to a specific stream key
+ * const { waitUntilComplete } = await streams.writer("my-custom-stream", {
+ *   execute: ({ write, merge }) => {
+ *     write("chunk 1");
+ *     write("chunk 2");
+ *     write("chunk 3");
+ *   },
+ * });
+ *
+ * // Wait for all chunks to be written
+ * await waitUntilComplete();
+ * ```
+ *
+ * @example
+ * ```ts
+ * // Write to a parent run
+ * await streams.writer("output", {
+ *   execute: ({ write, merge }) => {
+ *     write("chunk 1");
+ *     write("chunk 2");
+ *     write("chunk 3");
+ *   },
+ * });
+ *
+ * // Wait for all chunks to be written
+ * await waitUntilComplete();
+ * ```
+ *
+ * @example
+ * ```ts
+ * // Write to a specific stream key
+ * await streams.writer("my-custom-stream", {
+ *   execute: ({ write, merge }) => {
+ *     write("chunk 1");
+ *     write("chunk 2");
+ *     write("chunk 3");
+ *   },
+ * });
+ *
+ * // Wait for all chunks to be written
+ * await waitUntilComplete();
+ * ```
+ */
+function writer<TPart>(options: WriterStreamOptions<TPart>): PipeStreamResult<TPart>;
+/**
+ * Writes data to a realtime stream with a specific stream key.
+ *
+ * @template TPart - The type of data chunks in the stream
+ * @param key - The unique identifier of the stream to write to. Defaults to `"default"` if not provided.
+ * @param options - The options for writing to the stream
+ * @returns A promise that resolves to an object containing:
+ *   - `stream`: The original stream (can be consumed in your task)
+ *   - `waitUntilComplete`: A function that returns a promise resolving when the stream is fully sent
+ *
+ * @example
+ * ```ts
+ * import { streams } from "@trigger.dev/sdk";
+ *
+ * // Write to a specific stream key
+ * const { waitUntilComplete } = await streams.writer("my-custom-stream", {
+ *   execute: ({ write, merge }) => {
+ *     write("chunk 1");
+ *     write("chunk 2");
+ *     write("chunk 3");
+ *   },
+ * });
+ *
+ * // Wait for all chunks to be written
+ * await waitUntilComplete();
+ * ```
+ */
+function writer<TPart>(key: string, options: WriterStreamOptions<TPart>): PipeStreamResult<TPart>;
+function writer<TPart>(
+  keyOrOptions: string | WriterStreamOptions<TPart>,
+  valueOrOptions?: WriterStreamOptions<TPart>
+): PipeStreamResult<TPart> {
+  if (typeof keyOrOptions === "string") {
+    return writerInternal(keyOrOptions, valueOrOptions!);
+  }
+
+  return writerInternal(DEFAULT_STREAM_KEY, keyOrOptions);
+}
+
+function writerInternal<TPart>(key: string, options: WriterStreamOptions<TPart>) {
   let controller!: ReadableStreamDefaultController<TPart>;
 
   const ongoingStreamPromises: Promise<void>[] = [];
diff --git a/references/realtime-streams/src/app/actions.ts b/references/realtime-streams/src/app/actions.ts
@@ -14,7 +14,7 @@ export async function triggerStreamTask(
   const config = useDurableStreams
     ? {
         future: {
-          unstable_v2RealtimeStreams: true,
+          v2RealtimeStreams: true,
         },
       }
     : undefined;
@@ -52,7 +52,7 @@ export async function triggerAIChatTask(messages: UIMessage[]) {
     {
       clientConfig: {
         future: {
-          unstable_v2RealtimeStreams: true,
+          v2RealtimeStreams: true,
         },
       },
     }
diff --git a/references/realtime-streams/src/trigger/streams.ts b/references/realtime-streams/src/trigger/streams.ts

Original file line number	Diff line number	Diff line change
`@@ -334,7 +334,7 @@ export class RedisRealtimeStreams implements StreamIngestor, StreamResponder {`
`334`	`334`	`}`
`335`	`335`	`}`
`336`	`336`
`337`		`- async appendPart(part: string, runId: string, streamId: string): Promise<void> {`
	`337`	`+ async appendPart(part: string, partId: string, runId: string, streamId: string): Promise<void> {`
`338`	`338`	`const redis = new Redis(this.options.redis ?? {});`
`339`	`339`	const streamKey = `stream:${runId}:${streamId}`;
`340`	`340`
Original file line number	Diff line number	Diff line change
`@@ -13,7 +13,7 @@ export interface StreamIngestor {`
`13`	`13`	`resumeFromChunk?: number`
`14`	`14`	`): Promise<Response>;`
`15`	`15`
`16`		`- appendPart(part: string, runId: string, streamId: string): Promise<void>;`
	`16`	`+ appendPart(part: string, partId: string, runId: string, streamId: string): Promise<void>;`
`17`	`17`
`18`	`18`	`getLastChunkIndex(runId: string, streamId: string, clientId: string): Promise<number>;`
`19`	`19`	`}`