Add ability to stream into parent and root task runs

Author: ericallam

apps/webapp/app/routes/realtime.v1.streams.$runId.$streamId.ts

@@ -2,7 +2,6 @@ import { ActionFunctionArgs } from "@remix-run/server-runtime";
 import { z } from "zod";
 import { $replica } from "~/db.server";
 import { relayRealtimeStreams } from "~/services/realtime/relayRealtimeStreams.server";
-import { v1RealtimeStreams } from "~/services/realtime/v1StreamsGlobal.server";
 import { createLoaderApiRoute } from "~/services/routeBuilders/apiBuilder.server";
 
 const ParamsSchema = z.object({

apps/webapp/app/routes/realtime.v1.streams.$runId.$target.$streamId.ts

@@ -0,0 +1,61 @@
+import { z } from "zod";
+import { $replica } from "~/db.server";
+import { relayRealtimeStreams } from "~/services/realtime/relayRealtimeStreams.server";
+import { createActionApiRoute } from "~/services/routeBuilders/apiBuilder.server";
+
+const ParamsSchema = z.object({
+  runId: z.string(),
+  target: z.enum(["self", "parent", "root"]),
+  streamId: z.string(),
+});
+
+const { action } = createActionApiRoute(
+  {
+    params: ParamsSchema,
+  },
+  async ({ request, params, authentication }) => {
+    if (!request.body) {
+      return new Response("No body provided", { status: 400 });
+    }
+
+    const run = await $replica.taskRun.findFirst({
+      where: {
+        friendlyId: params.runId,
+        runtimeEnvironmentId: authentication.environment.id,
+      },
+      select: {
+        id: true,
+        friendlyId: true,
+        parentTaskRun: {
+          select: {
+            friendlyId: true,
+          },
+        },
+        rootTaskRun: {
+          select: {
+            friendlyId: true,
+          },
+        },
+      },
+    });
+
+    if (!run) {
+      return new Response("Run not found", { status: 404 });
+    }
+
+    const targetId =
+      params.target === "self"
+        ? run.friendlyId
+        : params.target === "parent"
+        ? run.parentTaskRun?.friendlyId
+        : run.rootTaskRun?.friendlyId;
+
+    if (!targetId) {
+      return new Response("Target not found", { status: 404 });
+    }
+
+    return relayRealtimeStreams.ingestData(request.body, targetId, params.streamId);
+  }
+);
+
+export { action };

packages/core/src/v3/runMetadata/manager.ts

@@ -60,7 +60,7 @@ export class StandardMetadataManager implements RunMetadataManager {
         this.queuedParentOperations.add({ type: "update", value });
         return this.parent;
       },
-      stream: (key, value, signal) => this.doStream(key, value, this.parent, signal),
+      stream: (key, value, signal) => this.doStream(key, value, "parent", this.parent, signal),
     };
   }
 
@@ -94,7 +94,7 @@ export class StandardMetadataManager implements RunMetadataManager {
         this.queuedParentOperations.add({ type: "update", value });
         return this.root;
       },
-      stream: (key, value, signal) => this.doStream(key, value, this.root, signal),
+      stream: (key, value, signal) => this.doStream(key, value, "root", this.root, signal),
     };
   }
 
@@ -194,12 +194,13 @@ export class StandardMetadataManager implements RunMetadataManager {
     value: AsyncIterable<T> | ReadableStream<T>,
     signal?: AbortSignal
   ): Promise<AsyncIterable<T>> {
-    return this.doStream(key, value, this, signal);
+    return this.doStream(key, value, "self", this, signal);
   }
 
   private async doStream<T>(
     key: string,
     value: AsyncIterable<T> | ReadableStream<T>,
+    target: "self" | "parent" | "root",
     updater: RunMetadataUpdater = this,
     signal?: AbortSignal
   ): Promise<AsyncIterable<T>> {
@@ -236,6 +237,7 @@ export class StandardMetadataManager implements RunMetadataManager {
         headers: this.apiClient.getHeaders(),
         signal,
         version: this.streamsVersion,
+        target,
       });
 
       this.activeStreams.set(key, streamInstance);

packages/core/src/v3/runMetadata/metadataStream.ts

@@ -6,6 +6,7 @@ export type MetadataOptions<T> = {
   headers?: Record<string, string>;
   signal?: AbortSignal;
   version?: "v1" | "v2";
+  target?: "self" | "parent" | "root";
 };
 
 export class MetadataStream<T> {
@@ -45,19 +46,14 @@ export class MetadataStream<T> {
       })
     );
 
-    return fetch(
-      `${this.options.baseUrl}/realtime/${this.options.version ?? "v1"}/streams/${
-        this.options.runId
-      }/${this.options.key}`,
-      {
-        method: "POST",
-        headers: this.options.headers ?? {},
-        body: serverStream,
-        // @ts-expect-error
-        duplex: "half",
-        signal: this.controller.signal,
-      }
-    );
+    return fetch(this.buildUrl(), {
+      method: "POST",
+      headers: this.options.headers ?? {},
+      body: serverStream,
+      signal: this.controller.signal,
+      // @ts-expect-error
+      duplex: "half",
+    });
   }
 
   public async wait(): Promise<void> {
@@ -67,6 +63,19 @@ export class MetadataStream<T> {
   public [Symbol.asyncIterator]() {
     return streamToAsyncIterator(this.consumerStream);
   }
+
+  private buildUrl(): string {
+    switch (this.options.version ?? "v1") {
+      case "v1": {
+        return `${this.options.baseUrl}/realtime/v1/streams/${this.options.runId}/${
+          this.options.target ?? "self"
+        }/${this.options.key}`;
+      }
+      case "v2": {
+        return `${this.options.baseUrl}/realtime/v2/streams/${this.options.runId}/${this.options.key}`;
+      }
+    }
+  }
 }
 
 async function* streamToAsyncIterator<T>(stream: ReadableStream<T>): AsyncIterableIterator<T> {

references/nextjs-realtime/src/trigger/csv.ts

@@ -83,6 +83,14 @@ export const handleCSVRow = schemaTask({
 
     metadata.parent.increment("processedRows", 1).append("rowRuns", ctx.run.id);
 
+    await metadata.parent.stream(
+      ctx.run.id,
+      (async function* () {
+        yield "hello";
+        yield "world";
+      })()
+    );
+
     return row;
   },
 });

references/v3-catalog/src/trigger/runMetadata.ts

@@ -21,8 +21,8 @@ export const runMetadataTask = task({
 export const runMetadataChildTask = task({
   id: "run-metadata-child-task",
   run: async (payload: any, { ctx }) => {
-    metadata.parent.incrementKey("numberOfChildren", 1);
-    metadata.root.incrementKey("numberOfChildren", 1);
+    metadata.parent.increment("numberOfChildren", 1);
+    metadata.root.increment("numberOfChildren", 1);
 
     logger.info("metadata", { metadata: metadata.current() });
 
@@ -63,7 +63,7 @@ export const runMetadataChildTask = task({
 export const runMetadataChildTask2 = task({
   id: "run-metadata-child-task-2",
   run: async (payload: any, { ctx }) => {
-    metadata.root.incrementKey("numberOfChildren", 1);
+    metadata.root.increment("numberOfChildren", 1);
   },
 });