Realtime streams now powered by electric (#1541)

* Realtime streams now powered by electric, and fix the streaming update duplicate issues by converting the electric Shape materialized view into a ReadableStream of changes

* Ensure realtime subscription stops when runs are finished, and add an onComplete handle to use realtime hooks

* Fix tests

Author: ericallam

.changeset/lemon-cherries-greet.md

@@ -0,0 +1,6 @@
+---
+"@trigger.dev/react-hooks": patch
+"@trigger.dev/sdk": patch
+---
+
+Realtime streams now powered by electric. Also, this change fixes a realtime bug that was causing too many re-renders, even on records that didn't change

apps/webapp/app/env.server.ts

@@ -243,6 +243,8 @@ const EnvironmentSchema = z.object({
   MAXIMUM_DEV_QUEUE_SIZE: z.coerce.number().int().optional(),
   MAXIMUM_DEPLOYED_QUEUE_SIZE: z.coerce.number().int().optional(),
   MAX_BATCH_V2_TRIGGER_ITEMS: z.coerce.number().int().default(500),
+
+  REALTIME_STREAM_VERSION: z.enum(["v1", "v2"]).default("v1"),
 });
 
 export type Environment = z.infer<typeof EnvironmentSchema>;

apps/webapp/app/presenters/v3/SpanPresenter.server.ts

@@ -215,7 +215,9 @@ export class SpanPresenter extends BasePresenter {
     const span = await eventRepository.getSpan(spanId, run.traceId);
 
     const metadata = run.metadata
-      ? await prettyPrintPacket(run.metadata, run.metadataType, { filteredKeys: ["$$streams"] })
+      ? await prettyPrintPacket(run.metadata, run.metadataType, {
+          filteredKeys: ["$$streams", "$$streamsVersion"],
+        })
       : undefined;
 
     const context = {

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

@@ -1,7 +1,7 @@
 import { ActionFunctionArgs } from "@remix-run/server-runtime";
 import { z } from "zod";
 import { $replica } from "~/db.server";
-import { realtimeStreams } from "~/services/realtimeStreamsGlobal.server";
+import { v1RealtimeStreams } from "~/services/realtime/v1StreamsGlobal.server";
 import { createLoaderApiRoute } from "~/services/routeBuilders/apiBuilder.server";
 
 const ParamsSchema = z.object({
@@ -16,7 +16,7 @@ export async function action({ request, params }: ActionFunctionArgs) {
     return new Response("No body provided", { status: 400 });
   }
 
-  return realtimeStreams.ingestData(request.body, $params.runId, $params.streamId);
+  return v1RealtimeStreams.ingestData(request.body, $params.runId, $params.streamId);
 }
 
 export const loader = createLoaderApiRoute(
@@ -50,7 +50,13 @@ export const loader = createLoaderApiRoute(
       superScopes: ["read:runs", "read:all", "admin"],
     },
   },
-  async ({ params, request, resource: run }) => {
-    return realtimeStreams.streamResponse(run.friendlyId, params.streamId, request.signal);
+  async ({ params, request, resource: run, authentication }) => {
+    return v1RealtimeStreams.streamResponse(
+      request,
+      run.friendlyId,
+      params.streamId,
+      authentication.environment,
+      request.signal
+    );
   }
 );

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

@@ -0,0 +1,87 @@
+import { z } from "zod";
+import { $replica } from "~/db.server";
+import {
+  createActionApiRoute,
+  createLoaderApiRoute,
+} from "~/services/routeBuilders/apiBuilder.server";
+import { v2RealtimeStreams } from "~/services/realtime/v2StreamsGlobal.server";
+
+const ParamsSchema = z.object({
+  runId: z.string(),
+  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,
+      },
+      include: {
+        batch: {
+          select: {
+            friendlyId: true,
+          },
+        },
+      },
+    });
+
+    if (!run) {
+      return new Response("Run not found", { status: 404 });
+    }
+
+    return v2RealtimeStreams.ingestData(request.body, run.id, params.streamId);
+  }
+);
+
+export { action };
+
+export const loader = createLoaderApiRoute(
+  {
+    params: ParamsSchema,
+    allowJWT: true,
+    corsStrategy: "all",
+    findResource: async (params, auth) => {
+      return $replica.taskRun.findFirst({
+        where: {
+          friendlyId: params.runId,
+          runtimeEnvironmentId: auth.environment.id,
+        },
+        include: {
+          batch: {
+            select: {
+              friendlyId: true,
+            },
+          },
+        },
+      });
+    },
+    authorization: {
+      action: "read",
+      resource: (run) => ({
+        runs: run.friendlyId,
+        tags: run.runTags,
+        batch: run.batch?.friendlyId,
+        tasks: run.taskIdentifier,
+      }),
+      superScopes: ["read:runs", "read:all", "admin"],
+    },
+  },
+  async ({ params, request, resource: run, authentication }) => {
+    return v2RealtimeStreams.streamResponse(
+      request,
+      run.id,
+      params.streamId,
+      authentication.environment,
+      request.signal
+    );
+  }
+);

apps/webapp/app/services/realtime/databaseRealtimeStreams.server.ts

@@ -0,0 +1,85 @@
+import { PrismaClient } from "@trigger.dev/database";
+import { AuthenticatedEnvironment } from "../apiAuth.server";
+import { logger } from "../logger.server";
+import { RealtimeClient } from "../realtimeClient.server";
+import { StreamIngestor, StreamResponder } from "./types";
+
+export type DatabaseRealtimeStreamsOptions = {
+  prisma: PrismaClient;
+  realtimeClient: RealtimeClient;
+};
+
+// Class implementing both interfaces
+export class DatabaseRealtimeStreams implements StreamIngestor, StreamResponder {
+  constructor(private options: DatabaseRealtimeStreamsOptions) {}
+
+  async streamResponse(
+    request: Request,
+    runId: string,
+    streamId: string,
+    environment: AuthenticatedEnvironment,
+    signal: AbortSignal
+  ): Promise<Response> {
+    return this.options.realtimeClient.streamChunks(
+      request.url,
+      environment,
+      runId,
+      streamId,
+      signal,
+      request.headers.get("x-trigger-electric-version") ?? undefined
+    );
+  }
+
+  async ingestData(
+    stream: ReadableStream<Uint8Array>,
+    runId: string,
+    streamId: string
+  ): Promise<Response> {
+    try {
+      const textStream = stream.pipeThrough(new TextDecoderStream());
+      const reader = textStream.getReader();
+      let sequence = 0;
+
+      while (true) {
+        const { done, value } = await reader.read();
+
+        if (done) {
+          break;
+        }
+
+        logger.debug("[DatabaseRealtimeStreams][ingestData] Reading data", {
+          streamId,
+          runId,
+          value,
+        });
+
+        const chunks = value
+          .split("\n")
+          .filter((chunk) => chunk) // Remove empty lines
+          .map((line) => {
+            return {
+              sequence: sequence++,
+              value: line,
+            };
+          });
+
+        await this.options.prisma.realtimeStreamChunk.createMany({
+          data: chunks.map((chunk) => {
+            return {
+              runId,
+              key: streamId,
+              sequence: chunk.sequence,
+              value: chunk.value,
+            };
+          }),
+        });
+      }
+
+      return new Response(null, { status: 200 });
+    } catch (error) {
+      logger.error("[DatabaseRealtimeStreams][ingestData] Error in ingestData:", { error });
+
+      return new Response(null, { status: 500 });
+    }
+  }
+}

apps/webapp/app/services/realtimeStreams.server.ts renamed to apps/webapp/app/services/realtime/redisRealtimeStreams.server.ts

@@ -1,16 +1,25 @@
 import Redis, { RedisKey, RedisOptions, RedisValue } from "ioredis";
-import { logger } from "./logger.server";
+import { logger } from "../logger.server";
+import { StreamIngestor, StreamResponder } from "./types";
+import { AuthenticatedEnvironment } from "../apiAuth.server";
 
 export type RealtimeStreamsOptions = {
   redis: RedisOptions | undefined;
 };
 
 const END_SENTINEL = "<<CLOSE_STREAM>>";
 
-export class RealtimeStreams {
+// Class implementing both interfaces
+export class RedisRealtimeStreams implements StreamIngestor, StreamResponder {
   constructor(private options: RealtimeStreamsOptions) {}
 
-  async streamResponse(runId: string, streamId: string, signal: AbortSignal): Promise<Response> {
+  async streamResponse(
+    request: Request,
+    runId: string,
+    streamId: string,
+    environment: AuthenticatedEnvironment,
+    signal: AbortSignal
+  ): Promise<Response> {
     const redis = new Redis(this.options.redis ?? {});
     const streamKey = `stream:${runId}:${streamId}`;
     let isCleanedUp = false;
@@ -115,11 +124,10 @@ export class RealtimeStreams {
     }
 
     try {
-      // Use TextDecoderStream to simplify text decoding
       const textStream = stream.pipeThrough(new TextDecoderStream());
       const reader = textStream.getReader();
 
-      const batchSize = 10; // Adjust this value based on performance testing
+      const batchSize = 10;
       let batchCommands: Array<[key: RedisKey, ...args: RedisValue[]]> = [];
 
       while (true) {
@@ -131,17 +139,13 @@ export class RealtimeStreams {
 
         logger.debug("[RealtimeStreams][ingestData] Reading data", { streamKey, value });
 
-        // 'value' is a string containing the decoded text
         const lines = value.split("\n");
 
         for (const line of lines) {
           if (line.trim()) {
-            // Avoid unnecessary parsing; assume 'line' is already a JSON string
-            // Add XADD command with MAXLEN option to limit stream size
             batchCommands.push([streamKey, "MAXLEN", "~", "2500", "*", "data", line]);
 
             if (batchCommands.length >= batchSize) {
-              // Send batch using a pipeline
               const pipeline = redis.pipeline();
               for (const args of batchCommands) {
                 pipeline.xadd(...args);
@@ -153,7 +157,6 @@ export class RealtimeStreams {
         }
       }
 
-      // Send any remaining commands
       if (batchCommands.length > 0) {
         const pipeline = redis.pipeline();
         for (const args of batchCommands) {
@@ -162,7 +165,6 @@ export class RealtimeStreams {
         await pipeline.exec();
       }
 
-      // Send the __end message to indicate the end of the stream
       await redis.xadd(streamKey, "MAXLEN", "~", "1000", "*", "data", END_SENTINEL);
 
       return new Response(null, { status: 200 });

apps/webapp/app/services/realtime/types.ts

@@ -0,0 +1,21 @@
+import { AuthenticatedEnvironment } from "../apiAuth.server";
+
+// Interface for stream ingestion
+export interface StreamIngestor {
+  ingestData(
+    stream: ReadableStream<Uint8Array>,
+    runId: string,
+    streamId: string
+  ): Promise<Response>;
+}
+
+// Interface for stream response
+export interface StreamResponder {
+  streamResponse(
+    request: Request,
+    runId: string,
+    streamId: string,
+    environment: AuthenticatedEnvironment,
+    signal: AbortSignal
+  ): Promise<Response>;
+}

apps/webapp/app/services/realtimeStreamsGlobal.server.ts renamed to apps/webapp/app/services/realtime/v1StreamsGlobal.server.ts

@@ -1,9 +1,9 @@
 import { env } from "~/env.server";
 import { singleton } from "~/utils/singleton";
-import { RealtimeStreams } from "./realtimeStreams.server";
+import { RedisRealtimeStreams } from "./redisRealtimeStreams.server";
 
-function initializeRealtimeStreams() {
-  return new RealtimeStreams({
+function initializeRedisRealtimeStreams() {
+  return new RedisRealtimeStreams({
     redis: {
       port: env.REDIS_PORT,
       host: env.REDIS_HOST,
@@ -16,4 +16,4 @@ function initializeRealtimeStreams() {
   });
 }
 
-export const realtimeStreams = singleton("realtimeStreams", initializeRealtimeStreams);
+export const v1RealtimeStreams = singleton("realtimeStreams", initializeRedisRealtimeStreams);

apps/webapp/app/services/realtime/v2StreamsGlobal.server.ts

@@ -0,0 +1,13 @@
+import { prisma } from "~/db.server";
+import { singleton } from "~/utils/singleton";
+import { realtimeClient } from "../realtimeClientGlobal.server";
+import { DatabaseRealtimeStreams } from "./databaseRealtimeStreams.server";
+
+function initializeDatabaseRealtimeStreams() {
+  return new DatabaseRealtimeStreams({
+    prisma,
+    realtimeClient,
+  });
+}
+
+export const v2RealtimeStreams = singleton("dbRealtimeStreams", initializeDatabaseRealtimeStreams);