@convex-dev/agent 0.7.1 → 0.7.3

package/src/component/messages.test.ts

@@ -164,6 +164,249 @@ describe("agent", () => {
     });
   });
 
+  test("an explicit order starts, appends, and avoids later order collisions", async () => {
+    const t = convexTest(schema, modules);
+    const thread = await t.mutation(api.threads.createThread, {
+      userId: "test",
+    });
+    await t.mutation(api.messages.addMessages, {
+      threadId: thread._id as Id<"threads">,
+      messages: [
+        { message: { role: "user", content: "hello" } },
+        { message: { role: "assistant", content: "agent reply" } },
+      ],
+    });
+
+    const { messages: firstHumanReply } = await t.mutation(
+      api.messages.addMessages,
+      {
+        threadId: thread._id as Id<"threads">,
+        order: 1,
+        agentName: "human:Alex",
+        messages: [{ message: { role: "assistant", content: "human reply" } }],
+      },
+    );
+    expect(firstHumanReply[0]).toMatchObject({
+      order: 1,
+      stepOrder: 0,
+      agentName: "human:Alex",
+    });
+
+    const { messages: secondHumanReply } = await t.mutation(
+      api.messages.addMessages,
+      {
+        threadId: thread._id as Id<"threads">,
+        order: 1,
+        agentName: "human:Sam",
+        messages: [{ message: { role: "assistant", content: "follow-up" } }],
+      },
+    );
+    expect(secondHumanReply[0]).toMatchObject({
+      order: 1,
+      stepOrder: 1,
+      agentName: "human:Sam",
+    });
+
+    const { messages: backdatedBatch } = await t.mutation(
+      api.messages.addMessages,
+      {
+        threadId: thread._id as Id<"threads">,
+        order: 0,
+        messages: [
+          { message: { role: "assistant", content: "backdated reply" } },
+          { message: { role: "user", content: "new user turn" } },
+        ],
+      },
+    );
+    expect(
+      backdatedBatch.map(({ order, stepOrder }) => [order, stepOrder]),
+    ).toEqual([
+      [0, 2],
+      [2, 0],
+    ]);
+  });
+
+  test("concurrent saves to an explicit order receive distinct step orders", async () => {
+    const t = convexTest(schema, modules);
+    const thread = await t.mutation(api.threads.createThread, {
+      userId: "test",
+    });
+    const saveReply = (agentName: string) =>
+      t.mutation(api.messages.addMessages, {
+        threadId: thread._id as Id<"threads">,
+        order: 3,
+        agentName,
+        messages: [
+          { message: { role: "assistant" as const, content: agentName } },
+        ],
+      });
+
+    const replies = await Promise.all([saveReply("Alex"), saveReply("Sam")]);
+
+    expect(
+      replies
+        .map(({ messages }) => messages[0].stepOrder)
+        .sort((a, b) => a - b),
+    ).toEqual([0, 1]);
+  });
+
+  test("next order is allocated after the latest message", async () => {
+    const t = convexTest(schema, modules);
+    const thread = await t.mutation(api.threads.createThread, {
+      userId: "test",
+    });
+    await t.mutation(api.messages.addMessages, {
+      threadId: thread._id as Id<"threads">,
+      messages: [{ message: { role: "user", content: "hello" } }],
+    });
+
+    const { messages } = await t.mutation(api.messages.addMessages, {
+      threadId: thread._id as Id<"threads">,
+      order: "next",
+      messages: [{ message: { role: "assistant", content: "separate reply" } }],
+    });
+
+    expect(messages[0]).toMatchObject({ order: 1, stepOrder: 0 });
+  });
+
+  test("concurrent next orders receive distinct orders", async () => {
+    const t = convexTest(schema, modules);
+    const thread = await t.mutation(api.threads.createThread, {
+      userId: "test",
+    });
+    await t.mutation(api.messages.addMessages, {
+      threadId: thread._id as Id<"threads">,
+      messages: [{ message: { role: "user", content: "hello" } }],
+    });
+    const saveReply = (agentName: string) =>
+      t.mutation(api.messages.addMessages, {
+        threadId: thread._id as Id<"threads">,
+        order: "next",
+        agentName,
+        messages: [
+          { message: { role: "assistant" as const, content: agentName } },
+        ],
+      });
+
+    const replies = await Promise.all([saveReply("Alex"), saveReply("Sam")]);
+
+    expect(
+      replies
+        .map(({ messages }) => messages[0])
+        .sort((a, b) => a.order - b.order)
+        .map(({ order, stepOrder }) => [order, stepOrder]),
+    ).toEqual([
+      [1, 0],
+      [2, 0],
+    ]);
+  });
+
+  test("an explicit order ahead of the thread places the batch there", async () => {
+    const t = convexTest(schema, modules);
+    const thread = await t.mutation(api.threads.createThread, {
+      userId: "test",
+    });
+    await t.mutation(api.messages.addMessages, {
+      threadId: thread._id as Id<"threads">,
+      messages: [{ message: { role: "user", content: "hello" } }],
+    });
+
+    const { messages } = await t.mutation(api.messages.addMessages, {
+      threadId: thread._id as Id<"threads">,
+      order: 5,
+      messages: [
+        { message: { role: "assistant", content: "imported reply" } },
+        { message: { role: "user", content: "imported question" } },
+      ],
+    });
+    expect(messages.map(({ order, stepOrder }) => [order, stepOrder])).toEqual([
+      [5, 0],
+      [6, 0],
+    ]);
+  });
+
+  test("an explicit order cannot conflict with another placement argument", 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" } }],
+    });
+
+    await expect(
+      t.mutation(api.messages.addMessages, {
+        threadId: thread._id as Id<"threads">,
+        order: 1,
+        promptMessageId: messages[0]._id as Id<"messages">,
+        messages: [{ message: { role: "assistant", content: "reply" } }],
+      }),
+    ).rejects.toThrow("order and promptMessageId cannot both be provided");
+
+    await expect(
+      t.mutation(api.messages.addMessages, {
+        threadId: thread._id as Id<"threads">,
+        order: "next",
+        promptMessageId: messages[0]._id as Id<"messages">,
+        messages: [{ message: { role: "assistant", content: "reply" } }],
+      }),
+    ).rejects.toThrow("order and promptMessageId cannot both be provided");
+
+    for (const order of [-1, 1.5, Number.MAX_SAFE_INTEGER]) {
+      await expect(
+        t.mutation(api.messages.addMessages, {
+          threadId: thread._id as Id<"threads">,
+          order,
+          messages: [{ message: { role: "assistant", content: "reply" } }],
+        }),
+      ).rejects.toThrow("order must be a non-negative safe integer");
+    }
+  });
+
+  test("derived message positions cannot exceed safe integers", 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">,
+      order: 1,
+      messages: [{ message: { role: "assistant", content: "reply" } }],
+    });
+    const messageId = messages[0]._id as Id<"messages">;
+
+    await t.run(async (ctx) => {
+      await ctx.db.patch("messages", messageId, {
+        stepOrder: Number.MAX_SAFE_INTEGER,
+      });
+    });
+    await expect(
+      t.mutation(api.messages.addMessages, {
+        threadId: thread._id as Id<"threads">,
+        order: 1,
+        messages: [{ message: { role: "assistant", content: "follow-up" } }],
+      }),
+    ).rejects.toThrow(
+      "stepOrder cannot be incremented past Number.MAX_SAFE_INTEGER",
+    );
+
+    await t.run(async (ctx) => {
+      await ctx.db.patch("messages", messageId, {
+        order: Number.MAX_SAFE_INTEGER,
+        stepOrder: 0,
+      });
+    });
+    await expect(
+      t.mutation(api.messages.addMessages, {
+        threadId: thread._id as Id<"threads">,
+        messages: [{ message: { role: "user", content: "new turn" } }],
+      }),
+    ).rejects.toThrow(
+      "order cannot be incremented past Number.MAX_SAFE_INTEGER",
+    );
+  });
+
   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, {
@@ -1017,3 +1260,164 @@ describe("agent", () => {
     });
   });
 });
