@convex-dev/agent 0.1.18 → 0.2.0

package/src/component/streams.ts

@@ -1,6 +1,7 @@
 import { omit, pick } from "convex-helpers";
 import { v } from "convex/values";
 import {
+  type MessageWithMetadataInternal,
   type StreamDelta,
   vStreamDelta,
   vStreamMessage,
@@ -18,6 +19,9 @@ import schema from "./schema.js";
 import { stream } from "convex-helpers/server/stream";
 import { mergedStream } from "convex-helpers/server/stream";
 import { paginator } from "convex-helpers/server/pagination";
+import type { WithoutSystemFields } from "convex/server";
+import { mergeDeltas } from "../react/deltas.js";
+import { serializeOrThrow } from "../mapping.js";
 
 const SECOND = 1000;
 const MINUTE = 60 * SECOND;
@@ -33,13 +37,16 @@ export const addDelta = mutation({
   args: deltaValidator,
   returns: v.boolean(),
   handler: async (ctx, args) => {
-    await ctx.db.insert("streamDeltas", args);
-    await heartbeatStream(ctx, { streamId: args.streamId });
     const stream = await ctx.db.get(args.streamId);
-    if (stream?.state.kind !== "streaming") {
-      console.warn(`Stream is not streaming: ${args.streamId}`);
+    if (!stream) {
+      console.warn("Stream not found", args.streamId);
+      return false;
+    }
+    if (stream.state.kind !== "streaming") {
       return false;
     }
+    await ctx.db.insert("streamDeltas", args);
+    await heartbeatStream(ctx, { streamId: args.streamId });
     return true;
   },
 });
@@ -48,10 +55,7 @@ export const listDeltas = query({
   args: {
     threadId: v.id("threads"),
     cursors: v.array(
-      v.object({
-        streamId: v.id("streamingMessages"),
-        cursor: v.number(),
-      }),
+      v.object({ streamId: v.id("streamingMessages"), cursor: v.number() }),
     ),
   },
   returns: v.array(vStreamDelta),
@@ -85,10 +89,7 @@ export const create = mutation({
   args: omit(schema.tables.streamingMessages.validator.fields, ["state"]),
   returns: v.id("streamingMessages"),
   handler: async (ctx, args) => {
-    const state = {
-      kind: "streaming" as const,
-      lastHeartbeat: Date.now(),
-    };
+    const state = { kind: "streaming" as const, lastHeartbeat: Date.now() };
     const streamId = await ctx.db.insert("streamingMessages", {
       ...args,
       state,
@@ -152,11 +153,7 @@ export const list = query({
 });
 
 export const abortByOrder = mutation({
-  args: {
-    threadId: v.id("threads"),
-    order: v.number(),
-    reason: v.string(),
-  },
+  args: { threadId: v.id("threads"), order: v.number(), reason: v.string() },
   returns: v.boolean(),
   handler: async (ctx, args) => {
     const streams = await ctx.db
@@ -169,10 +166,7 @@ export const abortByOrder = mutation({
       )
       .take(100);
     for (const stream of streams) {
-      await abortById(ctx, {
-        streamId: stream._id,
-        reason: args.reason,
-      });
+      await abortById(ctx, { streamId: stream._id, reason: args.reason });
     }
     return streams.length > 0;
   },
@@ -182,6 +176,7 @@ export const abort = mutation({
   args: {
     streamId: v.id("streamingMessages"),
     reason: v.string(),
+    finalDelta: v.optional(deltaValidator),
   },
   returns: v.boolean(),
   handler: abortById,
@@ -189,16 +184,20 @@ export const abort = mutation({
 
 async function abortById(
   ctx: MutationCtx,
-  args: { streamId: Id<"streamingMessages">; reason: string },
+  args: {
+    streamId: Id<"streamingMessages">;
+    reason: string;
+    finalDelta?: WithoutSystemFields<Doc<"streamDeltas">>;
+  },
 ) {
   const stream = await ctx.db.get(args.streamId);
   if (!stream) {
     throw new Error(`Stream not found: ${args.streamId}`);
   }
+  if (args.finalDelta) {
+    await ctx.db.insert("streamDeltas", args.finalDelta);
+  }
   if (stream.state.kind !== "streaming") {
-    console.warn(
-      `Stream trying to abort but not currently streaming (${stream.state.kind}): ${args.streamId}`,
-    );
     return false;
   }
   await cleanupTimeoutFn(ctx, stream);
@@ -220,49 +219,65 @@ async function cleanupTimeoutFn(
   }
 }
 
+// No longer used from the DeltaStreamer
 export const finish = mutation({
   args: {
     streamId: v.id("streamingMessages"),
     finalDelta: v.optional(deltaValidator),
   },
   returns: v.null(),
-  handler: async (ctx, args) => {
-    if (args.finalDelta) {
-      await ctx.db.insert("streamDeltas", args.finalDelta);
-    }
-    const stream = await ctx.db.get(args.streamId);
-    if (!stream) {
-      throw new Error(`Stream not found: ${args.streamId}`);
-    }
-    if (stream.state.kind !== "streaming") {
-      console.warn(
-        `Stream trying to finish but not currently streaming: ${args.streamId}`,
-      );
-      return;
-    }
-    await cleanupTimeoutFn(ctx, stream);
-    const cleanupFnId = await ctx.scheduler.runAfter(
-      DELETE_STREAM_DELAY,
-      api.streams.deleteStreamAsync,
-      { streamId: args.streamId },
-    );
-    await ctx.db.patch(args.streamId, {
-      state: { kind: "finished", endedAt: Date.now(), cleanupFnId },
-    });
+  handler: finishHandler,
+});
+
+export async function finishHandler(
+  ctx: MutationCtx,
+  args: {
+    streamId: Id<"streamingMessages">;
+    finalDelta?: WithoutSystemFields<Doc<"streamDeltas">>;
   },
+) {
+  if (args.finalDelta) {
+    await ctx.db.insert("streamDeltas", args.finalDelta);
+  }
+  const stream = await ctx.db.get(args.streamId);
+  if (!stream) {
+    throw new Error(`Stream not found: ${args.streamId}`);
+  }
+  if (stream.state.kind !== "streaming") {
+    console.warn(
+      `Stream trying to finish ${args.streamId} but is ${stream.state.kind}`,
+    );
+    return;
+  }
+  await cleanupTimeoutFn(ctx, stream);
+  const cleanupFnId = await ctx.scheduler.runAfter(
+    DELETE_STREAM_DELAY,
+    api.streams.deleteStreamAsync,
+    { streamId: args.streamId },
+  );
+  await ctx.db.patch(args.streamId, {
+    state: { kind: "finished", endedAt: Date.now(), cleanupFnId },
+  });
+}
+
+// TODO: use this heartbeat while streaming, every 30 seconds or so,
+// then reduce the timeout to 60 seconds.
+export const heartbeat = mutation({
+  args: { streamId: v.id("streamingMessages") },
+  returns: v.null(),
+  handler: heartbeatStream,
 });
 
 async function heartbeatStream(
   ctx: MutationCtx,
   args: { streamId: Id<"streamingMessages"> },
-) {
+): Promise<void> {
   const stream = await ctx.db.get(args.streamId);
   if (!stream) {
     console.warn("Stream not found", args.streamId);
     return;
   }
   if (stream.state.kind !== "streaming") {
-    console.warn("Stream is not streaming", args.streamId);
     return;
   }
   if (Date.now() - stream.state.lastHeartbeat < TIMEOUT_INTERVAL / 4) {
@@ -286,11 +301,7 @@ async function heartbeatStream(
     { streamId: args.streamId },
   );
   await ctx.db.patch(args.streamId, {
-    state: {
-      kind: "streaming",
-      lastHeartbeat: Date.now(),
-      timeoutFnId,
-    },
+    state: { kind: "streaming", lastHeartbeat: Date.now(), timeoutFnId },
   });
 }
 
@@ -304,10 +315,7 @@ export const timeoutStream = internalMutation({
       return;
     }
     await ctx.db.patch(args.streamId, {
-      state: {
-        kind: "aborted",
-        reason: "timeout",
-      },
+      state: { kind: "aborted", reason: "timeout" },
     });
   },
 });
@@ -362,11 +370,7 @@ export async function deleteStreamsPageForThreadId(
     "stepOrder",
   ]).first();
   if (!streamMessage) {
-    return {
-      isDone: true,
-      streamOrder: undefined,
-      deltaCursor: undefined,
-    };
+    return { isDone: true, streamOrder: undefined, deltaCursor: undefined };
   }
   const result = await deletePageForStreamId(ctx, {
     streamId: streamMessage._id,
@@ -375,11 +379,7 @@ export async function deleteStreamsPageForThreadId(
   if (result.isDone) {
     deltaCursor = undefined;
   }
-  return {
-    isDone: false,
-    streamOrder: streamMessage.order,
-    deltaCursor,
-  };
+  return { isDone: false, streamOrder: streamMessage.order, deltaCursor };
 }
 
 export const deleteStreamsPageForThreadIdMutation = internalMutation({
@@ -474,3 +474,114 @@ export const deleteAllStreamsForThreadIdSync = action({
     }
   },
 });
+
+export async function getStreamingMessages(
+  ctx: MutationCtx,
+  threadId: Id<"threads">,
+  order: number,
+  stepOrder: number,
+): Promise<Doc<"streamingMessages">[]> {
+  return mergedStream(
+    (["aborted", "streaming", "finished"] as const).map((state) =>
+      stream(ctx.db, schema)
+        .query("streamingMessages")
+        .withIndex("threadId_state_order_stepOrder", (q) =>
+          q
+            .eq("threadId", threadId)
+            .eq("state.kind", state)
+            .eq("order", order)
+            .lte("stepOrder", stepOrder),
+        )
+        .order("desc"),
+    ),
+    ["stepOrder"],
+  ).take(10);
+}
+
+export async function getStreamingMessagesWithMetadata(
+  ctx: MutationCtx,
+  {
+    threadId,
+    order,
+    stepOrder,
+  }: { threadId: Id<"threads">; order: number; stepOrder: number },
+  metadata: { status: "success" | "failed"; error?: string },
+): Promise<MessageWithMetadataInternal[]> {
+  // See if there are any streaming messages for this order
+  const streamingMessages = await getStreamingMessages(
+    ctx,
+    threadId,
+    order,
+    stepOrder,
+  );
+  const messages = (
+    await Promise.all(
+      streamingMessages.map((m) =>
+        getMessagesWithMetadataForStreamingMessage(
+          ctx,
+          threadId,
+          stepOrder,
+          m,
+          metadata,
+        ),
+      ),
+    )
+  ).flat();
+  return messages;
+}
+
+export async function getMessagesWithMetadataForStreamingMessage(
+  ctx: MutationCtx,
+  threadId: Id<"threads">,
+  stepOrder: number,
+  streamingMessage: Doc<"streamingMessages">,
+  metadata: { status: "success" | "failed"; error?: string },
+): Promise<MessageWithMetadataInternal[]> {
+  const deltas = await ctx.db
+    .query("streamDeltas")
+    .withIndex("streamId_start_end", (q) =>
+      q.eq("streamId", streamingMessage._id),
+    )
+    .take(1000);
+  const [messageDocs] = mergeDeltas(
+    threadId,
+    [
+      {
+        ...streamingMessage,
+        status: "streaming",
+        streamId: streamingMessage._id,
+      },
+    ],
+    [],
+    deltas,
+  );
+  // We don't save messages that have already been saved
+  const numToSkip = stepOrder - streamingMessage.stepOrder;
+  const messages = await Promise.all(
+    messageDocs
+      .slice(numToSkip)
+      .filter((m) => m.message !== undefined)
+      .map(async (msg) => {
+        const message = await serializeOrThrow(msg.message!);
+        return {
+          message,
+          ...pick(msg, [
+            "fileIds",
+            "status",
+            "finishReason",
+            "model",
+            "provider",
+            "providerMetadata",
+            "sources",
+            "reasoning",
+            "reasoningDetails",
+            "usage",
+            "warnings",
+            "error",
+          ]),
+          ...metadata,
+        } as MessageWithMetadataInternal;
+      }),
+  );
+  return messages;
+}