fix(assistant): mysql message storage, migration metadata and per-page bundle cost (#6530)

packages/api/src/assistant-message-order.ts

@@ -0,0 +1,101 @@
+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 = <TMessage extends OrderableMessage>(messages: TMessage[]): TMessage[] => {
+  if (messages.length <= 1) return messages;
+
+  const messagesById = new Map<string, TMessage>();
+  const incomingIndexById = new Map<string, number>();
+  messages.forEach((message, index) => {
+    messagesById.set(message.id, message);
+    incomingIndexById.set(message.id, index);
+  });
+
+  const childrenByParent = new Map<string | null, TMessage[]>();
+  const resolvedParentIds = new Map<string, string | null>();
+
+  for (const message of messages) {
+    // Treat a dangling parent reference as a root so the message is never dropped.
+    const parentId = message.parentId !== null && messagesById.has(message.parentId) ? message.parentId : null;
+    resolvedParentIds.set(message.id, parentId);
+    const siblings = childrenByParent.get(parentId);
+    if (siblings) siblings.push(message);
+    else childrenByParent.set(parentId, [message]);
+  }
+
+  const ordered: TMessage[] = [];
+  const visited = new Set<string>();
+
+  // A conversation is one long parent chain, so the traversal depth grows with the message count.
+  // An explicit stack keeps a long thread from exhausting the call stack.
+  const visitFrom = (start: TMessage) => {
+    const stack: TMessage[] = [start];
+    while (stack.length > 0) {
+      const message = stack.pop();
+      if (!message || visited.has(message.id)) continue;
+      visited.add(message.id);
+      ordered.push(message);
+      const children = childrenByParent.get(message.id);
+      if (!children) continue;
+      // Pushed in reverse so the stack pops siblings back in their incoming order.
+      for (let index = children.length - 1; index >= 0; index--) {
+        const child = children[index];
+        if (child) stack.push(child);
+      }
+    }
+  };
+
+  for (const root of childrenByParent.get(null) ?? []) {
+    visitFrom(root);
+  }
+
+  // Whatever is left belongs to a parent cycle, which by definition has no root to start from.
+  // Corrupted data should still come back complete and in a sensible order, so each cycle is
+  // broken at its earliest member in the incoming order and traversed from there. Only that one
+  // back-edge can violate parent-before-child; every other message in the component is fine.
+  const findCycleEntry = (start: TMessage) => {
+    const path: TMessage[] = [start];
+    const positionInPath = new Map<string, number>([[start.id, 0]]);
+    let current = start;
+
+    for (;;) {
+      const parentId = resolvedParentIds.get(current.id) ?? null;
+      if (parentId === null || visited.has(parentId)) return current;
+
+      const parent = messagesById.get(parentId);
+      if (!parent) return current;
+
+      const loopStart = positionInPath.get(parent.id);
+      if (loopStart !== undefined) {
+        // Closed the loop: the cycle is the tail of the walked path. Start from whichever of its
+        // members came first in the incoming list so the choice is deterministic.
+        return path.slice(loopStart).reduce((earliest, member) => {
+          const memberIndex = incomingIndexById.get(member.id) ?? 0;
+          const earliestIndex = incomingIndexById.get(earliest.id) ?? 0;
+          return memberIndex < earliestIndex ? member : earliest;
+        });
+      }
+
+      positionInPath.set(parent.id, path.length);
+      path.push(parent);
+      current = parent;
+    }
+  };
+
+  for (const message of messages) {
+    if (visited.has(message.id)) continue;
+    visitFrom(findCycleEntry(message));
+  }
+
+  return ordered;
+};

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,

packages/api/src/test/assistant-message-order.spec.ts

@@ -0,0 +1,89 @@
+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("keeps a child of a cycle behind its parent", () => {
+    // `child` hangs off the cycle a -> b -> a. Only the cycle's own back-edge can break the
+    // parent-before-child contract; `child` must still come after `a`.
+    const messages = [message("child", "a"), message("a", "b"), message("b", "a")];
+
+    const ordered = orderMessagesByParent(messages).map((entry) => entry.id);
+
+    expect(ordered.toSorted()).toEqual(["a", "b", "child"]);
+    expect(ordered.indexOf("a")).toBeLessThan(ordered.indexOf("child"));
+  });
+
+  test("handles a chain deeper than the call stack", () => {
+    // `getThread` loads every message in a thread, and a conversation is one long parent chain, so
+    // a recursive walk would blow the stack on a long thread.
+    const depth = 50_000;
+    const messages = Array.from({ length: depth }, (_, index) =>
+      message(`m${index}`, index === 0 ? null : `m${index - 1}`),
+    );
+
+    const ordered = orderMessagesByParent(messages);
+
+    expect(ordered).toHaveLength(depth);
+    expect(ordered[0]?.id).toBe("m0");
+    expect(ordered[depth - 1]?.id).toBe(`m${depth - 1}`);
+  });
+
+  test("returns short lists unchanged", () => {
+    expect(orderMessagesByParent([])).toEqual([]);
+    expect(orderMessagesByParent([message("only", null)]).map((entry) => entry.id)).toEqual(["only"]);
+  });
+});