feat(component): write new streams to core

src/client/index.ts

@@ -203,8 +203,8 @@ export class PersistentTextStreaming {
   }
 
   /**
-   * Delete a stream and all its chunks. Deletion happens asynchronously
-   * in batches, so this returns immediately.
+   * Delete a stream and its persisted events. Large backing stores may finish
+   * cleanup in background batches.
    *
    * @param ctx - A convex context capable of running mutations.
    * @param streamId - The ID of the stream to delete.
@@ -236,7 +236,7 @@ export class PersistentTextStreaming {
   private async setStreamStatus(
     ctx: MutationCtx | ActionCtx,
     streamId: StreamId,
-    status: StreamStatus,
+    status: "done" | "error" | "timeout",
   ) {
     await ctx.runMutation(this.component.lib.setStreamStatus, {
       streamId,

src/component/_generated/component.ts

@@ -87,7 +87,7 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
         "mutation",
         "internal",
         {
-          status: "pending" | "streaming" | "done" | "error" | "timeout";
+          status: "done" | "error" | "timeout";
           streamId: string;
         },
         any,

src/component/lib.test.ts

@@ -2,9 +2,9 @@
 
 import { validateStreamReadResult } from "@convex-dev/stream/testing";
 import { convexTest } from "convex-test";
-import { describe, expect, it } from "vitest";
+import { describe, expect, it, vi } from "vitest";
 
-import { api } from "./_generated/api.js";
+import { api, internal } from "./_generated/api.js";
 import schema from "./schema.js";
 import { textStreams } from "./streams.js";
 
@@ -174,7 +174,6 @@ describe("bounded chunk reads", () => {
       { text: events.join(""), status: "done" },
     );
   });
-
   it("requires paged reads for a live core-backed facade", async () => {
     const t = convexTest(schema, modules);
     const streamId = await t.run(async (ctx) => {
@@ -210,4 +209,133 @@ describe("bounded chunk reads", () => {
       "fullBodyLimitExceeded",
     );
   });
+
+  it("writes new streams to core storage only", async () => {
+    const t = convexTest(schema, modules);
+    const streamId = await t.mutation(api.lib.createStream, {});
+    await t.mutation(api.lib.addChunk, {
+      streamId,
+      text: "one",
+      final: false,
+    });
+    await t.mutation(api.lib.addChunk, {
+      streamId,
+      text: "two",
+      final: true,
+    });
+
+    const state = await t.run(async (ctx) => {
+      const facade = await ctx.db.get("streams", streamId);
+      const core = facade?.coreId
+        ? await ctx.db.get("textStreams", facade.coreId)
+        : null;
+      return {
+        core,
+        chunks: await ctx.db.query("chunks").collect(),
+        events: await ctx.db.query("textStreamsEvents").collect(),
+      };
+    });
+    expect(state).toMatchObject({
+      core: { status: "done", nextSeq: 2 },
+      chunks: [],
+      events: [
+        { attempt: 0, seq: 0, event: "one" },
+        { attempt: 0, seq: 1, event: "two" },
+      ],
+    });
+    await expect(t.query(api.lib.getStreamText, { streamId })).resolves.toEqual(
+      {
+        text: "onetwo",
+        status: "done",
+      },
+    );
+  });
+
+  it("mirrors errors and times out inactive streams", async () => {
+    vi.useFakeTimers();
+    try {
+      vi.setSystemTime(0);
+      const t = convexTest(schema, modules);
+      const failedId = await t.mutation(api.lib.createStream, {});
+      await t.mutation(api.lib.setStreamStatus, {
+        streamId: failedId,
+        status: "error",
+      });
+      const timedOutId = await t.mutation(api.lib.createStream, {});
+      const activeId = await t.mutation(api.lib.createStream, {});
+      const legacyId = await t.run((ctx) =>
+        ctx.db.insert("streams", { status: "streaming" }),
+      );
+      vi.setSystemTime(19 * 60 * 1000);
+      await t.mutation(api.lib.addChunk, {
+        streamId: activeId,
+        text: "active",
+        final: false,
+      });
+      await t.run((ctx) =>
+        ctx.db.insert("chunks", { streamId: legacyId, text: "active" }),
+      );
+      vi.setSystemTime(20 * 60 * 1000 + 1);
+      await t.mutation(internal.lib.cleanupExpiredStreams, {});
+
+      const firstStates = await t.run(async (ctx) =>
+        Promise.all(
+          [failedId, timedOutId, activeId, legacyId].map(async (streamId) => {
+            const facade = await ctx.db.get("streams", streamId);
+            return {
+              facade,
+              core: facade?.coreId
+                ? await ctx.db.get("textStreams", facade.coreId)
+                : null,
+            };
+          }),
+        ),
+      );
+      expect(firstStates).toMatchObject([
+        {
+          facade: { status: "error" },
+          core: { status: "failed", error: { code: "error" } },
+        },
+        {
+          facade: { status: "timeout" },
+          core: { status: "failed", error: { code: "timeout" } },
+        },
+        { facade: { status: "streaming" }, core: { status: "streaming" } },
+        { facade: { status: "streaming" }, core: null },
+      ]);
+    } finally {
+      vi.useRealTimers();
+    }
+  });
+
+  it("deletes core and legacy storage", async () => {
+    vi.useFakeTimers();
+    try {
+      const t = convexTest(schema, modules);
+      const currentId = await t.mutation(api.lib.createStream, {});
+      await t.mutation(api.lib.addChunk, {
+        streamId: currentId,
+        text: "current",
+        final: true,
+      });
+      const legacyId = await t.run(async (ctx) => {
+        const streamId = await ctx.db.insert("streams", { status: "done" });
+        await ctx.db.insert("chunks", { streamId, text: "legacy" });
+        return streamId;
+      });
+      await t.mutation(api.lib.deleteStream, { streamId: currentId });
+      await t.mutation(api.lib.deleteStream, { streamId: legacyId });
+      await t.finishAllScheduledFunctions(vi.runAllTimers);
+      await expect(
+        t.run(async (ctx) => ({
+          facades: await ctx.db.query("streams").collect(),
+          chunks: await ctx.db.query("chunks").collect(),
+          cores: await ctx.db.query("textStreams").collect(),
+          events: await ctx.db.query("textStreamsEvents").collect(),
+        })),
+      ).resolves.toEqual({ facades: [], chunks: [], cores: [], events: [] });
+    } finally {
+      vi.useRealTimers();
+    }
+  });
 });

src/component/lib.ts

@@ -133,12 +133,13 @@ function readResult(
   );
 }
 
-// Create a new stream with zero chunks.
 export const createStream = mutation({
   args: {},
   handler: async (ctx) => {
+    const core = await textStreams.create(ctx);
     const streamId = await ctx.db.insert("streams", {
       status: "pending",
+      coreId: core.streamId,
     });
     return streamId;
   },
@@ -158,33 +159,35 @@ export const addChunk = mutation({
     if (!stream) {
       throw new Error("Stream not found");
     }
-    if (stream.status === "pending") {
-      await ctx.db.patch("streams", args.streamId, {
-        status: "streaming",
-      });
-    } else if (stream.status !== "streaming") {
+    if (stream.status !== "pending" && stream.status !== "streaming") {
       throw new Error("Stream is not streaming; did it timeout?");
     }
-    await ctx.db.insert("chunks", {
-      streamId: args.streamId,
-      text: args.text,
-    });
-    if (args.final) {
+    if (stream.coreId !== undefined) {
+      await textStreams.append(ctx, {
+        streamId: stream.coreId,
+        event: args.text,
+        ...(args.final ? { complete: true as const } : {}),
+      });
+    } else {
+      await ctx.db.insert("chunks", {
+        streamId: args.streamId,
+        text: args.text,
+      });
+    }
+    if (args.final || stream.status === "pending") {
       await ctx.db.patch("streams", args.streamId, {
-        status: "done",
+        status: args.final ? "done" : "streaming",
       });
     }
   },
 });
 
-// Set the status of a stream.
+// Set a terminal status on a stream.
 // Can only be done on streams which are pending or streaming.
 export const setStreamStatus = mutation({
   args: {
     streamId: v.id("streams"),
     status: v.union(
-      v.literal("pending"),
-      v.literal("streaming"),
       v.literal("done"),
       v.literal("error"),
       v.literal("timeout"),
@@ -202,6 +205,22 @@ export const setStreamStatus = mutation({
       );
       return;
     }
+    if (stream.coreId !== undefined) {
+      if (args.status === "done") {
+        await textStreams.complete(ctx, { streamId: stream.coreId });
+      } else {
+        await textStreams.fail(ctx, {
+          streamId: stream.coreId,
+          error: {
+            code: args.status,
+            message:
+              args.status === "timeout"
+                ? "Stream generation timed out."
+                : "Stream generation failed.",
+          },
+        });
+      }
+    }
     await ctx.db.patch("streams", args.streamId, {
       status: args.status,
     });
@@ -352,8 +371,7 @@ const EXPIRATION_TIME = 20 * 60 * 1000; // 20 minutes in milliseconds
 const BATCH_SIZE = 100;
 const DELETE_BATCH_SIZE = 64;
 
-// Delete a stream and all its chunks.
-// The stream is deleted immediately; chunks are cleaned up asynchronously.
+// Delete a stream and its persisted events.
 export const deleteStream = mutation({
   args: {
     streamId: v.id("streams"),
@@ -364,11 +382,19 @@ export const deleteStream = mutation({
     if (!stream) {
       throw new Error(`Stream ${args.streamId} not found`);
     }
+    if (stream.coreId !== undefined) {
+      await textStreams.delete(
+        ctx,
+        { streamId: stream.coreId },
+        { run: internal.lib.run },
+      );
+    } else {
+      await ctx.scheduler.runAfter(0, internal.lib._deleteChunksPage, {
+        streamId: args.streamId,
+        cursor: null,
+      });
+    }
     await ctx.db.delete("streams", args.streamId);
-    await ctx.scheduler.runAfter(0, internal.lib._deleteChunksPage, {
-      streamId: args.streamId,
-      cursor: null,
-    });
   },
 });
 
@@ -398,8 +424,16 @@ export const _deleteChunksPage = internalMutation({
   },
 });
 
-// If the last chunk of a stream was added more than 20 minutes ago,
-// set the stream to timeout. The action feeding it has to be dead.
+export const run = internalMutation({
+  args: textStreams.args.run,
+  returns: textStreams.returns.run,
+  handler: async (ctx, args): Promise<{ isDone: boolean }> =>
+    await textStreams.run(ctx, args, {
+      run: internal.lib.run,
+    }),
+});
+
+// Time out streams after 20 minutes without a persisted append.
 export const cleanupExpiredStreams = internalMutation({
   args: {},
   handler: async (ctx) => {
@@ -414,8 +448,40 @@ export const cleanupExpiredStreams = internalMutation({
       .take(BATCH_SIZE);
 
     for (const stream of [...pendingStreams, ...streamingStreams]) {
-      if (now - stream._creationTime > EXPIRATION_TIME) {
+      if (now - stream._creationTime <= EXPIRATION_TIME) {
+        continue;
+      }
+      const core =
+        stream.coreId === undefined
+          ? null
+          : await ctx.db.get("textStreams", stream.coreId);
+      const latestChunk =
+        stream.coreId === undefined
+          ? await ctx.db
+              .query("chunks")
+              .withIndex("byStream", (q) => q.eq("streamId", stream._id))
+              .order("desc")
+              .first()
+          : null;
+      const lastActivity =
+        core?.lastAppendedAt ??
+        core?._creationTime ??
+        latestChunk?._creationTime ??
+        stream._creationTime;
+      if (now - lastActivity > EXPIRATION_TIME) {
         console.log("Cleaning up expired stream", stream._id);
+        if (
+          stream.coreId !== undefined &&
+          (core?.status === "pending" || core?.status === "streaming")
+        ) {
+          await textStreams.fail(ctx, {
+            streamId: stream.coreId,
+            error: {
+              code: "timeout",
+              message: "Stream generation timed out.",
+            },
+          });
+        }
         await ctx.db.patch("streams", stream._id, {
           status: "timeout",
         });