@convex-dev/agent 0.2.0-alpha.2 → 0.2.0

package/src/component/messages.ts

@@ -1,6 +1,9 @@
-import { assert, omit } from "convex-helpers";
+import { assert, omit, pick } from "convex-helpers";
 import { mergedStream, stream } from "convex-helpers/server/stream";
-import { paginationOptsValidator } from "convex/server";
+import {
+  paginationOptsValidator,
+  type WithoutSystemFields,
+} from "convex/server";
 import type { ObjectType } from "convex/values";
 import {
   DEFAULT_MESSAGE_RANGE,
@@ -26,11 +29,6 @@ import {
 } from "./_generated/server.js";
 import type { MessageDoc } from "./schema.js";
 import { schema, v, vMessageDoc } from "./schema.js";
-import {
-  getThread as _getThread,
-  listThreadsByUserId as _listThreadsByUserId,
-  updateThread as _updateThread,
-} from "./threads.js";
 import { insertVector, searchVectors } from "./vector/index.js";
 import {
   type VectorDimension,
@@ -39,6 +37,8 @@ import {
   vVectorId,
 } from "./vector/tables.js";
 import { changeRefcount } from "./files.js";
+import { getStreamingMessagesWithMetadata } from "./streams.js";
+import { partial } from "convex-helpers/validators";
 
 function publicMessage(message: Doc<"messages">): MessageDoc {
   return omit(message, ["parentMessageId", "stepId", "files"]);
@@ -127,6 +127,8 @@ const addMessagesArgs = {
   messages: v.array(vMessageWithMetadataInternal),
   embeddings: v.optional(vMessageEmbeddingsWithDimension),
   failPendingSteps: v.optional(v.boolean()),
+  // A pending message to update. If the pending message failed, abort.
+  pendingMessageId: v.optional(v.id("messages")),
 };
 export const addMessages = mutation({
   args: addMessagesArgs,
@@ -144,9 +146,15 @@ async function addMessagesHandler(
     assert(thread, `Thread ${args.threadId} not found`);
     userId = thread.userId;
   }
-  const { embeddings, failPendingSteps, messages, promptMessageId, ...rest } =
-    args;
-  const parentMessage = promptMessageId && (await ctx.db.get(promptMessageId));
+  const {
+    embeddings,
+    failPendingSteps,
+    messages,
+    promptMessageId,
+    pendingMessageId,
+    ...rest
+  } = args;
+  const promptMessage = promptMessageId && (await ctx.db.get(promptMessageId));
   if (failPendingSteps) {
     assert(args.threadId, "threadId is required to fail pending steps");
     const pendingMessages = await ctx.db
@@ -154,26 +162,37 @@ async function addMessagesHandler(
       .withIndex("threadId_status_tool_order_stepOrder", (q) =>
         q.eq("threadId", threadId).eq("status", "pending"),
       )
-      .collect();
+      .order("desc")
+      .take(100);
     await Promise.all(
       pendingMessages
-        .filter((m) => !parentMessage || m.order === parentMessage.order)
-        .map((m) =>
-          ctx.db.patch(m._id, { status: "failed", error: "Restarting" }),
-        ),
+        .filter((m) => !promptMessage || m.order === promptMessage.order)
+        .filter((m) => !pendingMessageId || m._id !== pendingMessageId)
+        .map(async (m) => {
+          if (m.embeddingId) {
+            await ctx.db.delete(m.embeddingId);
+          }
+          await ctx.db.patch(m._id, {
+            status: "failed",
+            error: "Restarting",
+            embeddingId: undefined,
+          });
+        }),
     );
   }
   let order, stepOrder;
   let fail = false;
+  let error: string | undefined;
   if (promptMessageId) {
-    assert(parentMessage, `Parent message ${promptMessageId} not found`);
-    if (parentMessage.status === "failed") {
+    assert(promptMessage, `Parent message ${promptMessageId} not found`);
+    if (promptMessage.status === "failed") {
       fail = true;
+      error = promptMessage.error ?? error ?? "The prompt message failed";
     }
-    order = parentMessage.order;
+    order = promptMessage.order;
     // Defend against there being existing messages with this parent.
     const maxMessage = await getMaxMessage(ctx, threadId, order);
-    stepOrder = maxMessage?.stepOrder ?? parentMessage.stepOrder;
+    stepOrder = maxMessage?.stepOrder ?? promptMessage.stepOrder;
   } else {
     const maxMessage = await getMaxMessage(ctx, threadId);
     order = maxMessage?.order ?? -1;
@@ -189,7 +208,12 @@ async function addMessagesHandler(
   for (let i = 0; i < messages.length; i++) {
     const message = messages[i];
     let embeddingId: VectorTableId | undefined;
-    if (embeddings && embeddings.vectors[i]) {
+    if (
+      embeddings &&
+      embeddings.vectors[i] &&
+      !fail &&
+      message.status !== "failed"
+    ) {
       embeddingId = await insertVector(ctx, embeddings.dimension, {
         vector: embeddings.vectors[i]!,
         model: embeddings.model,
@@ -198,8 +222,50 @@ async function addMessagesHandler(
         threadId,
       });
     }
+    const messageDoc = {
+      ...rest,
+      ...message,
+      embeddingId,
+      parentMessageId: promptMessageId,
+      userId,
+      tool: isTool(message.message),
+      text: extractText(message.message),
+      status: fail ? "failed" : (message.status ?? "success"),
+      error: fail ? error : message.error,
+    } satisfies Omit<
+      WithoutSystemFields<Doc<"messages">>,
+      "order" | "stepOrder"
+    >;
+    // If there is a pending message, we replace that one with the first message
+    // and subsequent ones will follow the regular order/subOrder advancement.
+    if (i === 0 && pendingMessageId) {
+      const pendingMessage = await ctx.db.get(pendingMessageId);
+      assert(pendingMessage, `Pending msg ${pendingMessageId} not found`);
+      if (pendingMessage.status === "failed") {
+        fail = true;
+        error =
+          `Trying to update a message that failed: ${pendingMessageId}, ` +
+          `error: ${pendingMessage.error ?? error}`;
+        messageDoc.status = "failed";
+        messageDoc.error = error;
+      }
+      if (message.fileIds) {
+        await changeRefcount(
+          ctx,
+          pendingMessage.fileIds ?? [],
+          message.fileIds,
+        );
+      }
+      await ctx.db.replace(pendingMessage._id, {
+        ...messageDoc,
+        order: pendingMessage.order,
+        stepOrder: pendingMessage.stepOrder,
+      });
+      toReturn.push(pendingMessage);
+      continue;
+    }
     if (message.message.role === "user") {
-      if (parentMessage && parentMessage.order === order) {
+      if (promptMessage && promptMessage.order === order) {
         // see if there's a later message than the parent message order
         const maxMessage = await getMaxMessage(ctx, threadId);
         order = (maxMessage?.order ?? order) + 1;
@@ -208,21 +274,14 @@ async function addMessagesHandler(
       }
       stepOrder = 0;
     } else {
+      if (order < 0) {
+        order = 0;
+      }
       stepOrder++;
     }
     const messageId = await ctx.db.insert("messages", {
-      ...rest,
-      ...message,
-      embeddingId,
-      parentMessageId: promptMessageId,
-      userId,
+      ...messageDoc,
       order,
-      tool: isTool(message.message),
-      text: extractText(message.message),
-      status: fail ? "failed" : (message.status ?? "success"),
-      error: fail
-        ? (parentMessage?.error ?? "Parent message failed")
-        : undefined,
       stepOrder,
     });
     if (message.fileIds) {
@@ -271,51 +330,77 @@ function orderedMessagesStream(
   );
 }
 
-export const rollbackMessage = mutation({
+export const finalizeMessage = mutation({
   args: {
     messageId: v.id("messages"),
-    error: v.optional(v.string()),
+    result: v.union(
+      v.object({ status: v.literal("success") }),
+      v.object({ status: v.literal("failed"), error: v.string() }),
+    ),
   },
   returns: v.null(),
-  handler: async (ctx, { messageId, error }) => {
+  handler: async (ctx, { messageId, result }) => {
     const message = await ctx.db.get(messageId);
     assert(message, `Message ${messageId} not found`);
-    const messages = await orderedMessagesStream(
-      ctx,
-      message.threadId,
-      "asc",
-      message.order,
-    ).collect();
-    for (const m of messages) {
-      if (m.status === "pending") {
-        await ctx.db.patch(m._id, { status: "failed", error });
+    if (message.status !== "pending") {
+      console.log(
+        "Trying to finalize a message that's already",
+        message.status,
+      );
+      return;
+    }
+    // See if we can add any in-progress data
+    if (message.message === undefined) {
+      const messages = await getStreamingMessagesWithMetadata(
+        ctx,
+        message,
+        result,
+      );
+      if (messages.length > 0) {
+        await addMessagesHandler(ctx, {
+          messages,
+          threadId: message.threadId,
+          agentName: message.agentName,
+          failPendingSteps: false,
+          pendingMessageId: messageId,
+          userId: message.userId,
+          embeddings: undefined,
+        });
+        return;
       }
     }
-
-    await ctx.db.patch(messageId, {
-      status: "failed",
-      error: error,
-    });
-  },
-});
-
-export const commitMessage = mutation({
-  args: {
-    messageId: v.id("messages"),
+    if (result.status === "failed") {
+      if (message.embeddingId) {
+        await ctx.db.delete(message.embeddingId);
+      }
+      await ctx.db.patch(messageId, {
+        status: "failed",
+        error: result.error,
+        embeddingId: undefined,
+      });
+    } else {
+      await ctx.db.patch(messageId, { status: "success" });
+    }
   },
-  returns: v.null(),
-  handler: commitMessageHandler,
 });
 
 export const updateMessage = mutation({
   args: {
     messageId: v.id("messages"),
-    patch: v.object({
-      message: v.optional(vMessageDoc.fields.message),
-      fileIds: v.optional(v.array(v.id("files"))),
-      status: v.optional(vMessageStatus),
-      error: v.optional(v.string()),
-    }),
+    patch: v.object(
+      partial(
+        pick(schema.tables.messages.validator.fields, [
+          "message",
+          "fileIds",
+          "status",
+          "error",
+          "model",
+          "provider",
+          "providerOptions",
+          "finishReason",
+        ]),
+      ),
+    ),
   },
   returns: vMessageDoc,
   handler: async (ctx, args) => {
@@ -326,9 +411,7 @@ export const updateMessage = mutation({
       await changeRefcount(ctx, message.fileIds ?? [], args.patch.fileIds);
     }
 
-    const patch: Partial<Doc<"messages">> = {
-      ...args.patch,
-    };
+    const patch: Partial<Doc<"messages">> = { ...args.patch };
 
     if (args.patch.message !== undefined) {
       patch.message = args.patch.message;
@@ -336,38 +419,18 @@ export const updateMessage = mutation({
       patch.text = extractText(args.patch.message);
     }
 
+    if (args.patch.status === "failed") {
+      if (message.embeddingId) {
+        await ctx.db.delete(message.embeddingId);
+      }
+      patch.embeddingId = undefined;
+    }
+
     await ctx.db.patch(args.messageId, patch);
     return publicMessage((await ctx.db.get(args.messageId))!);
   },
 });
 
-async function commitMessageHandler(
-  ctx: MutationCtx,
-  { messageId }: { messageId: Id<"messages"> },
-) {
-  const message = await ctx.db.get(messageId);
-  assert(message, `Message ${messageId} not found`);
-
-  const order = message.order!;
-  const messages = await mergedStream(
-    [true, false].map((tool) =>
-      stream(ctx.db, schema)
-        .query("messages")
-        .withIndex("threadId_status_tool_order_stepOrder", (q) =>
-          q
-            .eq("threadId", message.threadId)
-            .eq("status", "pending")
-            .eq("tool", tool)
-            .eq("order", order),
-        ),
-    ),
-    ["order", "stepOrder"],
-  ).collect();
-  for (const message of messages) {
-    await ctx.db.patch(message._id, { status: "success" });
-  }
-}
-
 export const listMessagesByThreadId = query({
   args: {
     threadId: v.id("threads"),
@@ -427,9 +490,7 @@ export const listMessagesByThreadId = query({
 });
 
 export const getMessagesByIds = query({
-  args: {
-    messageIds: v.array(v.id("messages")),
-  },
+  args: { messageIds: v.array(v.id("messages")) },
   handler: async (ctx, args) => {
     return (await Promise.all(args.messageIds.map((id) => ctx.db.get(id)))).map(
       (m) => (m ? publicMessage(m) : null),

package/src/component/schema.ts

@@ -11,7 +11,6 @@ import {
   vProviderOptions,
   vProviderMetadata,
   vReasoningDetails,
-  vTextStreamPart,
 } from "../validators.js";
 import { typedV } from "convex-helpers/validators";
 import vectorTables, { vVectorId } from "./vector/tables.js";
@@ -56,10 +55,11 @@ export const schema = defineSchema({
     usage: v.optional(vUsage),
     providerMetadata: v.optional(vProviderMetadata), // Received from model
     sources: v.optional(v.array(vSource)),
-    reasoning: v.optional(v.string()),
-    reasoningDetails: v.optional(vReasoningDetails),
     warnings: v.optional(v.array(vLanguageModelCallWarning)),
     finishReason: v.optional(vFinishReason),
+    // Likely deprecated soon
+    reasoning: v.optional(v.string()),
+    reasoningDetails: v.optional(vReasoningDetails),
     // DEPRECATED
     id: v.optional(v.string()), // external id, e.g. from Vercel AI SDK
     parentMessageId: v.optional(v.id("messages")),
@@ -129,7 +129,7 @@ export const schema = defineSchema({
     // the indexes work like: 0 <first> 1 <second> 2 <third> 3 ...
     start: v.number(), // inclusive
     end: v.number(), // exclusive
-    parts: v.array(vTextStreamPart),
+    parts: v.array(v.any()),
   }).index("streamId_start_end", ["streamId", "start", "end"]),
 
   memories: defineTable({