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

package/src/component/messages.test.ts

@@ -30,6 +30,29 @@ describe("agent", () => {
       stepOrder: 1,
     });
   });
+  test("getMaxMessage works for a specific order", async () => {
+    const t = convexTest(schema, modules);
+    const thread = await t.mutation(api.threads.createThread, {
+      userId: "test",
+    });
+    const { messages } = await t.mutation(api.messages.addMessages, {
+      threadId: thread._id as Id<"threads">,
+      messages: [
+        { message: { role: "user", content: "hello" } },
+        { message: { role: "assistant", content: "step 1" } },
+        { message: { role: "user", content: "hello2" } },
+      ],
+    });
+    const maxMessage = await t.run(async (ctx) => {
+      return await getMaxMessage(ctx, thread._id as Id<"threads">, 0);
+    });
+    expect(maxMessage).toMatchObject({
+      _id: messages.at(1)!._id,
+      order: 0,
+      stepOrder: 1,
+    });
+  });
+
   test("getMaxMessages works when there are tools involved", async () => {
     const t = convexTest(schema, modules);
     const thread = await t.mutation(api.threads.createThread, {
@@ -78,7 +101,7 @@ describe("agent", () => {
     });
   });
 
-  test("ordering is incremented on subsequent calls to addMessages", async () => {
+  test("ordering is incremented on subsequent calls to addMessages for user messages", async () => {
     const t = convexTest(schema, modules);
     const thread = await t.mutation(api.threads.createThread, {
       userId: "test",
@@ -109,7 +132,7 @@ describe("agent", () => {
     });
   });
 
-  test("sub order is incremented on subsequent calls to addMessages for the same promptMessageId", async () => {
+  test("ordering is incremented on subsequent calls to addMessages for assistant messages", async () => {
     const t = convexTest(schema, modules);
     const thread = await t.mutation(api.threads.createThread, {
       userId: "test",
@@ -127,8 +150,72 @@ describe("agent", () => {
       stepOrder: 0,
     });
     const { messages: messages2 } = await t.mutation(api.messages.addMessages, {
+      threadId: thread._id as Id<"threads">,
+      messages: [{ message: { role: "assistant", content: "hello" } }],
+    });
+    const maxMessage2 = await t.run(async (ctx) => {
+      return await getMaxMessage(ctx, thread._id as Id<"threads">);
+    });
+    expect(maxMessage2).toMatchObject({
+      _id: messages2.at(-1)!._id,
+      order: 0,
+      stepOrder: 1,
+    });
+  });
+
+  test("order is incremented for user messages on to addMessages for the same promptMessageId", async () => {
+    const t = convexTest(schema, modules);
+    const thread = await t.mutation(api.threads.createThread, {
+      userId: "test",
+    });
+    const { messages } = await t.mutation(api.messages.addMessages, {
       threadId: thread._id as Id<"threads">,
       messages: [{ message: { role: "user", content: "hello" } }],
+    });
+    const maxMessage = await t.run(async (ctx) => {
+      return await getMaxMessage(ctx, thread._id as Id<"threads">);
+    });
+    expect(maxMessage).toMatchObject({
+      _id: messages.at(-1)!._id,
+      order: 0,
+      stepOrder: 0,
+    });
+    const { messages: messages2 } = await t.mutation(api.messages.addMessages, {
+      threadId: thread._id as Id<"threads">,
+      messages: [{ message: { role: "user", content: "hello" } }],
+      agentName: "test",
+      promptMessageId: messages.at(-1)!._id as Id<"messages">,
+    });
+    const maxMessage2 = await t.run(async (ctx) => {
+      return await getMaxMessage(ctx, thread._id as Id<"threads">);
+    });
+    expect(maxMessage2).toMatchObject({
+      _id: messages2.at(-1)!._id,
+      order: 1,
+      stepOrder: 0,
+    });
+  });
+
+  test("sub order is incremented on subsequent calls to addMessages for the same promptMessageId", async () => {
+    const t = convexTest(schema, modules);
+    const thread = await t.mutation(api.threads.createThread, {
+      userId: "test",
+    });
+    const { messages } = await t.mutation(api.messages.addMessages, {
+      threadId: thread._id as Id<"threads">,
+      messages: [{ message: { role: "user", content: "hello" } }],
+    });
+    const maxMessage = await t.run(async (ctx) => {
+      return await getMaxMessage(ctx, thread._id as Id<"threads">);
+    });
+    expect(maxMessage).toMatchObject({
+      _id: messages.at(-1)!._id,
+      order: 0,
+      stepOrder: 0,
+    });
+    const { messages: messages2 } = await t.mutation(api.messages.addMessages, {
+      threadId: thread._id as Id<"threads">,
+      messages: [{ message: { role: "assistant", content: "hello" } }],
       agentName: "test",
       promptMessageId: messages.at(-1)!._id as Id<"messages">,
     });
@@ -142,6 +229,83 @@ describe("agent", () => {
     });
   });
 
+  test("adding multiple messages at a promptMessageId skips later messages", async () => {
+    const t = convexTest(schema, modules);
+    const thread = await t.mutation(api.threads.createThread, {
+      userId: "test",
+    });
+    const { messages } = await t.mutation(api.messages.addMessages, {
+      threadId: thread._id as Id<"threads">,
+      messages: [{ message: { role: "user", content: "hello" } }],
+    });
+
+    const { messages: messages2 } = await t.mutation(api.messages.addMessages, {
+      threadId: thread._id as Id<"threads">,
+      messages: [
+        { message: { role: "user", content: "hello2" } },
+        { message: { role: "assistant", content: "hello" } },
+      ],
+      agentName: "test",
+    });
+    expect(messages2.length).toBe(2);
+
+    const { messages: messages3 } = await t.mutation(api.messages.addMessages, {
+      threadId: thread._id as Id<"threads">,
+      messages: [
+        {
+          message: {
+            role: "assistant",
+            content: [
+              {
+                type: "tool-call",
+                args: { a: 1 },
+                toolCallId: "1",
+                toolName: "tool",
+              },
+            ],
+          },
+        },
+        {
+          message: {
+            role: "tool",
+            content: [
+              {
+                type: "tool-result",
+                toolName: "tool",
+                result: "foo",
+                toolCallId: "1",
+              },
+            ],
+          },
+        },
+        { message: { role: "user", content: "bye" } },
+      ],
+      agentName: "test",
+      promptMessageId: messages.at(-1)!._id as Id<"messages">,
+    });
+
+    expect(messages3.length).toBe(3);
+
+    const allMessages = await t.query(api.messages.listMessagesByThreadId, {
+      threadId: thread._id as Id<"threads">,
+      order: "asc",
+    });
+    expect(allMessages.page).toHaveLength(6);
+    expect(allMessages.page.map((m) => m.order)).toEqual([0, 0, 0, 1, 1, 2]);
+    expect(allMessages.page.map((m) => m.stepOrder)).toEqual([
+      0, 1, 2, 0, 1, 0,
+    ]);
+    expect(allMessages.page[0]!.message!.role).toBe("user");
+    expect(allMessages.page[0]!.message!.content).toBe("hello");
+    expect(allMessages.page[1]!.message!.role).toBe("assistant");
+    expect(allMessages.page[2]!.message!.role).toBe("tool");
+    expect(allMessages.page[3]!.message!.role).toBe("user");
+    expect(allMessages.page[3]!.message!.content).toBe("hello2");
+    expect(allMessages.page[4]!.message!.role).toBe("assistant");
+    expect(allMessages.page[5]!.message!.role).toBe("user");
+    expect(allMessages.page[5]!.message!.content).toBe("bye");
+  });
+
   test("updateMessage updates message content", async () => {
     const t = convexTest(schema, modules);
     const thread = await t.mutation(api.threads.createThread, {
@@ -155,9 +319,7 @@ describe("agent", () => {
 
     const updatedMessage = await t.mutation(api.messages.updateMessage, {
       messageId,
-      patch: {
-        message: { role: "user", content: "updated content" },
-      },
+      patch: { message: { role: "user", content: "updated content" } },
     });
 
     expect(updatedMessage.message).toEqual({
@@ -173,8 +335,9 @@ describe("agent", () => {
     });
     const { messages } = await t.mutation(api.messages.addMessages, {
       threadId: thread._id as Id<"threads">,
-      messages: [{ message: { role: "assistant", content: "hello" } }],
-      pending: true,
+      messages: [
+        { message: { role: "assistant", content: "hello" }, status: "pending" },
+      ],
     });
     const messageId = messages[0]._id as Id<"messages">;
 
@@ -184,9 +347,7 @@ describe("agent", () => {
     // Update to success
     const updatedMessage = await t.mutation(api.messages.updateMessage, {
       messageId,
-      patch: {
-        status: "success",
-      },
+      patch: { status: "success" },
     });
 
     expect(updatedMessage.status).toBe("success");
@@ -199,17 +360,15 @@ describe("agent", () => {
     });
     const { messages } = await t.mutation(api.messages.addMessages, {
       threadId: thread._id as Id<"threads">,
-      messages: [{ message: { role: "assistant", content: "hello" } }],
-      pending: true,
+      messages: [
+        { message: { role: "assistant", content: "hello" }, status: "pending" },
+      ],
     });
     const messageId = messages[0]._id as Id<"messages">;
 
     const updatedMessage = await t.mutation(api.messages.updateMessage, {
       messageId,
-      patch: {
-        status: "failed",
-        error: "Something went wrong",
-      },
+      patch: { status: "failed", error: "Something went wrong" },
     });
 
     expect(updatedMessage.status).toBe("failed");
@@ -277,9 +436,7 @@ describe("agent", () => {
     await expect(
       t.mutation(api.messages.updateMessage, {
         messageId: "invalidId" as Id<"messages">,
-        patch: {
-          message: { role: "user", content: "test" },
-        },
+        patch: { message: { role: "user", content: "test" } },
       }),
     ).rejects.toThrow();
   });
@@ -308,10 +465,7 @@ describe("agent", () => {
     // Verify messages are actually deleted
     const remainingMessages = await t.query(
       api.messages.listMessagesByThreadId,
-      {
-        threadId: thread._id as Id<"threads">,
-        order: "asc",
-      },
+      { threadId: thread._id as Id<"threads">, order: "asc" },
     );
     expect(remainingMessages.page).toHaveLength(1);
     expect(remainingMessages.page[0]._id).toBe(messageIds[1]);
@@ -343,10 +497,7 @@ describe("agent", () => {
     // Verify the valid message was deleted
     const remainingMessages = await t.query(
       api.messages.listMessagesByThreadId,
-      {
-        threadId: thread._id as Id<"threads">,
-        order: "asc",
-      },
+      { threadId: thread._id as Id<"threads">, order: "asc" },
     );
     expect(remainingMessages.page).toHaveLength(0);
   });
@@ -395,10 +546,7 @@ describe("agent", () => {
     // Verify only messages from order 0 were deleted
     const remainingMessages = await t.query(
       api.messages.listMessagesByThreadId,
-      {
-        threadId: thread._id as Id<"threads">,
-        order: "asc",
-      },
+      { threadId: thread._id as Id<"threads">, order: "asc" },
     );
 
     expect(remainingMessages.page).toHaveLength(3); // Should have messages from order 1 and 2
@@ -417,7 +565,7 @@ describe("agent", () => {
       messages: [
         { message: { role: "user", content: "step 0" } },
         { message: { role: "assistant", content: "step 1" } },
-        { message: { role: "user", content: "step 2" } },
+        { message: { role: "assistant", content: "step 2" } },
         { message: { role: "assistant", content: "step 3" } },
       ],
     });
@@ -438,10 +586,7 @@ describe("agent", () => {
     // Verify only step 1 and 2 were deleted (step 3 is excluded by upperBoundInclusive: false)
     const remainingMessages = await t.query(
       api.messages.listMessagesByThreadId,
-      {
-        threadId: thread._id as Id<"threads">,
-        order: "asc",
-      },
+      { threadId: thread._id as Id<"threads">, order: "asc" },
     );
 
     expect(remainingMessages.page).toHaveLength(2);
@@ -501,10 +646,7 @@ describe("agent", () => {
     // Verify original message is still there
     const remainingMessages = await t.query(
       api.messages.listMessagesByThreadId,
-      {
-        threadId: thread._id as Id<"threads">,
-        order: "asc",
-      },
+      { threadId: thread._id as Id<"threads">, order: "asc" },
     );
     expect(remainingMessages.page).toHaveLength(1);
   });

package/src/component/messages.ts

@@ -9,7 +9,7 @@ import {
   isTool,
 } from "../shared.js";
 import {
-  vMessageEmbeddings,
+  vMessageEmbeddingsWithDimension,
   vMessageStatus,
   vMessageWithMetadataInternal,
   vPaginationResult,
@@ -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,13 @@ 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()),
 };
 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,
@@ -149,14 +144,8 @@ async function addMessagesHandler(
     assert(thread, `Thread ${args.threadId} not found`);
     userId = thread.userId;
   }
-  const {
-    embeddings,
-    failPendingSteps,
-    pending,
-    messages,
-    promptMessageId,
-    ...rest
-  } = args;
+  const { embeddings, failPendingSteps, messages, promptMessageId, ...rest } =
+    args;
   const parentMessage = promptMessageId && (await ctx.db.get(promptMessageId));
   if (failPendingSteps) {
     assert(args.threadId, "threadId is required to fail pending steps");
@@ -187,8 +176,8 @@ async function addMessagesHandler(
     stepOrder = maxMessage?.stepOrder ?? parentMessage.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) {
@@ -209,7 +198,18 @@ async function addMessagesHandler(
         threadId,
       });
     }
-    stepOrder++;
+    if (message.message.role === "user") {
+      if (parentMessage && parentMessage.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 {
+      stepOrder++;
+    }
     const messageId = await ctx.db.insert("messages", {
       ...rest,
       ...message,
@@ -219,16 +219,12 @@ async function addMessagesHandler(
       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
+        ? (parentMessage?.error ?? "Parent message failed")
+        : undefined,
       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 +259,7 @@ function orderedMessagesStream(
               .eq("threadId", threadId)
               .eq("status", status)
               .eq("tool", tool);
-            if (order) {
+            if (order !== undefined) {
               return qq.eq("order", order);
             }
             return qq;
@@ -435,7 +431,9 @@ export const getMessagesByIds = query({
     messageIds: v.array(v.id("messages")),
   },
   handler: async (ctx, args) => {
-    return await Promise.all(args.messageIds.map((id) => ctx.db.get(id)));
+    return (await Promise.all(args.messageIds.map((id) => ctx.db.get(id)))).map(
+      (m) => (m ? publicMessage(m) : null),
+    );
   },
   returns: v.array(v.union(v.null(), vMessageDoc)),
 });

package/src/component/schema.ts

@@ -6,7 +6,7 @@ import {
   vMessageStatus,
   vUsage,
   vSource,
-  vLanguageModelV1CallWarning,
+  vLanguageModelCallWarning,
   vFinishReason,
   vProviderOptions,
   vProviderMetadata,
@@ -29,12 +29,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(),
@@ -62,9 +58,10 @@ export const schema = defineSchema({
     sources: v.optional(v.array(vSource)),
     reasoning: v.optional(v.string()),
     reasoningDetails: v.optional(vReasoningDetails),
-    warnings: v.optional(v.array(vLanguageModelV1CallWarning)),
+    warnings: v.optional(v.array(vLanguageModelCallWarning)),
     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
@@ -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);