From 5db95f2defcc267a51972ac89eaad78a2bb82e36 Mon Sep 17 00:00:00 2001
From: ajnart <thomas@ajnart.fr>
Date: Tue, 28 Jul 2026 22:56:28 +0200
Subject: [PATCH] fix: harden assistant integration

---
 .../src/app/api/assistant/chat/route.ts       |  72 +++++++------
 .../assistant/assistant-panel.module.css      |   4 -
 .../components/assistant/assistant-panel.tsx  |  12 ++-
 .../assistant/assistant-provider.tsx          |  45 ++++----
 .../assistant/assistant-tool-contracts.ts     |  19 ++++
 packages/api/src/router/assistant.ts          |  25 +++--
 packages/api/src/test/mcp.spec.ts             |  65 ++++++++++-
 .../migrations/mysql/0044_add_assistant.sql   |   8 +-
 .../postgresql/0010_workable_shard.sql        |   8 +-
 .../migrations/sqlite/0042_bumpy_tattoo.sql   |   4 +
 packages/db/schema/mysql.ts                   |  54 ++++++----
 packages/db/schema/postgresql.ts              |  54 ++++++----
 packages/db/schema/sqlite.ts                  |  66 +++++++-----
 packages/db/test/schema.spec.ts               | 102 +++++++++++-------
 14 files changed, 349 insertions(+), 189 deletions(-)
 create mode 100644 apps/nextjs/src/components/assistant/assistant-tool-contracts.ts

diff --git a/apps/nextjs/src/app/api/assistant/chat/route.ts b/apps/nextjs/src/app/api/assistant/chat/route.ts
index e4dd9efe9..cb58701d2 100644
--- a/apps/nextjs/src/app/api/assistant/chat/route.ts
+++ b/apps/nextjs/src/app/api/assistant/chat/route.ts
@@ -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." },
diff --git a/apps/nextjs/src/components/assistant/assistant-panel.module.css b/apps/nextjs/src/components/assistant/assistant-panel.module.css
index 0707d41e2..ba3d26f71 100644
--- a/apps/nextjs/src/components/assistant/assistant-panel.module.css
+++ b/apps/nextjs/src/components/assistant/assistant-panel.module.css
@@ -1,7 +1,3 @@
-.drawer {
-  --assistant-panel-width: min(32rem, 100vw);
-}
-
 .content {
   height: 100dvh;
   background: var(--mantine-color-body);
diff --git a/apps/nextjs/src/components/assistant/assistant-panel.tsx b/apps/nextjs/src/components/assistant/assistant-panel.tsx
index a99a3a164..3aebdc1fe 100644
--- a/apps/nextjs/src/components/assistant/assistant-panel.tsx
+++ b/apps/nextjs/src/components/assistant/assistant-panel.tsx
@@ -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}>
diff --git a/apps/nextjs/src/components/assistant/assistant-provider.tsx b/apps/nextjs/src/components/assistant/assistant-provider.tsx
index 444d891ad..98b29f6b6 100644
--- a/apps/nextjs/src/components/assistant/assistant-provider.tsx
+++ b/apps/nextjs/src/components/assistant/assistant-provider.tsx
@@ -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 };
diff --git a/apps/nextjs/src/components/assistant/assistant-tool-contracts.ts b/apps/nextjs/src/components/assistant/assistant-tool-contracts.ts
new file mode 100644
index 000000000..9dccd7fe6
--- /dev/null
+++ b/apps/nextjs/src/components/assistant/assistant-tool-contracts.ts
@@ -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;
diff --git a/packages/api/src/router/assistant.ts b/packages/api/src/router/assistant.ts
index acf2eada4..990cf3c18 100644
--- a/packages/api/src/router/assistant.ts
+++ b/packages/api/src/router/assistant.ts
@@ -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);
diff --git a/packages/api/src/test/mcp.spec.ts b/packages/api/src/test/mcp.spec.ts
index 09e6c077e..14dfc3264 100644
--- a/packages/api/src/test/mcp.spec.ts
+++ b/packages/api/src/test/mcp.spec.ts
@@ -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" }),
+    );
   });
 });
diff --git a/packages/db/migrations/mysql/0044_add_assistant.sql b/packages/db/migrations/mysql/0044_add_assistant.sql
index 50897f051..f52838255 100644
--- a/packages/db/migrations/mysql/0044_add_assistant.sql
+++ b/packages/db/migrations/mysql/0044_add_assistant.sql
@@ -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`);
diff --git a/packages/db/migrations/postgresql/0010_workable_shard.sql b/packages/db/migrations/postgresql/0010_workable_shard.sql
index df037b8e9..706fd1da9 100644
--- a/packages/db/migrations/postgresql/0010_workable_shard.sql
+++ b/packages/db/migrations/postgresql/0010_workable_shard.sql
@@ -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;
\ No newline at end of file
+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");
diff --git a/packages/db/migrations/sqlite/0042_bumpy_tattoo.sql b/packages/db/migrations/sqlite/0042_bumpy_tattoo.sql
index 256d86886..d4c304d33 100644
--- a/packages/db/migrations/sqlite/0042_bumpy_tattoo.sql
+++ b/packages/db/migrations/sqlite/0042_bumpy_tattoo.sql
@@ -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`);
diff --git a/packages/db/schema/mysql.ts b/packages/db/schema/mysql.ts
index bfb38111d..760a03098 100644
--- a/packages/db/schema/mysql.ts
+++ b/packages/db/schema/mysql.ts
@@ -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, {
diff --git a/packages/db/schema/postgresql.ts b/packages/db/schema/postgresql.ts
index db1adecaf..199244a0e 100644
--- a/packages/db/schema/postgresql.ts
+++ b/packages/db/schema/postgresql.ts
@@ -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, {
diff --git a/packages/db/schema/sqlite.ts b/packages/db/schema/sqlite.ts
index 99f83d789..76b2dba91 100644
--- a/packages/db/schema/sqlite.ts
+++ b/packages/db/schema/sqlite.ts
@@ -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, {
diff --git a/packages/db/test/schema.spec.ts b/packages/db/test/schema.spec.ts
index 5e1eeab98..f5503e516 100644
--- a/packages/db/test/schema.spec.ts
+++ b/packages/db/test/schema.spec.ts
@@ -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;