diff --git a/packages/api/src/assistant-message-order.ts b/packages/api/src/assistant-message-order.ts new file mode 100644 index 000000000..7a31cfcf2 --- /dev/null +++ b/packages/api/src/assistant-message-order.ts @@ -0,0 +1,51 @@ +interface OrderableMessage { + id: string; + parentId: string | null; +} + +/** + * Orders thread messages so that a parent always precedes its children. + * + * The database can only order by `created_at`, which is second-granular on MySQL and SQLite. Two + * messages written in the same second therefore tie, and the client rebuilds the conversation tree + * from `parent_id`, so a child arriving before its parent corrupts the thread. Sorting here keeps + * the relative order of siblings from the incoming (already deterministic) list. + */ +export const orderMessagesByParent = (messages: TMessage[]): TMessage[] => { + if (messages.length <= 1) return messages; + + const known = new Set(messages.map((message) => message.id)); + const childrenByParent = new Map(); + + for (const message of messages) { + // Treat a dangling parent reference as a root so the message is never dropped. + const parentId = message.parentId !== null && known.has(message.parentId) ? message.parentId : null; + const siblings = childrenByParent.get(parentId); + if (siblings) siblings.push(message); + else childrenByParent.set(parentId, [message]); + } + + const ordered: TMessage[] = []; + const visited = new Set(); + + const visit = (message: TMessage) => { + // Guards against a cycle introduced by corrupted data. + if (visited.has(message.id)) return; + visited.add(message.id); + ordered.push(message); + for (const child of childrenByParent.get(message.id) ?? []) { + visit(child); + } + }; + + for (const root of childrenByParent.get(null) ?? []) { + visit(root); + } + + // Any message left over is part of a cycle; append it so nothing is silently lost. + for (const message of messages) { + if (!visited.has(message.id)) ordered.push(message); + } + + return ordered; +}; diff --git a/packages/api/src/router/assistant.ts b/packages/api/src/router/assistant.ts index 3c7957d29..3544038a1 100644 --- a/packages/api/src/router/assistant.ts +++ b/packages/api/src/router/assistant.ts @@ -29,6 +29,7 @@ import { import type { createTRPCContext } from "../trpc"; import { fetchOpenRouterGenerationTelemetryAsync } from "../assistant-generation-telemetry"; +import { orderMessagesByParent } from "../assistant-message-order"; import { verifyAssistantGenerationAccessToken } from "../assistant-generation-access"; import { createTRPCRouter, isDemoMode, permissionRequiredProcedure, protectedProcedure } from "../trpc"; import { boardRouter } from "./board"; @@ -710,11 +711,13 @@ export const assistantRouter = createTRPCRouter({ const thread = await ownedThreadAsync(ctx.db, input.threadId, ctx.session.user.id); const messages = await ctx.db.query.assistantMessages.findMany({ where: eq(assistantMessages.threadId, thread.id), - orderBy: assistantMessages.createdAt, + // `created_at` is only second-granular on MySQL and SQLite, so sibling messages written in the + // same second tie. The id tiebreaker makes the database order at least deterministic. + orderBy: [asc(assistantMessages.createdAt), asc(assistantMessages.id)], }); return { thread, - messages: messages.map((message) => ({ + messages: orderMessagesByParent(messages).map((message) => ({ id: message.id, parentId: message.parentId, format: message.format, diff --git a/packages/api/src/test/assistant-message-order.spec.ts b/packages/api/src/test/assistant-message-order.spec.ts new file mode 100644 index 000000000..f97ab9934 --- /dev/null +++ b/packages/api/src/test/assistant-message-order.spec.ts @@ -0,0 +1,63 @@ +import { describe, expect, test } from "vitest"; + +import { orderMessagesByParent } from "../assistant-message-order"; + +const message = (id: string, parentId: string | null) => ({ id, parentId }); + +describe("orderMessagesByParent", () => { + test("keeps an already correct order untouched", () => { + const messages = [message("a", null), message("b", "a"), message("c", "b")]; + + expect(orderMessagesByParent(messages).map((entry) => entry.id)).toEqual(["a", "b", "c"]); + }); + + test("moves a child that the database returned before its parent", () => { + // Happens when created_at ties on MySQL/SQLite, which store whole seconds only. + const messages = [message("c", "b"), message("a", null), message("b", "a")]; + + expect(orderMessagesByParent(messages).map((entry) => entry.id)).toEqual(["a", "b", "c"]); + }); + + test("preserves sibling order from the incoming list", () => { + const messages = [message("root", null), message("first", "root"), message("second", "root")]; + + expect(orderMessagesByParent(messages).map((entry) => entry.id)).toEqual(["root", "first", "second"]); + }); + + test("keeps branches grouped under their parent", () => { + const messages = [ + message("root", null), + message("branch-a", "root"), + message("branch-b", "root"), + message("branch-a-child", "branch-a"), + ]; + + expect(orderMessagesByParent(messages).map((entry) => entry.id)).toEqual([ + "root", + "branch-a", + "branch-a-child", + "branch-b", + ]); + }); + + test("treats a dangling parent reference as a root instead of dropping the message", () => { + const messages = [message("orphan", "deleted-parent"), message("root", null)]; + + expect(orderMessagesByParent(messages).map((entry) => entry.id)).toEqual(["orphan", "root"]); + }); + + test("does not lose messages that form a cycle", () => { + const messages = [message("a", "b"), message("b", "a")]; + + expect( + orderMessagesByParent(messages) + .map((entry) => entry.id) + .toSorted(), + ).toEqual(["a", "b"]); + }); + + test("returns short lists unchanged", () => { + expect(orderMessagesByParent([])).toEqual([]); + expect(orderMessagesByParent([message("only", null)]).map((entry) => entry.id)).toEqual(["only"]); + }); +});