+
+describe("late saves racing a failed pending message (issue #320)", () => {
+  const PROVIDER_ERROR = "invalid_prompt: Invalid prompt: flagged by policy.";
+
+  test("keeps the first durable failure authoritative", async () => {
+    const t = initConvexTest();
+    const thread = await t.mutation(api.threads.createThread, {
+      userId: "u1",
+    });
+    const threadId = thread._id as Id<"threads">;
+
+    const { messages: seeded } = await t.mutation(api.messages.addMessages, {
+      threadId,
+      messages: [
+        { message: { role: "user", content: "hello" } },
+        { message: { role: "assistant", content: [] }, status: "pending" },
+      ],
+    });
+    const pending = seeded.at(-1)!;
+    expect(pending.status).toBe("pending");
+
+    const streamId = await t.mutation(api.streams.create, {
+      threadId,
+      order: pending.order,
+      stepOrder: pending.stepOrder,
+      format: "UIMessageChunk",
+    });
+
+    await t.mutation(api.messages.finalizeMessage, {
+      messageId: pending._id as Id<"messages">,
+      result: { status: "failed", error: PROVIDER_ERROR },
+    });
+    await t.mutation(api.streams.abort, { streamId, reason: PROVIDER_ERROR });
+
+    const { messages: late } = await t.mutation(api.messages.addMessages, {
+      threadId,
+      pendingMessageId: pending._id as Id<"messages">,
+      finishStreamId: streamId,
+      failPendingSteps: false,
+      messages: [
+        { message: { role: "assistant", content: "partial response" } },
+      ],
+    });
+
+    const assistants = (
+      await t.run(async (ctx) =>
+        ctx.db
+          .query("messages")
+          .withIndex("threadId_status_tool_order_stepOrder", (q) =>
+            q.eq("threadId", threadId),
+          )
+          .collect(),
+      )
+    ).filter((message) => message.message?.role === "assistant");
+
+    expect(late).toHaveLength(1);
+    expect(assistants).toHaveLength(1);
+    expect(assistants[0]!._id).toBe(pending._id);
+    expect(assistants[0]!.status).toBe("failed");
+    expect(assistants[0]!.error).toBe(PROVIDER_ERROR);
+    expect(assistants[0]!.text).toBe("partial response");
+
+    const stream = await t.run((ctx) =>
+      ctx.db.get("streamingMessages", streamId),
+    );
+    expect(stream?.state.kind).toBe("aborted");
+  });
+});
+
+describe("deleting a message aborts generation writing to it (issue #300)", () => {
+  test("a stream at the deleted order is aborted", async () => {
+    const t = initConvexTest();
+    const thread = await t.mutation(api.threads.createThread, { userId: "u" });
+    const threadId = thread._id as Id<"threads">;
+
+    const { messages } = await t.mutation(api.messages.addMessages, {
+      threadId,
+      messages: [{ message: { role: "user", content: "hello" } }],
+    });
+    const prompt = messages[0];
+
+    await t.mutation(api.streams.create, {
+      threadId,
+      order: prompt.order,
+      stepOrder: prompt.stepOrder + 1,
+      userId: "u",
+      agentName: "a",
+      model: "m",
+      provider: "p",
+      format: "UIMessageChunk",
+    });
+
+    await t.mutation(api.messages.deleteByIds, {
+      messageIds: [prompt._id as Id<"messages">],
+    });
+
+    const streaming = await t.query(api.streams.list, {
+      threadId,
+      statuses: ["streaming"],
+    });
+    const aborted = await t.query(api.streams.list, {
+      threadId,
+      statuses: ["aborted"],
+    });
+    expect(streaming).toHaveLength(0);
+    expect(aborted).toHaveLength(1);
+  });
+});
+
+describe("abandoning a save whose prompt was deleted (issue #300)", () => {
+  test("abandons instead of throwing when the caller opts in", async () => {
+    const t = initConvexTest();
+    const thread = await t.mutation(api.threads.createThread, { userId: "u" });
+    const threadId = thread._id as Id<"threads">;
+
+    const { messages } = await t.mutation(api.messages.addMessages, {
+      threadId,
+      messages: [{ message: { role: "user", content: "hello" } }],
+    });
+    const promptMessageId = messages[0]._id as Id<"messages">;
+
+    await t.mutation(api.messages.deleteByIds, { messageIds: [promptMessageId] });
+
+    const saved = await t.mutation(api.messages.addMessages, {
+      threadId,
+      promptMessageId,
+      abandonIfPromptMissing: true,
+      messages: [{ message: { role: "assistant", content: "answer" } }],
+    });
+    expect(saved.messages).toEqual([]);
+
+    // Nothing was grafted onto the thread.
+    const all = await t.query(api.messages.listMessagesByThreadId, {
+      threadId,
+      order: "asc",
+      paginationOpts: { cursor: null, numItems: 10 },
+    });
+    expect(all.page).toHaveLength(0);
+  });
+
+  test("still throws for a caller that did not opt in", async () => {
+    const t = initConvexTest();
+    const thread = await t.mutation(api.threads.createThread, { userId: "u" });
+    const threadId = thread._id as Id<"threads">;
+
+    const { messages } = await t.mutation(api.messages.addMessages, {
+      threadId,
+      messages: [{ message: { role: "user", content: "hello" } }],
+    });
+    const promptMessageId = messages[0]._id as Id<"messages">;
+    await t.mutation(api.messages.deleteByIds, { messageIds: [promptMessageId] });
+
+    await expect(
+      t.mutation(api.messages.addMessages, {
+        threadId,
+        promptMessageId,
+        messages: [{ message: { role: "assistant", content: "answer" } }],
+      }),
+    ).rejects.toThrow("not found");
+  });
+});

