@convex-dev/agent 0.1.18 → 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,
@@ -9,7 +12,7 @@ import {
   isTool,
 } from "../shared.js";
 import {
-  vMessageEmbeddings,
+  vMessageEmbeddingsWithDimension,
   vMessageStatus,
   vMessageWithMetadataInternal,
   vPaginationResult,
@@ -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"]);
@@ -58,9 +58,7 @@ export async function deleteMessage(
 }
 
 export const deleteByIds = mutation({
-  args: {
-    messageIds: v.array(v.id("messages")),
-  },
+  args: { messageIds: v.array(v.id("messages")) },
   returns: v.array(v.id("messages")),
   handler: async (ctx, args) => {
     const deletedMessageIds = await Promise.all(
@@ -115,8 +113,8 @@ export const deleteByOrder = mutation({
     await Promise.all(messages.map((m) => deleteMessage(ctx, m)));
     return {
       isDone: messages.length < 64,
-      lastOrder: messages[messages.length - 1]?.order,
-      lastStepOrder: messages[messages.length - 1]?.stepOrder,
+      lastOrder: messages.at(-1)?.order,
+      lastStepOrder: messages.at(-1)?.stepOrder,
     };
   },
 });
@@ -127,16 +125,15 @@ const addMessagesArgs = {
   promptMessageId: v.optional(v.id("messages")),
   agentName: v.optional(v.string()),
   messages: v.array(vMessageWithMetadataInternal),
-  embeddings: v.optional(vMessageEmbeddings),
-  pending: v.optional(v.boolean()),
+  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,
   handler: addMessagesHandler,
-  returns: v.object({
-    messages: v.array(vMessageDoc),
-  }),
+  returns: v.object({ messages: v.array(vMessageDoc) }),
 });
 async function addMessagesHandler(
   ctx: MutationCtx,
@@ -152,12 +149,12 @@ async function addMessagesHandler(
   const {
     embeddings,
     failPendingSteps,
-    pending,
     messages,
     promptMessageId,
+    pendingMessageId,
     ...rest
   } = args;
-  const parentMessage = promptMessageId && (await ctx.db.get(promptMessageId));
+  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
@@ -165,30 +162,41 @@ 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 ? maxMessage.order + 1 : 0;
-    stepOrder = -1;
+    order = maxMessage?.order ?? -1;
+    stepOrder = maxMessage?.stepOrder ?? -1;
   }
   const toReturn: Doc<"messages">[] = [];
   if (embeddings) {
@@ -200,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,
@@ -209,26 +222,68 @@ async function addMessagesHandler(
         threadId,
       });
     }
-    stepOrder++;
-    const messageId = await ctx.db.insert("messages", {
+    const messageDoc = {
       ...rest,
       ...message,
       embeddingId,
       parentMessageId: promptMessageId,
       userId,
-      order,
       tool: isTool(message.message),
       text: extractText(message.message),
-      status: fail ? "failed" : pending ? "pending" : "success",
-      error: fail ? "Parent message failed" : undefined,
+      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 (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;
+      } else {
+        order++;
+      }
+      stepOrder = 0;
+    } else {
+      if (order < 0) {
+        order = 0;
+      }
+      stepOrder++;
+    }
+    const messageId = await ctx.db.insert("messages", {
+      ...messageDoc,
+      order,
       stepOrder,
     });
-    // Let's just not set the id field and have it set only in explicit cases.
-    // if (!message.id) {
-    //   await ctx.db.patch(messageId, {
-    //     id: messageId,
-    //   });
-    // }
     if (message.fileIds) {
       await changeRefcount(ctx, [], message.fileIds);
     }
@@ -263,7 +318,7 @@ function orderedMessagesStream(
               .eq("threadId", threadId)
               .eq("status", status)
               .eq("tool", tool);
-            if (order) {
+            if (order !== undefined) {
               return qq.eq("order", order);
             }
             return qq;
@@ -275,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) => {
@@ -330,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;
@@ -340,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"),
@@ -431,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

@@ -6,12 +6,11 @@ import {
   vMessageStatus,
   vUsage,
   vSource,
-  vLanguageModelV1CallWarning,
+  vLanguageModelCallWarning,
   vFinishReason,
   vProviderOptions,
   vProviderMetadata,
   vReasoningDetails,
-  vTextStreamPart,
 } from "../validators.js";
 import { typedV } from "convex-helpers/validators";
 import vectorTables, { vVectorId } from "./vector/tables.js";
@@ -29,12 +28,8 @@ export const schema = defineSchema({
     order: /*DEPRECATED*/ v.optional(v.number()),
   })
     .index("userId", ["userId"])
-    .searchIndex("title", {
-      searchField: "title",
-      filterFields: ["userId"],
-    }),
+    .searchIndex("title", { searchField: "title", filterFields: ["userId"] }),
   messages: defineTable({
-    id: v.optional(v.string()), // external id, e.g. from Vercel AI SDK
     userId: v.optional(v.string()), // useful for searching across threads
     threadId: v.id("threads"),
     order: v.number(),
@@ -60,11 +55,13 @@ export const schema = defineSchema({
     usage: v.optional(vUsage),
     providerMetadata: v.optional(vProviderMetadata), // Received from model
     sources: v.optional(v.array(vSource)),
+    warnings: v.optional(v.array(vLanguageModelCallWarning)),
+    finishReason: v.optional(vFinishReason),
+    // Likely deprecated soon
     reasoning: v.optional(v.string()),
     reasoningDetails: v.optional(vReasoningDetails),
-    warnings: v.optional(v.array(vLanguageModelV1CallWarning)),
-    finishReason: v.optional(vFinishReason),
     // DEPRECATED
+    id: v.optional(v.string()), // external id, e.g. from Vercel AI SDK
     parentMessageId: v.optional(v.id("messages")),
     stepId: v.optional(v.string()),
     files: v.optional(v.array(v.any())),
@@ -115,10 +112,7 @@ export const schema = defineSchema({
         endedAt: v.number(),
         cleanupFnId: v.optional(v.id("_scheduled_functions")),
       }),
-      v.object({
-        kind: v.literal("aborted"),
-        reason: v.string(),
-      }),
+      v.object({ kind: v.literal("aborted"), reason: v.string() }),
     ),
   })
     // There should only be one per "order" index
@@ -135,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({
@@ -162,9 +156,9 @@ export const schema = defineSchema({
   // To authenticate playground usage
   // Delete a key to invalidate it
   // Provide a name to easily identify it / invalidate by name
-  apiKeys: defineTable({
-    name: v.optional(v.string()),
-  }).index("name", ["name"]),
+  apiKeys: defineTable({ name: v.optional(v.string()) }).index("name", [
+    "name",
+  ]),
 });
 
 export const vv = typedV(schema);