fix: harden assistant integration
This commit is contained in:
1 parent
10ef871777
commit
5db95f2def
14 files changed
+349
-189
No files matched your search
@@ -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;
|
||||
@@ -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);
|
||||
|
||||
@@ -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
@@ -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, {
|
||||
|
||||
@@ -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, {
|
||||
|
||||
@@ -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, {
|
||||
|
||||
@@ -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;
|
||||
|
||||
Reference in new issue
Block a user