fix(assistant): mysql message storage, migration metadata and per-page bundle cost (#6530)
This commit is contained in:
1 parent
f6f755d5c9
commit
a9b0c2ae47
14 files changed
+527
-129
No files matched your search
@@ -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;
|
||||
};
|
||||
@@ -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,
|
||||
|
||||
@@ -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"]);
|
||||
});
|
||||
});
|
||||
@@ -26,7 +26,7 @@ CREATE TABLE `assistant_message` (
|
||||
`thread_id` varchar(64) NOT NULL,
|
||||
`parent_id` varchar(128),
|
||||
`format` varchar(64) DEFAULT 'ai-sdk/v6' NOT NULL,
|
||||
`content` text DEFAULT ('{"json": {}}') NOT NULL,
|
||||
`content` mediumtext DEFAULT ('{"json": {}}') NOT NULL,
|
||||
`created_at` timestamp DEFAULT (now()) NOT NULL,
|
||||
CONSTRAINT `assistant_message_id` PRIMARY KEY(`id`)
|
||||
);
|
||||
|
||||
@@ -9,6 +9,7 @@ import {
|
||||
foreignKey,
|
||||
index,
|
||||
int,
|
||||
mediumtext,
|
||||
mysqlTable,
|
||||
primaryKey,
|
||||
smallint,
|
||||
@@ -487,6 +488,7 @@ export const assistantConfigurations = mysqlTable("assistant_configuration", {
|
||||
webSearchEnabled: boolean().notNull().default(false),
|
||||
provider: varchar({ length: 32 })
|
||||
.$type<
|
||||
| "homarr"
|
||||
| "openrouter"
|
||||
| "openai"
|
||||
| "anthropic"
|
||||
@@ -536,7 +538,9 @@ export const assistantMessages = mysqlTable(
|
||||
.references(() => assistantThreads.id, { onDelete: "cascade" }),
|
||||
parentId: varchar({ length: 128 }),
|
||||
format: varchar({ length: 64 }).notNull().default("ai-sdk/v6"),
|
||||
content: text().default(emptySuperJSON).notNull(),
|
||||
// `text` only holds 64KB on MySQL, which a single base64 image attachment already exceeds.
|
||||
// `mediumtext` holds 16MB, comfortably above the 2MB request cap enforced by the chat route.
|
||||
content: mediumtext().notNull(),
|
||||
createdAt: timestamp().notNull().defaultNow(),
|
||||
},
|
||||
(message) => ({
|
||||
|
||||
@@ -0,0 +1,98 @@
|
||||
import fs from "fs";
|
||||
import path from "path";
|
||||
import { describe, expect, test } from "vitest";
|
||||
|
||||
/**
|
||||
* Guards the migration folders against the failure modes that have actually shipped:
|
||||
*
|
||||
* - v1.70.0 shipped `0042_add_widget_secret.sql` for MySQL without a `_journal.json` entry, so the
|
||||
* migrator never ran it and MySQL instances upgraded from that release had no `widget_secret`
|
||||
* table.
|
||||
* - The follow-up fix hand-wrote `0042_snapshot.json` with a `checkConstraints` key instead of
|
||||
* `checkConstraint`, which made `drizzle-kit generate` reject the folder and silently exit 0 for
|
||||
* every later MySQL migration.
|
||||
*/
|
||||
const dialects = ["mysql", "postgresql", "sqlite"] as const;
|
||||
|
||||
const migrationsRoot = path.join(__dirname, "..", "migrations");
|
||||
|
||||
interface Journal {
|
||||
entries: { idx: number; tag: string }[];
|
||||
}
|
||||
|
||||
const readDialect = (dialect: (typeof dialects)[number]) => {
|
||||
const folder = path.join(migrationsRoot, dialect);
|
||||
const journal = JSON.parse(fs.readFileSync(path.join(folder, "meta", "_journal.json"), "utf8")) as Journal;
|
||||
const sqlFiles = fs
|
||||
.readdirSync(folder)
|
||||
.filter((file) => file.endsWith(".sql"))
|
||||
.map((file) => file.replace(/\.sql$/u, ""));
|
||||
const snapshots = fs
|
||||
.readdirSync(path.join(folder, "meta"))
|
||||
.filter((file) => /^\d+_snapshot\.json$/u.test(file))
|
||||
.map((file) => file.split("_")[0]);
|
||||
|
||||
return { folder, journal, sqlFiles, snapshots };
|
||||
};
|
||||
|
||||
describe.each(dialects)("%s migrations", (dialect) => {
|
||||
test("every journal entry has a matching .sql file", () => {
|
||||
const { journal, sqlFiles } = readDialect(dialect);
|
||||
const missing = journal.entries.filter((entry) => !sqlFiles.includes(entry.tag)).map((entry) => entry.tag);
|
||||
|
||||
expect(missing).toEqual([]);
|
||||
});
|
||||
|
||||
test("every .sql file is registered in the journal", () => {
|
||||
const { journal, sqlFiles } = readDialect(dialect);
|
||||
const tags = new Set(journal.entries.map((entry) => entry.tag));
|
||||
const orphans = sqlFiles.filter((file) => !tags.has(file));
|
||||
|
||||
expect(orphans).toEqual([]);
|
||||
});
|
||||
|
||||
test("every journal entry has a matching snapshot", () => {
|
||||
const { journal, snapshots } = readDialect(dialect);
|
||||
const missing = journal.entries
|
||||
.filter((entry) => !snapshots.includes(String(entry.idx).padStart(4, "0")))
|
||||
.map((entry) => entry.tag);
|
||||
|
||||
expect(missing).toEqual([]);
|
||||
});
|
||||
|
||||
test("journal indexes are unique and gap-free", () => {
|
||||
const { journal } = readDialect(dialect);
|
||||
const indexes = journal.entries.map((entry) => entry.idx);
|
||||
|
||||
expect(indexes).toEqual(indexes.map((_, position) => position));
|
||||
});
|
||||
|
||||
// The check-constraint key is spelled differently per snapshot format version (MySQL's v5 uses
|
||||
// `checkConstraint`, the newer SQLite/Postgres formats use `checkConstraints`), so rather than
|
||||
// hard-coding a spelling this asserts every snapshot in a dialect agrees with the others. A
|
||||
// hand-written snapshot that gets it wrong stands out as the odd one and fails drizzle-kit.
|
||||
test("all snapshots agree on the check-constraint key spelling", () => {
|
||||
const { folder, snapshots } = readDialect(dialect);
|
||||
|
||||
const spellingsPerSnapshot = snapshots.map((snapshot) => {
|
||||
const file = path.join(folder, "meta", `${snapshot}_snapshot.json`);
|
||||
const parsed = JSON.parse(fs.readFileSync(file, "utf8")) as { tables?: Record<string, object> };
|
||||
const spellings = new Set<string>();
|
||||
|
||||
for (const table of Object.values(parsed.tables ?? {})) {
|
||||
if ("checkConstraint" in table) spellings.add("checkConstraint");
|
||||
if ("checkConstraints" in table) spellings.add("checkConstraints");
|
||||
}
|
||||
|
||||
return { snapshot, spellings: [...spellings].toSorted().join("+") };
|
||||
});
|
||||
|
||||
const used = [...new Set(spellingsPerSnapshot.map((entry) => entry.spellings))].filter(
|
||||
(spelling) => spelling.length > 0,
|
||||
);
|
||||
|
||||
expect(used, `snapshots disagree: ${JSON.stringify(spellingsPerSnapshot.filter((e) => e.spellings))}`).toHaveLength(
|
||||
1,
|
||||
);
|
||||
});
|
||||
});
|
||||
@@ -1,6 +1,7 @@
|
||||
import path from "path";
|
||||
import { readFileSync } from "node:fs";
|
||||
import { MySqlContainer } from "@testcontainers/mysql";
|
||||
import { eq } from "drizzle-orm";
|
||||
import { drizzle } from "drizzle-orm/mysql2";
|
||||
import { migrate } from "drizzle-orm/mysql2/migrator";
|
||||
import mysql from "mysql2";
|
||||
@@ -56,6 +57,22 @@ describe("Mysql Migration", () => {
|
||||
.query<mysql.RowDataPacket[]>("SHOW TABLES LIKE 'custom_widget%'");
|
||||
expect(customWidgetTables).toHaveLength(4);
|
||||
|
||||
// A single base64 image attachment already exceeds the 64KB that `text` holds on MySQL, so
|
||||
// `assistant_message.content` has to be `mediumtext`. Everything below 64KB passes either way.
|
||||
const largeContent = "x".repeat(200_000);
|
||||
await database.insert(mysqlSchema.users).values({ id: "size-check-user", name: "size-check" });
|
||||
await database.insert(mysqlSchema.assistantThreads).values({ id: "size-check-thread", userId: "size-check-user" });
|
||||
await database.insert(mysqlSchema.assistantMessages).values({
|
||||
id: "size-check-message",
|
||||
threadId: "size-check-thread",
|
||||
parentId: null,
|
||||
content: largeContent,
|
||||
});
|
||||
const [storedMessage] = await database.query.assistantMessages.findMany({
|
||||
where: eq(mysqlSchema.assistantMessages.id, "size-check-message"),
|
||||
});
|
||||
expect(storedMessage?.content).toHaveLength(largeContent.length);
|
||||
|
||||
connection.end();
|
||||
await mysqlContainer.stop();
|
||||
}, 120_000);
|
||||
|
||||
Reference in new issue
Block a user