fix(component): extend message pages to the order boundary

src/component/messages.test.ts

@@ -1381,7 +1381,9 @@ describe("abandoning a save whose prompt was deleted (issue #300)", () => {
     });
     const promptMessageId = messages[0]._id as Id<"messages">;
 
-    await t.mutation(api.messages.deleteByIds, { messageIds: [promptMessageId] });
+    await t.mutation(api.messages.deleteByIds, {
+      messageIds: [promptMessageId],
+    });
 
     const saved = await t.mutation(api.messages.addMessages, {
       threadId,
@@ -1410,7 +1412,9 @@ describe("abandoning a save whose prompt was deleted (issue #300)", () => {
       messages: [{ message: { role: "user", content: "hello" } }],
     });
     const promptMessageId = messages[0]._id as Id<"messages">;
-    await t.mutation(api.messages.deleteByIds, { messageIds: [promptMessageId] });
+    await t.mutation(api.messages.deleteByIds, {
+      messageIds: [promptMessageId],
+    });
 
     await expect(
       t.mutation(api.messages.addMessages, {
@@ -1421,3 +1425,177 @@ describe("abandoning a save whose prompt was deleted (issue #300)", () => {
     ).rejects.toThrow("not found");
   });
 });
+
+describe("listMessagesByThreadId completes split orders", () => {
+  // Three generations. Orders 0 and 2 are tool loops (four rows each), order
+  // 1 is a plain exchange (two rows).
+  async function seedThread() {
+    const t = initConvexTest();
+    const thread = await t.mutation(api.threads.createThread, {
+      userId: "test",
+    });
+    const threadId = thread._id as Id<"threads">;
+    const toolLoop = (tag: string) => [
+      { message: { role: "user" as const, content: `${tag} prompt` } },
+      {
+        message: {
+          role: "assistant" as const,
+          content: [
+            {
+              type: "tool-call" as const,
+              toolCallId: `${tag}-call`,
+              toolName: "echo",
+              input: {},
+            },
+          ],
+        },
+      },
+      {
+        message: {
+          role: "tool" as const,
+          content: [
+            {
+              type: "tool-result" as const,
+              toolCallId: `${tag}-call`,
+              toolName: "echo",
+              output: { type: "text" as const, value: "ok" },
+            },
+          ],
+        },
+      },
+      { message: { role: "assistant" as const, content: `${tag} answer` } },
+    ];
+    await t.mutation(api.messages.addMessages, {
+      threadId,
+      messages: toolLoop("a"),
+    });
+    await t.mutation(api.messages.addMessages, {
+      threadId,
+      messages: [
+        { message: { role: "user", content: "b prompt" } },
+        { message: { role: "assistant", content: "b answer" } },
+      ],
+    });
+    await t.mutation(api.messages.addMessages, {
+      threadId,
+      messages: toolLoop("c"),
+    });
+    return { t, threadId };
+  }
+
+  const positions = (page: { order: number; stepOrder: number }[]) =>
+    page.map((m) => `${m.order}/${m.stepOrder}`);
+
+  test("a page that ends inside an order is extended to the order boundary", async () => {
+    const { t, threadId } = await seedThread();
+    const first = await t.query(api.messages.listMessagesByThreadId, {
+      threadId,
+      order: "desc",
+      paginationOpts: { cursor: null, numItems: 3 },
+    });
+    expect(positions(first.page)).toEqual(["2/3", "2/2", "2/1", "2/0"]);
+    expect(first.isDone).toBe(false);
+
+    const second = await t.query(api.messages.listMessagesByThreadId, {
+      threadId,
+      order: "desc",
+      paginationOpts: { cursor: first.continueCursor, numItems: 3 },
+    });
+    expect(positions(second.page)).toEqual([
+      "1/1",
+      "1/0",
+      "0/3",
+      "0/2",
+      "0/1",
+      "0/0",
+    ]);
+    const third = await t.query(api.messages.listMessagesByThreadId, {
+      threadId,
+      order: "desc",
+      paginationOpts: { cursor: second.continueCursor, numItems: 3 },
+    });
+    expect(third.page).toEqual([]);
+    expect(third.isDone).toBe(true);
+  });
+
+  test("a page that already ends on an order boundary is not extended", async () => {
+    const { t, threadId } = await seedThread();
+    const page = await t.query(api.messages.listMessagesByThreadId, {
+      threadId,
+      order: "desc",
+      paginationOpts: { cursor: null, numItems: 4 },
+    });
+    expect(positions(page.page)).toEqual(["2/3", "2/2", "2/1", "2/0"]);
+  });
+
+  test("the completion request is exact when a filter hides rows of the order", async () => {
+    const { t, threadId } = await seedThread();
+    // Tool rows are the tool result and the assistant tool call, so with them
+    // excluded the stream for order 2 is 2/3, 2/0. A page of one row ends at
+    // 2/3, whose stepOrder would suggest three remaining rows, but only one
+    // exists in the filtered stream. Fetching three would pull 1/1 and 1/0
+    // into the page and the cursor would skip past them.
+    const first = await t.query(api.messages.listMessagesByThreadId, {
+      threadId,
+      order: "desc",
+      excludeToolMessages: true,
+      paginationOpts: { cursor: null, numItems: 1 },
+    });
+    expect(positions(first.page)).toEqual(["2/3", "2/0"]);
+
+    const second = await t.query(api.messages.listMessagesByThreadId, {
+      threadId,
+      order: "desc",
+      excludeToolMessages: true,
+      paginationOpts: { cursor: first.continueCursor, numItems: 1 },
+    });
+    expect(positions(second.page)).toEqual(["1/1", "1/0"]);
+  });
+
+  test("re-reading a page with its endCursor returns the same extended range", async () => {
+    const { t, threadId } = await seedThread();
+    const first = await t.query(api.messages.listMessagesByThreadId, {
+      threadId,
+      order: "desc",
+      paginationOpts: { cursor: null, numItems: 3 },
+    });
+    const reread = await t.query(api.messages.listMessagesByThreadId, {
+      threadId,
+      order: "desc",
+      paginationOpts: {
+        cursor: null,
+        numItems: 3,
+        endCursor: first.continueCursor,
+      },
+    });
+    expect(reread.page.map((m) => m._id)).toEqual(first.page.map((m) => m._id));
+    expect(reread.continueCursor).toBe(first.continueCursor);
+  });
+
+  test("completeOrders: false returns exactly the requested rows", async () => {
+    const { t, threadId } = await seedThread();
+    const page = await t.query(api.messages.listMessagesByThreadId, {
+      threadId,
+      order: "desc",
+      completeOrders: false,
+      paginationOpts: { cursor: null, numItems: 1 },
+    });
+    expect(positions(page.page)).toEqual(["2/3"]);
+  });
+
+  test("a caller-supplied read limit is honored instead of extending", async () => {
+    const { t, threadId } = await seedThread();
+    const rows = await t.query(api.messages.listMessagesByThreadId, {
+      threadId,
+      order: "desc",
+      paginationOpts: { cursor: null, numItems: 3, maximumRowsRead: 1 },
+    });
+    expect(rows.page.length).toBeLessThanOrEqual(1);
+    const bytes = await t.query(api.messages.listMessagesByThreadId, {
+      threadId,
+      order: "desc",
+      paginationOpts: { cursor: null, numItems: 3, maximumBytesRead: 1 },
+    });
+    expect(bytes.page.length).toBeLessThanOrEqual(1);
+  });
+});

src/component/messages.ts

@@ -618,6 +618,7 @@ export const cloneMessageBatch = internalMutation({
       paginationOpts: args.paginationOpts,
       statuses: args.statuses,
       upToAndIncludingMessageId: args.upToAndIncludingMessageId,
+      completeOrders: false,
     });
 
     const existing =
@@ -736,6 +737,12 @@ export const listMessagesByThreadIdArgs = {
   paginationOpts: v.optional(paginationOptsValidator),
   statuses: v.optional(v.array(vMessageStatus)),
   upToAndIncludingMessageId: v.optional(v.id("messages")),
+  /**
+   * Extend each descending page to the end of its oldest order so a UI
+   * message is never split across pages. Defaults to true. Pass false when
+   * only the newest rows matter and the page size must be exact.
+   */
+  completeOrders: v.optional(v.boolean()),
 };
 export const listMessagesByThreadId = query({
   args: listMessagesByThreadIdArgs,
@@ -781,18 +788,67 @@ async function listMessagesByThreadIdHandler(
         ),
     ),
   );
+  const paginationOpts = args.paginationOpts ?? {
+    numItems: DEFAULT_RECENT_MESSAGES,
+    cursor: null,
+  };
   const messages = await mergedStream(streams, ["order", "stepOrder"]).paginate(
-    args.paginationOpts ?? {
-      numItems: DEFAULT_RECENT_MESSAGES,
-      cursor: null,
-    },
+    paginationOpts,
   );
   if (messages.page.length === 0) {
     messages.isDone = true;
+    return messages;
+  }
+  // A UI message is assembled from every row sharing one order, so a page
+  // that ends partway through an order leaves the client with a fragment it
+  // cannot render or reconcile with its streaming counterpart. Extend the
+  // page to the order boundary so the continue cursor never splits an order.
+  // Only "desc" is handled: the remaining rows of the oldest order are the
+  // ones with a lower stepOrder, a bounded set. With endCursor the range is
+  // fixed by the caller, and a caller-supplied row or byte limit, or a page
+  // the paginator already asked to split, means the caller is managing
+  // transaction budget and the page must not grow. An order with more than
+  // MAX_ORDER_COMPLETION_MESSAGES remaining rows is completed across
+  // subsequent pages instead.
+  const oldest = messages.page[messages.page.length - 1];
+  if (
+    order === "desc" &&
+    args.completeOrders !== false &&
+    !messages.isDone &&
+    !messages.pageStatus &&
+    !paginationOpts.endCursor &&
+    paginationOpts.maximumRowsRead === undefined &&
+    paginationOpts.maximumBytesRead === undefined &&
+    oldest.stepOrder > 0
+  ) {
+    const paginateFrom = (numItems: number) =>
+      mergedStream(streams, ["order", "stepOrder"]).paginate({
+        numItems,
+        cursor: messages.continueCursor,
+      });
+    let completion = await paginateFrom(
+      Math.min(oldest.stepOrder, MAX_ORDER_COMPLETION_MESSAGES),
+    );
+    // stepOrder overestimates when a filter or deletion removed rows from the
+    // order. The stream is ordered, so the order's rows are a prefix of the
+    // page; re-read exactly that prefix so the cursor stops at the boundary
+    // rather than past rows the caller never received.
+    const kept = completion.page.findIndex((m) => m.order !== oldest.order);
+    if (kept === 0) {
+      return messages;
+    }
+    if (kept > 0) {
+      completion = await paginateFrom(kept);
+    }
+    messages.page.push(...completion.page);
+    messages.continueCursor = completion.continueCursor;
+    messages.isDone = completion.isDone;
   }
   return messages;
 }
 
+const MAX_ORDER_COMPLETION_MESSAGES = 1_000;
+
 export const getMessagesByIds = query({
   args: { messageIds: v.array(v.id("messages")) },
   handler: async (ctx, args) => {

src/vercel/client/definePlaygroundAPI.ts

@@ -171,6 +171,7 @@ export function definePlaygroundAPI<DataModel extends GenericDataModel>(
               threadId: thread._id,
               order: "desc",
               paginationOpts: { numItems: 1, cursor: null },
+              completeOrders: false,
             });
             return {
               ...thread,