fix: harden assistant integration

This commit is contained in:
ajnart committed 2026-08-14 11:32:09 +02:00
1 parent 10ef871777
commit 5db95f2def
14 files changed
+349 -189

No files matched your search

+39 -33
View File
@@ -1,3 +1,5 @@
import { hkdfSync } from "node:crypto";
import { createOpenAICompatible } from "@ai-sdk/openai-compatible";
import type { ToolSet, UIMessage } from "ai";
import { convertToModelMessages, jsonSchema, stepCountIs, streamText, tool } from "ai";
@@ -13,11 +15,17 @@ import { and, eq } from "@homarr/db";
import { db } from "@homarr/db";
import { assistantConfigurations, assistantThreads } from "@homarr/db/schema";
import { browserToolContracts } from "~/components/assistant/assistant-tool-contracts";
import { extractMcpTools } from "../../mcp/_extract-tools";
export const maxDuration = 60;
const logger = createLogger({ module: "assistant" });
const toolApprovalSecret = Buffer.from(
hkdfSync("sha256", Buffer.from(env.SECRET_ENCRYPTION_KEY, "hex"), "", "assistant-tool-approval", 32),
).toString("base64url");
const safeStreamError = "The model endpoint stopped the response. Check its URL, model, and credentials.";
const requestSchema = z.object({
id: z.string().min(1).max(64),
@@ -45,26 +53,6 @@ const requestSchema = z.object({
.optional(),
});
const browserToolDefinitions = {
navigate_to_route: {
description: "Navigate the current Homarr tab to a safe internal route.",
parameters: {
type: "object",
properties: { path: { type: "string", description: "An internal Homarr path beginning with /" } },
required: ["path"],
additionalProperties: false,
},
},
open_command_menu: {
description: "Open Homarr's command and search menu.",
parameters: { type: "object", properties: {}, additionalProperties: false },
},
open_media_request_search: {
description: "Open Homarr's media request search interface.",
parameters: { type: "object", properties: {}, additionalProperties: false },
},
} as const;
const assistantInstructions = `You are Homarr Assistant, embedded in the user's self-hosted Homarr dashboard.
Use the available Homarr tools whenever live instance data is needed. Never invent integrations, boards, apps, users, media, system status, or action results.
@@ -164,11 +152,10 @@ export async function POST(request: Request) {
const context = createTRPCContext({ headers: request.headers, session });
const caller = mcpRouter.createCaller(context);
const procedureTypes = getProcedureTypeMap();
const mcpTools = extractMcpTools();
const homarrTools = Object.fromEntries(
extractMcpTools().map((mcpTool) => {
const procedurePath = mcpTool.pathInRouter.join(".");
const isMutation = procedureTypes.get(procedurePath) === "mutation";
mcpTools.map((mcpTool) => {
return [
mcpTool.name,
tool({
@@ -176,7 +163,6 @@ export async function POST(request: Request) {
inputSchema: jsonSchema(
(mcpTool.inputSchema ?? { type: "object", properties: {} }) as Parameters<typeof jsonSchema>[0],
),
needsApproval: isMutation,
execute: async (input) => {
try {
const procedure = mcpTool.pathInRouter.reduce<unknown>(
@@ -205,17 +191,24 @@ export async function POST(request: Request) {
];
}),
) satisfies ToolSet;
const toolApproval = Object.fromEntries(
mcpTools.flatMap((mcpTool) =>
procedureTypes.get(mcpTool.pathInRouter.join(".")) === "mutation"
? [[mcpTool.name, "user-approval" as const]]
: [],
),
);
const frontendTools = Object.fromEntries(
Object.keys(parsed.data.tools ?? {}).flatMap((name) => {
if (!(name in browserToolDefinitions)) return [];
const definition = browserToolDefinitions[name as keyof typeof browserToolDefinitions];
if (!(name in browserToolContracts)) return [];
const definition = browserToolContracts[name as keyof typeof browserToolContracts];
return [
[
name,
tool({
description: definition.description,
inputSchema: jsonSchema(definition.parameters as Parameters<typeof jsonSchema>[0]),
inputSchema: jsonSchema(z.toJSONSchema(definition.parameters) as Parameters<typeof jsonSchema>[0]),
}),
] as const,
];
@@ -249,16 +242,29 @@ export async function POST(request: Request) {
abortSignal: request.signal,
timeout: { totalMs: 55_000, stepMs: 30_000, toolMs: 30_000 },
maxRetries: 2,
experimental_toolApprovalSecret: env.SECRET_ENCRYPTION_KEY,
toolApproval,
experimental_toolApprovalSecret: toolApprovalSecret,
onError: ({ error }) => {
logger.error("Assistant response stream failed", {
error: error instanceof Error ? error.message : String(error),
});
},
onFinish: async () => {
await db
.update(assistantThreads)
.set({ modelId: configuration.modelId, updatedAt: new Date() })
.where(eq(assistantThreads.id, thread.id));
try {
await db
.update(assistantThreads)
.set({ modelId: configuration.modelId, updatedAt: new Date() })
.where(eq(assistantThreads.id, thread.id));
} catch (error) {
logger.error("Failed to update assistant conversation metadata", {
threadId: thread.id,
error: error instanceof Error ? error.message : String(error),
});
}
},
});
return result.toUIMessageStreamResponse();
return result.toUIMessageStreamResponse({ onError: () => safeStreamError });
} catch {
return Response.json(
{ error: "The configured model endpoint could not start this response. Check its URL, model, and credentials." },
@@ -1,7 +1,3 @@
.drawer {
--assistant-panel-width: min(32rem, 100vw);
}
.content {
height: 100dvh;
background: var(--mantine-color-body);
@@ -13,6 +13,7 @@ import {
useAuiState,
} from "@assistant-ui/react";
import { ActionIcon, Box, Button, Drawer, Group, ScrollArea, Stack, Text, ThemeIcon, Tooltip } from "@mantine/core";
import { useMediaQuery } from "@mantine/hooks";
import {
IconArchive,
IconArrowUp,
@@ -206,8 +207,17 @@ const ArchivedThreadListItem = () => {
const ThreadHistory = ({ opened, onClose }: { opened: boolean; onClose: () => void }) => {
const t = useScopedI18n("common.assistant");
const isMobile = useMediaQuery("(max-width: 48em)");
const hidden = isMobile && !opened;
return (
<Stack className={classes.history} gap="xs" p="sm" data-opened={opened}>
<Stack
className={classes.history}
gap="xs"
p="sm"
data-opened={opened}
inert={hidden ? true : undefined}
aria-hidden={hidden || undefined}
>
<Group gap="xs" wrap="nowrap">
<ThreadListPrimitive.New asChild>
<Button variant="light" leftSection={<IconPlus size={16} />} fullWidth onClick={onClose}>
@@ -3,7 +3,6 @@
import type { PropsWithChildren } from "react";
import { createContext, useCallback, useContext, useMemo, useState } from "react";
import { useRouter } from "next/navigation";
import { useChat } from "@ai-sdk/react";
import type { MessageFormatAdapter, RemoteThreadListAdapter, ThreadHistoryAdapter, Toolkit } from "@assistant-ui/react";
import {
AssistantRuntimeProvider,
@@ -13,12 +12,11 @@ import {
useAuiState,
useRemoteThreadListRuntime,
} from "@assistant-ui/react";
import { AssistantChatTransport, useAISDKRuntime } from "@assistant-ui/react-ai-sdk";
import { AssistantChatTransport, useChatRuntime } from "@assistant-ui/react-ai-sdk";
import { useHotkeys } from "@mantine/hooks";
import { createAssistantStream } from "assistant-stream";
import type { UIMessage } from "ai";
import { lastAssistantMessageIsCompleteWithApprovalResponses } from "ai";
import { z } from "zod/v4";
import { clientApi, fetchApi } from "@homarr/api/client";
import { useSession } from "@homarr/auth/client";
@@ -32,6 +30,7 @@ import {
} from "@homarr/spotlight";
import { AssistantPanel } from "./assistant-panel";
import { browserToolContracts } from "./assistant-tool-contracts";
interface AssistantContextValue {
enabled: boolean;
@@ -178,22 +177,15 @@ const createHistoryAdapter = (threadId: string | undefined): ThreadHistoryAdapte
});
const AssistantThreadRuntime = () => {
const localThreadId = useAuiState((state) => state.threadListItem.id);
const threadId = useAuiState((state) => state.threadListItem.remoteId);
const aui = useAui();
const transport = useMemo(() => new AssistantChatTransport({ api: "/api/assistant/chat" }), []);
const chat = useChat<UIMessage>({
id: localThreadId,
const history = useMemo(() => createHistoryAdapter(threadId), [threadId]);
return useChatRuntime<UIMessage>({
transport,
sendAutomaticallyWhen: lastAssistantMessageIsCompleteWithApprovalResponses,
adapters: { history },
});
const history = useMemo(() => createHistoryAdapter(threadId), [threadId]);
const runtime = useAISDKRuntime(chat, { adapters: { history } });
transport.setRuntime(runtime);
transport["__internal_setGetThreadListItem"](() => (aui.threadListItem.source ? aui.threadListItem() : undefined));
return runtime;
};
const AssistantRuntime = ({ children }: PropsWithChildren) => {
@@ -208,17 +200,20 @@ const AssistantRuntime = ({ children }: PropsWithChildren) => {
defineToolkit({
navigate_to_route: {
type: "frontend",
description:
"Navigate the current Homarr tab to a safe internal route. Only paths beginning with a single slash are accepted.",
parameters: z.object({
path: z.string().describe("An internal Homarr route, for example /manage/apps"),
}),
...browserToolContracts.navigate_to_route,
execute: async ({ path }) => {
if (!path.startsWith("/") || path.startsWith("//")) {
const target = URL.canParse(path, window.location.origin) ? new URL(path, window.location.origin) : null;
if (
!path.startsWith("/") ||
path.startsWith("/\\") ||
target === null ||
target.origin !== window.location.origin
) {
return { success: false, error: "Only internal Homarr paths are allowed." };
}
router.push(path);
return { success: true, path };
const internalPath = `${target.pathname}${target.search}${target.hash}`;
router.push(internalPath);
return { success: true, path: internalPath };
},
renderText: {
running: ({ args }) => `Opening ${args.path}…`,
@@ -227,8 +222,7 @@ const AssistantRuntime = ({ children }: PropsWithChildren) => {
},
open_command_menu: {
type: "frontend",
description: "Open Homarr's command and search menu.",
parameters: z.object({}),
...browserToolContracts.open_command_menu,
execute: async () => {
openSpotlight();
return { success: true };
@@ -237,8 +231,7 @@ const AssistantRuntime = ({ children }: PropsWithChildren) => {
},
open_media_request_search: {
type: "frontend",
description: "Open Homarr's media request search interface.",
parameters: z.object({}),
...browserToolContracts.open_media_request_search,
execute: async () => {
openMediaRequestSearch();
return { success: true };
@@ -0,0 +1,19 @@
import { z } from "zod/v4";
export const browserToolContracts = {
navigate_to_route: {
description:
"Navigate the current Homarr tab to a safe internal route. Only same-origin paths beginning with a single slash are accepted.",
parameters: z.object({
path: z.string().describe("An internal Homarr route, for example /manage/apps"),
}),
},
open_command_menu: {
description: "Open Homarr's command and search menu.",
parameters: z.object({}),
},
open_media_request_search: {
description: "Open Homarr's media request search interface.",
parameters: z.object({}),
},
} as const;
+16 -9
View File
@@ -195,15 +195,22 @@ const ownedThreadAsync = async (db: Database, threadId: string, userId: string)
};
export const assistantRouter = createTRPCRouter({
getAvailability: protectedProcedure.query(async ({ ctx }) => {
const configuration = await getConfigurationAsync(ctx.db);
const requiresApiKey = configuration?.provider === "openrouter" || configuration?.provider === "openai";
return {
enabled: Boolean(
configuration?.enabled && configuration.modelId && (!requiresApiKey || configuration.encryptedApiKey),
),
};
}),
getAvailability: protectedProcedure
.meta({
mcp: {
enabled: true,
description: "Check whether Homarr Assistant is enabled and configured for the current user",
},
})
.query(async ({ ctx }) => {
const configuration = await getConfigurationAsync(ctx.db);
const requiresApiKey = configuration?.provider === "openrouter" || configuration?.provider === "openai";
return {
enabled: Boolean(
configuration?.enabled && configuration.modelId && (!requiresApiKey || configuration.encryptedApiKey),
),
};
}),
getAdminConfiguration: adminProcedure.query(async ({ ctx }) => {
const configuration = await getConfigurationAsync(ctx.db);
+61 -4
View File
@@ -1,8 +1,9 @@
import type { TRPCError } from "@trpc/server";
import { describe, expect, test, vi } from "vitest";
import { extractToolsFromProcedures } from "trpc-to-mcp";
import { z } from "zod/v4";
import { mcpRouter } from "../mcp";
import { extractMcpToolsFromProcedures } from "../mcp-tools";
vi.mock("@homarr/auth", () => ({}));
@@ -18,9 +19,14 @@ const MCP_TOOL_ALLOWLIST = {
"beszel_getSystems",
"beszel_getSystemStats",
"board_getAllBoards",
"board_getBoardSettings",
"calendar_findAllEvents",
"customWidget_get",
"customWidget_getAuthoringPrompt",
"customWidget_getComponent",
"customWidget_getComponentCatalog",
"customWidget_getExample",
"customWidget_getSharedProps",
"customWidget_getSkill",
"customWidget_legacyMigrationPrompt",
"customWidget_list",
@@ -97,6 +103,7 @@ const MCP_TOOL_ALLOWLIST = {
"apiKeys_delete",
"customWidget_configurationRequestUser",
"customWidget_create",
"customWidget_createFromPreview",
"customWidget_migrateLegacy",
"customWidget_previewCreate",
"customWidget_secretSet",
@@ -124,7 +131,7 @@ function actualToolInventory() {
string,
{ _def?: { type?: "query" | "mutation" } }
>;
return extractToolsFromProcedures(mcpRouter)
return extractMcpToolsFromProcedures(mcpRouter)
.map((tool) => {
const procedure = procedures[tool.pathInRouter.join(".")];
const type = procedure?.["_def"]?.type;
@@ -143,7 +150,7 @@ describe("production MCP router", () => {
});
test("keeps every secret-bearing tool behind a mutation procedure", () => {
const tools = extractToolsFromProcedures(mcpRouter);
const tools = extractMcpToolsFromProcedures(mcpRouter);
const procedures = mcpRouter["_def"].procedures as unknown as Record<
string,
{ _def?: { type?: "query" | "mutation" } }
@@ -157,7 +164,7 @@ describe("production MCP router", () => {
});
test("gives every production tool a description", () => {
const tools = extractToolsFromProcedures(mcpRouter);
const tools = extractMcpToolsFromProcedures(mcpRouter);
for (const tool of tools) expect(tool.description, `Tool ${tool.name} should have a description`).toBeTruthy();
expect(tools.find((tool) => tool.name === "customWidget_getAuthoringPrompt")?.description).toBe(
"Get the current Custom Widget authoring instructions.",
@@ -165,6 +172,14 @@ describe("production MCP router", () => {
});
});
test("MCP tools are deterministically ordered and retain executable schemas", () => {
const tools = extractMcpToolsFromProcedures(mcpRouter);
const toolNames = tools.map((tool) => tool.name);
expect(toolNames).toEqual(toolNames.toSorted((left, right) => left.localeCompare(right)));
expect(tools.every((tool) => tool.inputValidator instanceof z.ZodObject)).toBe(true);
});
describe("custom widget authoring procedure access", () => {
const unauthenticatedCaller = mcpRouter.createCaller({
db: null as never,
@@ -176,9 +191,14 @@ describe("custom widget authoring procedure access", () => {
const calls = [
unauthenticatedCaller.customWidget.schema(),
unauthenticatedCaller.customWidget.getAuthoringPrompt(),
unauthenticatedCaller.customWidget.getComponentCatalog(),
unauthenticatedCaller.customWidget.getComponent({ name: "Stack" }),
unauthenticatedCaller.customWidget.getExample({ name: "service-dashboard" }),
unauthenticatedCaller.customWidget.getSharedProps({ names: ["p"] }),
unauthenticatedCaller.customWidget.getSkill(),
unauthenticatedCaller.customWidget.list(),
unauthenticatedCaller.customWidget.validate({ widget: {} }),
unauthenticatedCaller.customWidget.createFromPreview({ previewSessionId: "preview" }),
unauthenticatedCaller.customWidget.secretSet({
definitionId: "widget",
secret: { sourceId: "default", kind: "apiKey", value: "secret" },
@@ -198,7 +218,12 @@ describe("custom widget authoring procedure access", () => {
const calls = [
nonAdminCaller.customWidget.schema(),
nonAdminCaller.customWidget.getAuthoringPrompt(),
nonAdminCaller.customWidget.getComponentCatalog(),
nonAdminCaller.customWidget.getComponent({ name: "Stack" }),
nonAdminCaller.customWidget.getExample({ name: "service-dashboard" }),
nonAdminCaller.customWidget.getSharedProps({ names: ["p"] }),
nonAdminCaller.customWidget.getSkill(),
nonAdminCaller.customWidget.createFromPreview({ previewSessionId: "preview" }),
nonAdminCaller.customWidget.secretSet({
definitionId: "widget",
secret: { sourceId: "default", kind: "apiKey", value: "secret" },
@@ -223,5 +248,37 @@ describe("custom widget authoring procedure access", () => {
await expect(adminCaller.customWidget.getSkill()).resolves.toMatchObject({
name: expect.any(String),
});
await expect(adminCaller.customWidget.getComponentCatalog()).resolves.toMatchObject({
components: expect.arrayContaining([expect.objectContaining({ name: "Stack" })]),
examples: expect.arrayContaining([expect.objectContaining({ id: "service-dashboard" })]),
});
await expect(adminCaller.customWidget.getComponent({ name: "Stack" })).resolves.toMatchObject({
name: "Stack",
props: expect.anything(),
});
await expect(adminCaller.customWidget.getExample({ name: "service-dashboard" })).resolves.toMatchObject({
id: "service-dashboard",
widget: expect.objectContaining({
$schema: "homarr-custom-widget-v2",
templateLines: expect.any(Array),
}),
});
await expect(adminCaller.customWidget.getExample({ name: "pokedex" })).resolves.toMatchObject({
id: "pokedex",
widget: expect.objectContaining({
name: "Pokédex",
templateLines: expect.arrayContaining([expect.stringContaining("<Stack")]),
}),
});
await expect(adminCaller.customWidget.getSharedProps({ names: ["p", "m"] })).resolves.toMatchObject({
props: [expect.objectContaining({ name: "p" }), expect.objectContaining({ name: "m" })],
notFound: [],
});
await expect(adminCaller.customWidget.getComponent({ name: "NotAComponent" })).rejects.toEqual(
expect.objectContaining<Partial<TRPCError>>({ code: "NOT_FOUND" }),
);
await expect(adminCaller.customWidget.getExample({ name: "not-an-example" })).rejects.toEqual(
expect.objectContaining<Partial<TRPCError>>({ code: "NOT_FOUND" }),
);
});
});
@@ -23,9 +23,9 @@ CREATE TABLE `assistant_thread` (
);
--> statement-breakpoint
CREATE TABLE `assistant_message` (
`id` varchar(64) NOT NULL,
`id` varchar(128) NOT NULL,
`thread_id` varchar(64) NOT NULL,
`parent_id` varchar(64),
`parent_id` varchar(128),
`format` varchar(64) DEFAULT 'ai-sdk/v6' NOT NULL,
`content` text DEFAULT ('{"json": {}}') NOT NULL,
`created_at` timestamp DEFAULT (now()) NOT NULL,
@@ -35,3 +35,7 @@ CREATE TABLE `assistant_message` (
ALTER TABLE `assistant_thread` ADD CONSTRAINT `assistant_thread_user_id_user_id_fk` FOREIGN KEY (`user_id`) REFERENCES `user`(`id`) ON DELETE cascade ON UPDATE no action;
--> statement-breakpoint
ALTER TABLE `assistant_message` ADD CONSTRAINT `assistant_message_thread_id_assistant_thread_id_fk` FOREIGN KEY (`thread_id`) REFERENCES `assistant_thread`(`id`) ON DELETE cascade ON UPDATE no action;
--> statement-breakpoint
CREATE INDEX `assistant_message__thread_id_created_at_idx` ON `assistant_message` (`thread_id`,`created_at`);
--> statement-breakpoint
CREATE INDEX `assistant_thread__user_id_updated_at_idx` ON `assistant_thread` (`user_id`,`updated_at`);
@@ -11,9 +11,9 @@ CREATE TABLE "assistant_configuration" (
);
--> statement-breakpoint
CREATE TABLE "assistant_message" (
"id" varchar(64) PRIMARY KEY NOT NULL,
"id" varchar(128) PRIMARY KEY NOT NULL,
"thread_id" varchar(64) NOT NULL,
"parent_id" varchar(64),
"parent_id" varchar(128),
"format" varchar(64) DEFAULT 'ai-sdk/v6' NOT NULL,
"content" text DEFAULT '{"json": {}}' NOT NULL,
"created_at" timestamp DEFAULT now() NOT NULL
@@ -30,4 +30,6 @@ CREATE TABLE "assistant_thread" (
);
--> statement-breakpoint
ALTER TABLE "assistant_message" ADD CONSTRAINT "assistant_message_thread_id_assistant_thread_id_fk" FOREIGN KEY ("thread_id") REFERENCES "public"."assistant_thread"("id") ON DELETE cascade ON UPDATE no action;--> statement-breakpoint
ALTER TABLE "assistant_thread" ADD CONSTRAINT "assistant_thread_user_id_user_id_fk" FOREIGN KEY ("user_id") REFERENCES "public"."user"("id") ON DELETE cascade ON UPDATE no action;
ALTER TABLE "assistant_thread" ADD CONSTRAINT "assistant_thread_user_id_user_id_fk" FOREIGN KEY ("user_id") REFERENCES "public"."user"("id") ON DELETE cascade ON UPDATE no action;--> statement-breakpoint
CREATE INDEX "assistant_message__thread_id_created_at_idx" ON "assistant_message" USING btree ("thread_id","created_at");--> statement-breakpoint
CREATE INDEX "assistant_thread__user_id_updated_at_idx" ON "assistant_thread" USING btree ("user_id","updated_at");
@@ -30,3 +30,7 @@ CREATE TABLE `assistant_thread` (
`updated_at` integer DEFAULT (unixepoch()) NOT NULL,
FOREIGN KEY (`user_id`) REFERENCES `user`(`id`) ON UPDATE no action ON DELETE cascade
);
--> statement-breakpoint
CREATE INDEX `assistant_message__thread_id_created_at_idx` ON `assistant_message` (`thread_id`,`created_at`);
--> statement-breakpoint
CREATE INDEX `assistant_thread__user_id_updated_at_idx` ON `assistant_thread` (`user_id`,`updated_at`);
+33 -21
View File
@@ -496,28 +496,40 @@ export const assistantConfigurations = mysqlTable("assistant_configuration", {
updatedAt: timestamp().notNull().defaultNow(),
});
export const assistantThreads = mysqlTable("assistant_thread", {
id: varchar({ length: 64 }).notNull().primaryKey(),
userId: varchar({ length: 64 })
.notNull()
.references(() => users.id, { onDelete: "cascade" }),
title: varchar({ length: 256 }),
modelId: varchar({ length: 256 }),
status: varchar({ length: 16 }).$type<"regular" | "archived">().notNull().default("regular"),
createdAt: timestamp().notNull().defaultNow(),
updatedAt: timestamp().notNull().defaultNow(),
});
export const assistantThreads = mysqlTable(
"assistant_thread",
{
id: varchar({ length: 64 }).notNull().primaryKey(),
userId: varchar({ length: 64 })
.notNull()
.references(() => users.id, { onDelete: "cascade" }),
title: varchar({ length: 256 }),
modelId: varchar({ length: 256 }),
status: varchar({ length: 16 }).$type<"regular" | "archived">().notNull().default("regular"),
createdAt: timestamp().notNull().defaultNow(),
updatedAt: timestamp().notNull().defaultNow(),
},
(thread) => ({
userUpdatedAtIdx: index("assistant_thread__user_id_updated_at_idx").on(thread.userId, thread.updatedAt),
}),
);
export const assistantMessages = mysqlTable("assistant_message", {
id: varchar({ length: 64 }).notNull().primaryKey(),
threadId: varchar({ length: 64 })
.notNull()
.references(() => assistantThreads.id, { onDelete: "cascade" }),
parentId: varchar({ length: 64 }),
format: varchar({ length: 64 }).notNull().default("ai-sdk/v6"),
content: text().default(emptySuperJSON).notNull(),
createdAt: timestamp().notNull().defaultNow(),
});
export const assistantMessages = mysqlTable(
"assistant_message",
{
id: varchar({ length: 128 }).notNull().primaryKey(),
threadId: varchar({ length: 64 })
.notNull()
.references(() => assistantThreads.id, { onDelete: "cascade" }),
parentId: varchar({ length: 128 }),
format: varchar({ length: 64 }).notNull().default("ai-sdk/v6"),
content: text().default(emptySuperJSON).notNull(),
createdAt: timestamp().notNull().defaultNow(),
},
(message) => ({
threadCreatedAtIdx: index("assistant_message__thread_id_created_at_idx").on(message.threadId, message.createdAt),
}),
);
export const apiKeyRelations = relations(apiKeys, ({ one }) => ({
user: one(users, {
+33 -21
View File
@@ -495,28 +495,40 @@ export const assistantConfigurations = pgTable("assistant_configuration", {
updatedAt: timestamp().notNull().defaultNow(),
});
export const assistantThreads = pgTable("assistant_thread", {
id: varchar({ length: 64 }).notNull().primaryKey(),
userId: varchar({ length: 64 })
.notNull()
.references(() => users.id, { onDelete: "cascade" }),
title: varchar({ length: 256 }),
modelId: varchar({ length: 256 }),
status: varchar({ length: 16 }).$type<"regular" | "archived">().notNull().default("regular"),
createdAt: timestamp().notNull().defaultNow(),
updatedAt: timestamp().notNull().defaultNow(),
});
export const assistantThreads = pgTable(
"assistant_thread",
{
id: varchar({ length: 64 }).notNull().primaryKey(),
userId: varchar({ length: 64 })
.notNull()
.references(() => users.id, { onDelete: "cascade" }),
title: varchar({ length: 256 }),
modelId: varchar({ length: 256 }),
status: varchar({ length: 16 }).$type<"regular" | "archived">().notNull().default("regular"),
createdAt: timestamp().notNull().defaultNow(),
updatedAt: timestamp().notNull().defaultNow(),
},
(thread) => ({
userUpdatedAtIdx: index("assistant_thread__user_id_updated_at_idx").on(thread.userId, thread.updatedAt),
}),
);
export const assistantMessages = pgTable("assistant_message", {
id: varchar({ length: 64 }).notNull().primaryKey(),
threadId: varchar({ length: 64 })
.notNull()
.references(() => assistantThreads.id, { onDelete: "cascade" }),
parentId: varchar({ length: 64 }),
format: varchar({ length: 64 }).notNull().default("ai-sdk/v6"),
content: text().default(emptySuperJSON).notNull(),
createdAt: timestamp().notNull().defaultNow(),
});
export const assistantMessages = pgTable(
"assistant_message",
{
id: varchar({ length: 128 }).notNull().primaryKey(),
threadId: varchar({ length: 64 })
.notNull()
.references(() => assistantThreads.id, { onDelete: "cascade" }),
parentId: varchar({ length: 128 }),
format: varchar({ length: 64 }).notNull().default("ai-sdk/v6"),
content: text().default(emptySuperJSON).notNull(),
createdAt: timestamp().notNull().defaultNow(),
},
(message) => ({
threadCreatedAtIdx: index("assistant_message__thread_id_created_at_idx").on(message.threadId, message.createdAt),
}),
);
export const apiKeyRelations = relations(apiKeys, ({ one }) => ({
user: one(users, {
+39 -27
View File
@@ -479,34 +479,46 @@ export const assistantConfigurations = sqliteTable("assistant_configuration", {
.default(sql`(unixepoch())`),
});
export const assistantThreads = sqliteTable("assistant_thread", {
id: text().notNull().primaryKey(),
userId: text()
.notNull()
.references(() => users.id, { onDelete: "cascade" }),
title: text(),
modelId: text(),
status: text().$type<"regular" | "archived">().notNull().default("regular"),
createdAt: int({ mode: "timestamp" })
.notNull()
.default(sql`(unixepoch())`),
updatedAt: int({ mode: "timestamp" })
.notNull()
.default(sql`(unixepoch())`),
});
export const assistantThreads = sqliteTable(
"assistant_thread",
{
id: text().notNull().primaryKey(),
userId: text()
.notNull()
.references(() => users.id, { onDelete: "cascade" }),
title: text(),
modelId: text(),
status: text().$type<"regular" | "archived">().notNull().default("regular"),
createdAt: int({ mode: "timestamp" })
.notNull()
.default(sql`(unixepoch())`),
updatedAt: int({ mode: "timestamp" })
.notNull()
.default(sql`(unixepoch())`),
},
(thread) => ({
userUpdatedAtIdx: index("assistant_thread__user_id_updated_at_idx").on(thread.userId, thread.updatedAt),
}),
);
export const assistantMessages = sqliteTable("assistant_message", {
id: text().notNull().primaryKey(),
threadId: text()
.notNull()
.references(() => assistantThreads.id, { onDelete: "cascade" }),
parentId: text(),
format: text().notNull().default("ai-sdk/v6"),
content: text().default(emptySuperJSON).notNull(),
createdAt: int({ mode: "timestamp" })
.notNull()
.default(sql`(unixepoch())`),
});
export const assistantMessages = sqliteTable(
"assistant_message",
{
id: text().notNull().primaryKey(),
threadId: text()
.notNull()
.references(() => assistantThreads.id, { onDelete: "cascade" }),
parentId: text(),
format: text().notNull().default("ai-sdk/v6"),
content: text().default(emptySuperJSON).notNull(),
createdAt: int({ mode: "timestamp" })
.notNull()
.default(sql`(unixepoch())`),
},
(message) => ({
threadCreatedAtIdx: index("assistant_message__thread_id_created_at_idx").on(message.threadId, message.createdAt),
}),
);
export const apiKeyRelations = relations(apiKeys, ({ one }) => ({
user: one(users, {
+64 -38
View File
@@ -1,8 +1,11 @@
/* eslint-disable @typescript-eslint/no-non-null-assertion */
/* eslint-disable import/namespace -- This test compares equivalent schema exports by shared dynamic keys. */
import { is } from "drizzle-orm";
import type { Column, InferSelectModel } from "drizzle-orm";
import { MySqlTable, getTableConfig as getMysqlTableConfig } from "drizzle-orm/mysql-core";
import type { ForeignKey as MysqlForeignKey, MySqlTableWithColumns } from "drizzle-orm/mysql-core";
import type { PgTableWithColumns, ForeignKey as PostgresqlForeignKey } from "drizzle-orm/pg-core";
import { SQLiteTable, getTableConfig as getSqliteTableConfig } from "drizzle-orm/sqlite-core";
import type { ForeignKey as SqliteForeignKey, SQLiteTableWithColumns } from "drizzle-orm/sqlite-core";
import { expect, expectTypeOf, test } from "vitest";
@@ -12,6 +15,9 @@ import * as mysqlSchema from "../schema/mysql";
import * as postgresqlSchema from "../schema/postgresql";
import * as sqliteSchema from "../schema/sqlite";
const mysqlTables = { ...mysqlSchema };
const postgresqlTables = { ...postgresqlSchema };
// We need the following three types as there is currently no support for Buffer in mysql & pg and
// so we use a custom type which results in the config beeing different
type FixedMysqlConfig = {
@@ -51,6 +57,21 @@ type FixedSqliteConfig = {
};
};
/** Identifies a foreign key by what it links rather than by its generated name. */
const referenceSignature = (foreignKey: {
reference: () => { columns: { name: string }[]; foreignTable: object; foreignColumns: { name: string }[] };
}) => {
const reference = foreignKey.reference();
const target = Object.keys(reference.foreignTable).sort().join("+");
return `${reference.columns
.map((column) => column.name)
.sort()
.join(",")}->${target}(${reference.foreignColumns
.map((column) => column.name)
.sort()
.join(",")})`;
};
test("schemas should match", () => {
expectTypeOf<SqliteTables>().toEqualTypeOf<MysqlTables>();
expectTypeOf<MysqlTables>().toEqualTypeOf<SqliteTables>();
@@ -63,8 +84,7 @@ test("schemas should match", () => {
if (!("uniqueName" in sqliteColumn)) return;
if (!("primary" in sqliteColumn)) return;
// oxlint-disable-next-line import/namespace -- this test intentionally compares matching dynamic schema keys.
const mysqlTable = mysqlSchema[tableName];
const mysqlTable = mysqlTables[tableName];
const mysqlColumn = mysqlTable[columnName as keyof typeof mysqlTable] as object;
if (!("isUnique" in mysqlColumn)) return;
@@ -85,58 +105,65 @@ test("schemas should match", () => {
).toEqual(mysqlColumn.primary);
});
// oxlint-disable-next-line import/namespace -- this test intentionally compares matching dynamic schema keys.
const mysqlTable = mysqlSchema[tableName];
const sqliteForeignKeys = sqliteTable[Symbol.for("drizzle:SQLiteInlineForeignKeys") as keyof typeof sqliteTable] as
| SqliteForeignKey[]
| undefined;
const mysqlForeignKeys = mysqlTable[Symbol.for("drizzle:MySqlInlineForeignKeys") as keyof typeof mysqlTable] as
| MysqlForeignKey[]
| undefined;
const mysqlTable = mysqlTables[tableName];
// Read through getTableConfig so keys declared at table level count too, not just inline
// `.references()` ones. MySQL needs a table-level declaration wherever drizzle's derived name
// would exceed its 64-character identifier limit.
if (!is(sqliteTable, SQLiteTable) || !is(mysqlTable, MySqlTable)) return;
if (!sqliteForeignKeys && !mysqlForeignKeys) return;
const sqliteForeignKeys = getSqliteTableConfig(sqliteTable).foreignKeys as SqliteForeignKey[];
const mysqlForeignKeys = getMysqlTableConfig(mysqlTable).foreignKeys as MysqlForeignKey[];
expect(mysqlForeignKeys, `mysql foreign key for ${tableName} to be defined`).toBeDefined();
expect(sqliteForeignKeys, `sqlite foreign key for ${tableName} to be defined`).toBeDefined();
if (sqliteForeignKeys.length === 0 && mysqlForeignKeys.length === 0) return;
expect(
sqliteForeignKeys!.length,
sqliteForeignKeys.length,
`expect number of foreign keys in table ${tableName} to be the same for both schemas`,
).toEqual(mysqlForeignKeys!.length);
).toEqual(mysqlForeignKeys.length);
sqliteForeignKeys?.forEach((sqliteForeignKey) => {
sqliteForeignKey.getName();
const mysqlForeignKey = mysqlForeignKeys!.find((key) => key.getName() === sqliteForeignKey.getName());
expect(
mysqlForeignKey,
`expect foreign key ${sqliteForeignKey.getName()} to be defined in mysql schema`,
).toBeDefined();
sqliteForeignKeys.forEach((sqliteForeignKey) => {
const describeKey = `${tableName}.${referenceSignature(sqliteForeignKey)}`;
// Matched on what the key actually does rather than on its generated name, because the name
// is allowed to differ per dialect.
const mysqlForeignKey = mysqlForeignKeys.find(
(key) => referenceSignature(key) === referenceSignature(sqliteForeignKey),
);
expect(mysqlForeignKey, `expect foreign key ${describeKey} to be defined in mysql schema`).toBeDefined();
expect(
sqliteForeignKey.onDelete,
`expect foreign key (${sqliteForeignKey.getName()}) onDelete to be the same for both schemas`,
`expect foreign key (${describeKey}) onDelete to be the same for both schemas`,
).toEqual(mysqlForeignKey!.onDelete);
expect(
sqliteForeignKey.onUpdate,
`expect foreign key (${sqliteForeignKey.getName()}) onUpdate to be the same for both schemas`,
`expect foreign key (${describeKey}) onUpdate to be the same for both schemas`,
).toEqual(mysqlForeignKey!.onUpdate);
sqliteForeignKey.reference().foreignColumns.forEach((column) => {
expect(
mysqlForeignKey!.reference().foreignColumns.map((column) => column.name),
`expect foreign key (${sqliteForeignKey.getName()}) columns to be the same for both schemas`,
).toContainEqual(column.name);
});
expect(
Object.keys(sqliteForeignKey.reference().foreignTable),
`expect foreign key (${sqliteForeignKey.getName()}) table to be the same for both schemas`,
).toEqual(Object.keys(mysqlForeignKey!.reference().foreignTable));
});
});
});
test("mysql identifiers stay within the 64 character limit", () => {
// MySQL rejects longer identifiers with error 1059, and drizzle derives foreign key names from
// the table and column names, so a long pairing silently produces a migration that cannot apply.
const tooLong: string[] = [];
objectEntries(mysqlSchema).forEach(([, table]) => {
if (!is(table, MySqlTable)) return;
const config = getMysqlTableConfig(table);
if (config.name.length > 64) tooLong.push(config.name);
for (const foreignKey of config.foreignKeys) {
if (foreignKey.getName().length > 64) tooLong.push(foreignKey.getName());
}
for (const index of config.indexes) {
if (index.config.name.length > 64) tooLong.push(index.config.name);
}
});
expect(tooLong).toEqual([]);
});
test("schemas should match for postgresql", () => {
expectTypeOf<SqliteTables>().toEqualTypeOf<PostgresqlTables>();
expectTypeOf<PostgresqlTables>().toEqualTypeOf<SqliteTables>();
@@ -146,8 +173,7 @@ test("schemas should match for postgresql", () => {
objectEntries(sqliteSchema).forEach(([tableName, sqliteTable]) => {
// keys of sqliteSchema and postgresqlSchema are the same, so we can safely use tableName as key
// skipcq: JS-E1007
// oxlint-disable-next-line import/namespace -- this test intentionally compares matching dynamic schema keys.
const postgresqlTable = postgresqlSchema[tableName];
const postgresqlTable = postgresqlTables[tableName];
Object.entries(sqliteTable).forEach(([columnName, sqliteColumn]: [string, object]) => {
if (!("isUnique" in sqliteColumn)) return;
if (!("uniqueName" in sqliteColumn)) return;