@convex-dev/agent 0.7.0 → 0.7.2

package/src/component/messages.test.ts

@@ -5,6 +5,7 @@ import { describe, expect, test, vi } from "vitest";
 import { api } from "./_generated/api.js";
 import type { Id } from "./_generated/dataModel.js";
 import { getMaxMessage } from "./messages.js";
+import { timeoutStreamHandler } from "./streams.js";
 import schema from "./schema.js";
 import { initConvexTest, modules } from "./setup.test.js";
 
@@ -163,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, {
@@ -924,4 +1168,256 @@ describe("agent", () => {
     expect(state.messages[0].embeddingId).toBeUndefined();
     expect(state.embedding).toBeNull();
   });
+
+  test("transfers recovered file ownership and releases timed-out streams", async () => {
+    const t = initConvexTest();
+    const thread = await t.mutation(api.threads.createThread, {
+      userId: "stream-file-recovery",
+    });
+    const { messages } = await t.mutation(api.messages.addMessages, {
+      threadId: thread._id as Id<"threads">,
+      messages: [
+        { message: { role: "assistant", content: [] }, status: "pending" },
+      ],
+    });
+    const pending = messages[0]!;
+    const { fileId: recoveredFileId } = await t.mutation(api.files.addFile, {
+      storageId: "recovered-storage",
+      hash: "recovered-hash",
+      filename: "recovered.txt",
+    });
+    await t.mutation(api.files.copyFile, { fileId: recoveredFileId });
+    const recoveredStreamId = await t.run((ctx) =>
+      ctx.db.insert("streamingMessages", {
+        threadId: thread._id as Id<"threads">,
+        order: pending.order,
+        stepOrder: pending.stepOrder,
+        format: "UIMessageChunk",
+        state: { kind: "aborted", reason: "interrupted" },
+        fileRefs: [
+          { url: "https://files.example/recovered", fileId: recoveredFileId },
+        ],
+      }),
+    );
+    await t.run((ctx) =>
+      ctx.db.insert("streamDeltas", {
+        streamId: recoveredStreamId,
+        start: 0,
+        end: 3,
+        parts: [
+          { type: "start" },
+          {
+            type: "file",
+            url: "https://files.example/recovered",
+            mediaType: "text/plain",
+          },
+          { type: "finish" },
+        ],
+      }),
+    );
+    await t.mutation(api.messages.finalizeMessage, {
+      messageId: pending._id as Id<"messages">,
+      result: { status: "success" },
+    });
+    await expect(
+      t.query(api.files.get, { fileId: recoveredFileId }),
+    ).resolves.toMatchObject({
+      refcount: 1,
+    });
+    expect(
+      (await t.run((ctx) => ctx.db.get("streamingMessages", recoveredStreamId)))
+        ?.fileRefs,
+    ).toBeUndefined();
+
+    const { fileId: timedOutFileId } = await t.mutation(api.files.addFile, {
+      storageId: "timeout-storage",
+      hash: "timeout-hash",
+      filename: "timeout.txt",
+    });
+    await t.mutation(api.files.copyFile, { fileId: timedOutFileId });
+    const timedOutStreamId = await t.run((ctx) =>
+      ctx.db.insert("streamingMessages", {
+        threadId: thread._id as Id<"threads">,
+        order: pending.order + 1,
+        stepOrder: 0,
+        format: "UIMessageChunk",
+        state: { kind: "streaming", lastHeartbeat: Date.now() },
+        fileRefs: [
+          { url: "https://files.example/timeout", fileId: timedOutFileId },
+        ],
+      }),
+    );
+    await t.run((ctx) =>
+      timeoutStreamHandler(ctx, { streamId: timedOutStreamId }),
+    );
+    await t.mutation(api.streams.deleteStreamSync, {
+      streamId: timedOutStreamId,
+    });
+    await expect(
+      t.query(api.files.get, { fileId: timedOutFileId }),
+    ).resolves.toMatchObject({
+      refcount: 0,
+    });
+  });
+});
+
+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");
+  });
 });