package/src/component/messages.ts

@@ -44,6 +44,7 @@ import {
   getStreamingMessagesWithMetadata,
   finishHandler,
   releaseStreamFileOwnershipByIds,
+  abortStreamsAtOrder,
 } from "./streams.js";
 import { partial } from "convex-helpers/validators";
 
@@ -65,21 +66,45 @@ export async function deleteMessage(
   }
 }
 
+/**
+ * Deleting a message strands any generation still writing to its order, which
+ * would otherwise only surface as a missing-parent failure when that generation
+ * finalizes. Aborting the stream lets the in-flight run stop on its own.
+ */
+async function abortStreamsForDeleted(
+  ctx: MutationCtx,
+  deleted: (Doc<"messages"> | null)[],
+) {
+  const seen = new Set<string>();
+  for (const message of deleted) {
+    if (!message) continue;
+    const key = `${message.threadId}:${message.order}`;
+    if (seen.has(key)) continue;
+    seen.add(key);
+    await abortStreamsAtOrder(ctx, {
+      threadId: message.threadId,
+      order: message.order,
+      reason: "Message deleted",
+    });
+  }
+}
+
 export const deleteByIds = mutation({
   args: { messageIds: v.array(v.id("messages")) },
   returns: v.array(v.id("messages")),
   handler: async (ctx, args) => {
-    const deletedMessageIds = await Promise.all(
+    const deleted = await Promise.all(
       args.messageIds.map(async (id) => {
         const message = await ctx.db.get("messages", id);
         if (message) {
           await deleteMessage(ctx, message);
-          return id;
+          return message;
         }
         return null;
       }),
     );
-    return deletedMessageIds.filter((id) => id !== null);
+    await abortStreamsForDeleted(ctx, deleted);
+    return deleted.filter((m) => m !== null).map((m) => m._id);
   },
 });
 
@@ -145,6 +170,7 @@ export const deleteByOrder = mutation({
       })
       .take(64);
     await Promise.all(messages.map((m) => deleteMessage(ctx, m)));
+    await abortStreamsForDeleted(ctx, messages);
     return {
       isDone: messages.length < 64,
       lastOrder: messages.at(-1)?.order,
@@ -157,6 +183,13 @@ const addMessagesArgs = {
   userId: v.optional(v.string()),
   threadId: v.id("threads"),
   promptMessageId: v.optional(v.id("messages")),
+  /**
+   * For saves that belong to a run anchored on promptMessageId: if that
+   * message is gone the run is obsolete, so abandon the save instead of
+   * throwing. A caller passing an id that never existed still gets an error.
+   */
+  abandonIfPromptMissing: v.optional(v.boolean()),
+  order: v.optional(v.union(v.number(), v.literal("next"))),
   agentName: v.optional(v.string()),
   messages: v.array(vMessageWithMetadataInternal),
   embeddings: v.optional(vMessageEmbeddingsWithDimension),
@@ -175,6 +208,15 @@ export const addMessages = mutation({
   handler: addMessagesHandler,
   returns: v.object({ messages: v.array(vMessageDoc) }),
 });
+
+function incrementMessagePosition(value: number, field: "order" | "stepOrder") {
+  assert(
+    Number.isSafeInteger(value) && value < Number.MAX_SAFE_INTEGER,
+    `${field} cannot be incremented past Number.MAX_SAFE_INTEGER`,
+  );
+  return value + 1;
+}
+
 async function addMessagesHandler(
   ctx: MutationCtx,
   args: ObjectType<typeof addMessagesArgs>,
@@ -193,12 +235,30 @@ async function addMessagesHandler(
     finishStreamId,
     messages,
     promptMessageId,
+    abandonIfPromptMissing,
+    order: requestedOrder,
     pendingMessageId,
     hideFromUserIdSearch,
     ...rest
   } = args;
   const promptMessage =
     promptMessageId && (await ctx.db.get("messages", promptMessageId));
+  assert(
+    requestedOrder === undefined ||
+      requestedOrder === "next" ||
+      (Number.isSafeInteger(requestedOrder) &&
+        requestedOrder >= 0 &&
+        requestedOrder < Number.MAX_SAFE_INTEGER),
+    "order must be a non-negative safe integer less than Number.MAX_SAFE_INTEGER",
+  );
+  assert(
+    requestedOrder === undefined || !promptMessageId,
+    "order and promptMessageId cannot both be provided",
+  );
+  assert(
+    requestedOrder === undefined || !pendingMessageId,
+    "order and pendingMessageId cannot both be provided",
+  );
   if (failPendingSteps) {
     assert(args.threadId, "threadId is required to fail pending steps");
     const pendingMessages = await ctx.db
@@ -228,7 +288,21 @@ async function addMessagesHandler(
   let order, stepOrder;
   let fail = false;
   let error: string | undefined;
-  if (promptMessageId) {
+  const startsAtNextOrder = requestedOrder === "next";
+  const explicitOrder =
+    typeof requestedOrder === "number" ? requestedOrder : undefined;
+  if (startsAtNextOrder) {
+    const maxMessage = await getMaxMessage(ctx, threadId);
+    order = incrementMessagePosition(maxMessage?.order ?? -1, "order");
+    stepOrder = -1;
+  } else if (explicitOrder !== undefined) {
+    order = explicitOrder;
+    const maxMessage = await getMaxMessage(ctx, threadId, order);
+    stepOrder = maxMessage?.stepOrder ?? -1;
+  } else if (promptMessageId) {
+    if (!promptMessage && abandonIfPromptMissing) {
+      return { messages: [] };
+    }
     assert(promptMessage, `Parent message ${promptMessageId} not found`);
     if (promptMessage.status === "failed") {
       fail = true;
@@ -289,8 +363,8 @@ async function addMessagesHandler(
       if (pendingMessage.status === "failed") {
         fail = true;
         error =
-          `Trying to update a message that failed: ${pendingMessageId}, ` +
-          `error: ${pendingMessage.error ?? error}`;
+          pendingMessage.error ??
+          `Trying to update a message that failed: ${pendingMessageId}`;
         messageDoc.status = "failed";
         messageDoc.error = error;
       }
@@ -309,20 +383,28 @@ async function addMessagesHandler(
       toReturn.push((await ctx.db.get("messages", pendingMessage._id))!);
       continue;
     }
-    if (message.message.role === "user") {
-      if (promptMessage && promptMessage.order === order) {
-        // see if there's a later message than the parent message order
+    if ((startsAtNextOrder || explicitOrder !== undefined) && i === 0) {
+      stepOrder = incrementMessagePosition(stepOrder, "stepOrder");
+    } else if (message.message.role === "user") {
+      if (
+        (explicitOrder !== undefined && order === explicitOrder) ||
+        (promptMessage && promptMessage.order === order)
+      ) {
+        // Avoid colliding with a later order when saving from an older one.
         const maxMessage = await getMaxMessage(ctx, threadId);
-        order = (maxMessage?.order ?? order) + 1;
+        order = incrementMessagePosition(
+          Math.max(maxMessage?.order ?? order, order),
+          "order",
+        );
       } else {
-        order++;
+        order = incrementMessagePosition(order, "order");
       }
       stepOrder = 0;
     } else {
       if (order < 0) {
         order = 0;
       }
-      stepOrder++;
+      stepOrder = incrementMessagePosition(stepOrder, "stepOrder");
     }
     const messageId = await ctx.db.insert("messages", {
       ...messageDoc,

package/src/component/streams.test.ts

@@ -1,10 +1,12 @@
 /// <reference types="vite/client" />
 
-import { describe, expect, test } from "vitest";
+import { afterEach, describe, expect, test, vi } from "vitest";
 import { api } from "./_generated/api.js";
 import type { Id } from "./_generated/dataModel.js";
 import { initConvexTest } from "./setup.test.js";
 
+afterEach(() => vi.useRealTimers());
+
 async function seedStream(t: ReturnType<typeof initConvexTest>) {
   const thread = await t.mutation(api.threads.createThread, {
     userId: "stream-files",
@@ -106,4 +108,55 @@ describe("streams", () => {
     expect(stream?.state.kind).toBe("finished");
     expect(stream?.fileRefs).toBeUndefined();
   });
+
+  test("bulk deletion preserves global stream order across status lanes", async () => {
+    vi.useFakeTimers();
+
+    const t = initConvexTest();
+    const thread = await t.mutation(api.threads.createThread, {
+      userId: "stream-cleanup",
+    });
+    const threadId = thread._id as Id<"threads">;
+    const earlyStreamId = await t.run(async (ctx) => {
+      const streamId = await ctx.db.insert("streamingMessages", {
+        threadId,
+        order: 1,
+        stepOrder: 0,
+        format: "UIMessageChunk",
+        state: { kind: "streaming", lastHeartbeat: Date.now() },
+      });
+      await ctx.db.insert("streamingMessages", {
+        threadId,
+        order: 9,
+        stepOrder: 0,
+        format: "UIMessageChunk",
+        state: { kind: "aborted", reason: "retained regeneration" },
+      });
+      await ctx.db.insert("streamDeltas", {
+        streamId,
+        start: 0,
+        end: 1,
+        parts: [{ type: "text-start", id: "text" }],
+      });
+      return streamId;
+    });
+
+    await t.mutation(api.streams.deleteAllStreamsForThreadIdAsync, {
+      threadId,
+    });
+    await t.finishAllScheduledFunctions(vi.runAllTimers);
+
+    expect(
+      await t.query(api.streams.list, {
+        threadId,
+        statuses: ["streaming", "finished", "aborted"],
+      }),
+    ).toEqual([]);
+    expect(
+      await t.query(api.streams.listDeltas, {
+        threadId,
+        cursors: [{ streamId: earlyStreamId, cursor: 0 }],
+      }),
+    ).toEqual([]);
+  });
 });