Write to s2 from the client instead of the server

Author: ericallam

apps/webapp/app/env.server.ts

@@ -1211,7 +1211,7 @@ const EnvironmentSchema = z
       .enum(["log", "error", "warn", "info", "debug"])
       .default("info"),
     REALTIME_STREAMS_S2_FLUSH_INTERVAL_MS: z.coerce.number().int().default(100),
-    REALTIME_STREAMS_S2_RESUME_TTL_SECONDS: z.coerce.number().int().default(86400),
+    REALTIME_STREAMS_S2_MAX_RETRIES: z.coerce.number().int().default(10),
   })
   .and(GithubAppEnvSchema);
 

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

@@ -1,5 +1,6 @@
+import { json } from "@remix-run/server-runtime";
 import { z } from "zod";
-import { $replica } from "~/db.server";
+import { $replica, prisma } from "~/db.server";
 import { getRealtimeStreamInstance } from "~/services/realtime/v1StreamsGlobal.server";
 import {
   createActionApiRoute,
@@ -53,26 +54,58 @@ const { action } = createActionApiRoute(
       return new Response("Target not found", { status: 404 });
     }
 
-    // Extract client ID from header, default to "default" if not provided
-    const clientId = request.headers.get("X-Client-Id") || "default";
-    const streamVersion = request.headers.get("X-Stream-Version") || "v1";
-
-    if (!request.body) {
-      return new Response("No body provided", { status: 400 });
-    }
+    if (request.method === "PUT") {
+      // This is the "create" endpoint
+      const updatedRun = await prisma.taskRun.update({
+        where: {
+          friendlyId: targetId,
+          runtimeEnvironmentId: authentication.environment.id,
+        },
+        data: {
+          realtimeStreams: {
+            push: params.streamId,
+          },
+        },
+        select: {
+          realtimeStreamsVersion: true,
+        },
+      });
 
-    const resumeFromChunk = request.headers.get("X-Resume-From-Chunk");
-    const resumeFromChunkNumber = resumeFromChunk ? parseInt(resumeFromChunk, 10) : undefined;
+      const realtimeStream = getRealtimeStreamInstance(
+        authentication.environment,
+        updatedRun.realtimeStreamsVersion
+      );
 
-    const realtimeStream = getRealtimeStreamInstance(authentication.environment, streamVersion);
+      const { responseHeaders } = await realtimeStream.initializeStream(targetId, params.streamId);
 
-    return realtimeStream.ingestData(
-      request.body,
-      targetId,
-      params.streamId,
-      clientId,
-      resumeFromChunkNumber
-    );
+      return json(
+        {
+          version: updatedRun.realtimeStreamsVersion,
+        },
+        { status: 202, headers: responseHeaders }
+      );
+    } else {
+      // Extract client ID from header, default to "default" if not provided
+      const clientId = request.headers.get("X-Client-Id") || "default";
+      const streamVersion = request.headers.get("X-Stream-Version") || "v1";
+
+      if (!request.body) {
+        return new Response("No body provided", { status: 400 });
+      }
+
+      const resumeFromChunk = request.headers.get("X-Resume-From-Chunk");
+      const resumeFromChunkNumber = resumeFromChunk ? parseInt(resumeFromChunk, 10) : undefined;
+
+      const realtimeStream = getRealtimeStreamInstance(authentication.environment, streamVersion);
+
+      return realtimeStream.ingestData(
+        request.body,
+        targetId,
+        params.streamId,
+        clientId,
+        resumeFromChunkNumber
+      );
+    }
   }
 );
 

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

@@ -29,6 +29,13 @@ export class RedisRealtimeStreams implements StreamIngestor, StreamResponder {
     this.inactivityTimeoutMs = options.inactivityTimeoutMs ?? 60000; // Default: 60 seconds
   }
 
+  async initializeStream(
+    runId: string,
+    streamId: string
+  ): Promise<{ responseHeaders?: Record<string, string> }> {
+    return {};
+  }
+
   async streamResponse(
     request: Request,
     runId: string,