From 18b98d64c3991b2ee1e14a65fc596d84d89275d5 Mon Sep 17 00:00:00 2001 From: Robel Estifanos Date: Fri, 18 Sep 2026 09:26:46 -0700 Subject: [PATCH 1/2] test: paginated message pages end on order boundaries Claude-Session: https://claude.ai/code/session_01YGMZR56ff1JhEEoMnuwa9e --- src/component/messages.test.ts | 182 ++++++++++++++++++++++++++++- src/vercel/client/search.test.ts | 192 ++++++++++++------------------- 2 files changed, 255 insertions(+), 119 deletions(-) diff --git a/src/component/messages.test.ts b/src/component/messages.test.ts index ce380eb8..68474e92 100644 --- a/src/component/messages.test.ts +++ b/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); + }); +}); diff --git a/src/vercel/client/search.test.ts b/src/vercel/client/search.test.ts index b8557ba6..9b51894d 100644 --- a/src/vercel/client/search.test.ts +++ b/src/vercel/client/search.test.ts @@ -616,78 +616,61 @@ describe("search.ts", () => { }); }); - it("should complete an order truncated by the recent messages limit", async () => { + it("keeps an order the component completed past the recent messages limit", async () => { const toolCallId = "call_issue_309"; - vi.mocked(mockCtx.runQuery) - .mockResolvedValueOnce({ - page: [ - createMockMessageDoc( - "assistant-final", - "assistant", - [{ type: "text", text: "The result is 72." }], - 1, - 3, - ), - createMockMessageDoc( - "tool-result", - "tool", - [ - { - type: "tool-result", - toolCallId, - toolName: "getWeather", - output: { - type: "json", - value: { temperature: 72 }, - }, - }, - ], - 1, - 2, - ), - createMockMessageDoc( - "assistant-tool-call", - "assistant", - [ - { - type: "tool-call", - toolCallId, - toolName: "getWeather", - input: { city: "New York" }, + vi.mocked(mockCtx.runQuery).mockResolvedValueOnce({ + page: [ + createMockMessageDoc( + "assistant-final", + "assistant", + [{ type: "text", text: "The result is 72." }], + 1, + 3, + ), + createMockMessageDoc( + "tool-result", + "tool", + [ + { + type: "tool-result", + toolCallId, + toolName: "getWeather", + output: { + type: "json", + value: { temperature: 72 }, }, - ], - 1, - 1, - ), - ], - continueCursor: "complete-order", - isDone: false, - }) - .mockResolvedValueOnce({ - page: [ - createMockMessageDoc( - "user", - "user", - "What is the weather?", - 1, - 0, - ), - ], - continueCursor: "done", - isDone: true, - }); + }, + ], + 1, + 2, + ), + createMockMessageDoc( + "assistant-tool-call", + "assistant", + [ + { + type: "tool-call", + toolCallId, + toolName: "getWeather", + input: { city: "New York" }, + }, + ], + 1, + 1, + ), + createMockMessageDoc("user", "user", "What is the weather?", 1, 0), + ], + continueCursor: "next-order", + isDone: false, + }); - const result = await fetchContextWithPrompt( - mockCtx, - components.agent, - { - ...baseArgs, - prompt: "What should I do next?", - messages: undefined, - promptMessageId: undefined, - contextOptions: { recentMessages: 3 }, - }, - ); + const result = await fetchContextWithPrompt(mockCtx, components.agent, { + ...baseArgs, + prompt: "What should I do next?", + messages: undefined, + promptMessageId: undefined, + contextOptions: { recentMessages: 3 }, + }); expect(result.messages.map((message) => message.role)).toEqual([ "user", @@ -696,61 +679,36 @@ describe("search.ts", () => { "assistant", "user", ]); - expect(mockCtx.runQuery).toHaveBeenNthCalledWith(2, expect.anything(), { - threadId: "thread123", - excludeToolMessages: undefined, - paginationOpts: { numItems: 1, cursor: "complete-order" }, - upToAndIncludingMessageId: undefined, - order: "desc", - statuses: ["success"], - }); + expect(mockCtx.runQuery).toHaveBeenCalledTimes(1); }); - it("should omit an order when its boundary cannot be fetched", async () => { - vi.mocked(mockCtx.runQuery) - .mockResolvedValueOnce({ - page: [ - createMockMessageDoc( - "assistant-final", - "assistant", - "Partial answer", - 1, - 2, - ), - ], - continueCursor: "complete-order", - isDone: false, - }) - .mockResolvedValueOnce({ - page: [ - createMockMessageDoc( - "assistant-step", - "assistant", - "Still partial", - 1, - 1, - ), - ], - continueCursor: "done", - isDone: true, - }); + it("omits an order the component left incomplete", async () => { + vi.mocked(mockCtx.runQuery).mockResolvedValueOnce({ + page: [ + createMockMessageDoc( + "assistant-final", + "assistant", + "Partial answer", + 1, + 2, + ), + ], + continueCursor: "complete-order", + isDone: false, + }); - const result = await fetchContextWithPrompt( - mockCtx, - components.agent, - { - ...baseArgs, - prompt: "New prompt", - messages: undefined, - promptMessageId: undefined, - contextOptions: { recentMessages: 1 }, - }, - ); + const result = await fetchContextWithPrompt(mockCtx, components.agent, { + ...baseArgs, + prompt: "New prompt", + messages: undefined, + promptMessageId: undefined, + contextOptions: { recentMessages: 1 }, + }); expect(result.messages).toEqual([ { role: "user", content: "New prompt" }, ]); - expect(mockCtx.runQuery).toHaveBeenCalledTimes(2); + expect(mockCtx.runQuery).toHaveBeenCalledTimes(1); }); it("should handle input messages correctly", async () => { From 2eab6e88bc928287ef3d02898be0b0a99f8e2304 Mon Sep 17 00:00:00 2001 From: Robel Estifanos Date: Fri, 18 Sep 2026 09:26:46 -0700 Subject: [PATCH 2/2] fix(component): extend message pages to the order boundary Fixes https://github.com/get-convex/agent/issues/193. Complete the oldest order inside listMessagesByThreadId so the continue cursor never splits a UI message. search.ts drops its own completion fetch since the query now does it. completeOrders: false opts out for callers that only want the newest rows. Supersedes the client-side loop in https://github.com/get-convex/agent/pull/321. Claude-Session: https://claude.ai/code/session_01YGMZR56ff1JhEEoMnuwa9e --- src/component/_generated/component.ts | 1 + src/component/messages.ts | 64 ++++++++++++++++++++++-- src/vercel/client/definePlaygroundAPI.ts | 1 + src/vercel/client/search.ts | 22 +------- src/vercel/react/useThreadMessages.ts | 4 ++ src/vercel/react/useUIMessages.ts | 4 ++ 6 files changed, 71 insertions(+), 25 deletions(-) diff --git a/src/component/_generated/component.ts b/src/component/_generated/component.ts index 143108e3..59170e0a 100644 --- a/src/component/_generated/component.ts +++ b/src/component/_generated/component.ts @@ -2232,6 +2232,7 @@ export type ComponentApi = "query", "internal", { + completeOrders?: boolean; excludeToolMessages?: boolean; order: "asc" | "desc"; paginationOpts?: { diff --git a/src/component/messages.ts b/src/component/messages.ts index 9b87d257..2dc94fcb 100644 --- a/src/component/messages.ts +++ b/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) => { diff --git a/src/vercel/client/definePlaygroundAPI.ts b/src/vercel/client/definePlaygroundAPI.ts index 151346e2..ecb00278 100644 --- a/src/vercel/client/definePlaygroundAPI.ts +++ b/src/vercel/client/definePlaygroundAPI.ts @@ -171,6 +171,7 @@ export function definePlaygroundAPI( threadId: thread._id, order: "desc", paginationOpts: { numItems: 1, cursor: null }, + completeOrders: false, }); return { ...thread, diff --git a/src/vercel/client/search.ts b/src/vercel/client/search.ts index 9107ec62..27bc1c55 100644 --- a/src/vercel/client/search.ts +++ b/src/vercel/client/search.ts @@ -37,7 +37,6 @@ import { const DEFAULT_VECTOR_SCORE_THRESHOLD = 0.0; // Bound the rare boundary-extension query; incomplete orders are trimmed below. -const MAX_ORDER_COMPLETION_MESSAGES = 1_000; // 10k characters should be more than enough for most cases, and stays under // the 8k token limit for some models. const MAX_EMBEDDING_TEXT_LENGTH = 10_000; @@ -164,26 +163,7 @@ export async function fetchRecentAndSearchMessages( opts.recentMessages ?? DEFAULT_RECENT_MESSAGES, null, ); - let page = firstPage.page; - const oldest = page.at(-1); - if ( - oldest && - oldest.stepOrder > 0 && - !firstPage.isDone && - firstPage.continueCursor - ) { - const completionPage = await fetchRecentPage( - Math.min(oldest.stepOrder, MAX_ORDER_COMPLETION_MESSAGES), - firstPage.continueCursor, - ); - page = [ - ...page, - ...completionPage.page.filter( - (message) => message.order === oldest.order, - ), - ]; - } - const retained = trimIncompleteOldestOrder(sorted(page)); + const retained = trimIncompleteOldestOrder(sorted(firstPage.page)); included = new Set(retained.map((m) => m._id)); recentMessages = filterOutOrphanedToolMessages(retained); } diff --git a/src/vercel/react/useThreadMessages.ts b/src/vercel/react/useThreadMessages.ts index 6a1b1f14..d7a40336 100644 --- a/src/vercel/react/useThreadMessages.ts +++ b/src/vercel/react/useThreadMessages.ts @@ -120,6 +120,10 @@ export function useThreadMessages>( query: Query, args: ThreadMessagesArgs | "skip", options: { + /** + * Minimum page size. A page is extended to the end of its oldest `order` + * so a UI message is never split across pages, so it can hold more rows. + */ initialNumItems: number; stream?: Query extends StreamQuery ? boolean diff --git a/src/vercel/react/useUIMessages.ts b/src/vercel/react/useUIMessages.ts index 0396adc7..5fa3c791 100644 --- a/src/vercel/react/useUIMessages.ts +++ b/src/vercel/react/useUIMessages.ts @@ -126,6 +126,10 @@ export function useUIMessages>( query: Query, args: UIMessagesQueryArgs | "skip", options: { + /** + * Minimum page size. A page is extended to the end of its oldest `order` + * so a UI message is never split across pages, so it can hold more rows. + */ initialNumItems: number; stream?: Query extends StreamQuery ? boolean