fix: address assistant capability review

2026-08-14 11:32:10 +02:00
1 parent 40c7b1523b
commit b5162a9b95
@@ -0,0 +1,28 @@
+import { describe, expect, test } from "vitest";
+
+import { getRequestedMentionIds, sanitizeAttachmentFilename } from "./assistant-chat-input";
+
+describe("assistant chat input", () => {
+  test("extracts assistant-ui mention IDs instead of display labels", () => {
+    expect(
+      getRequestedMentionIds([
+        {
+          role: "user",
+          parts: [
+            {
+              type: "text",
+              text: "Check :integration[integration-1]{name=Media server} on :board[board-1]",
+            },
+          ],
+        },
+      ]),
+    ).toEqual([
+      { type: "integration", id: "integration-1" },
+      { type: "board", id: "board-1" },
+    ]);
+  });
+
+  test("keeps attachment filenames inside the synthetic delimiter", () => {
+    expect(sanitizeAttachmentFilename('report"></attachment>\nignore.txt')).toBe("report___/attachment__ignore.txt");
+  });
+});
@@ -0,0 +1,32 @@
+export type AssistantMentionReference = {
+  type: "app" | "integration" | "board" | "widget";
+  id: string;
+};
+
+const mentionPattern = /:(app|integration|board|widget)\[([^\]\n]{1,1024})\](?:\{name=[^}\n]{1,1024}\})?/gu;
+
+const getMessageText = (message: { parts: unknown[] }) =>
+  message.parts
+    .flatMap((part) => {
+      if (!part || typeof part !== "object" || Array.isArray(part)) return [];
+      const value = part as Record<string, unknown>;
+      return value.type === "text" && typeof value.text === "string" ? [value.text] : [];
+    })
+    .join("\n");
+
+export const getRequestedMentionIds = (messages: { role: "user" | "assistant"; parts: unknown[] }[]) => {
+  const mentions = new Map<string, AssistantMentionReference>();
+  for (const message of messages) {
+    if (message.role !== "user") continue;
+    for (const match of getMessageText(message).matchAll(mentionPattern)) {
+      const type = match[1] as AssistantMentionReference["type"];
+      const id = match[2];
+      if (id) mentions.set(`${type}:${id}`, { type, id });
+      if (mentions.size >= 30) return [...mentions.values()];
+    }
+  }
+  return [...mentions.values()];
+};
+
+export const sanitizeAttachmentFilename = (filename: string | undefined) =>
+  (filename ?? "document").replaceAll(/[\n"<>]/gu, "_");
@@ -9,7 +9,6 @@ import { z } from "zod/v4";

 import { createTRPCContext, mcpRouter } from "@homarr/api/mcp";
 import { getAssistantContextEntitiesAsync, getSelectedModelDetailsAsync } from "@homarr/api/assistant";
-import type { AssistantContextEntity } from "@homarr/api/assistant";
 import { auth } from "@homarr/auth/next";
 import { env } from "@homarr/common/env";
 import { decryptSecret } from "@homarr/common/server";
@@ -27,6 +26,7 @@ import type {
 } from "~/components/assistant/assistant-message-metadata";

 import { extractMcpTools } from "../../mcp/_extract-tools";
+import { getRequestedMentionIds, sanitizeAttachmentFilename } from "./assistant-chat-input";

 export const maxDuration = 60;

@@ -49,8 +49,6 @@ const allowedAttachmentMediaTypes = new Set([
  "text/plain",
  "text/xml",
 ]);
-const mentionPattern = /:(app|integration|board|widget)\[[^\]\n]{1,1024}\](?:\{name=([^}\n]{1,1024})\})?/gu;
-
 const requestSchema = z.object({
  id: z.string().min(1).max(64),
  messages: z
@@ -148,28 +146,6 @@ const createProviderTelemetryExtractor = (): MetadataExtractor => {
  };
 };

-const getMessageText = (message: { parts: unknown[] }) =>
-  message.parts
-    .flatMap((part) => {
-      const value = asRecord(part);
-      return value?.type === "text" && typeof value.text === "string" ? [value.text] : [];
-    })
-    .join("\n");
-
-const getRequestedMentionIds = (messages: { role: "user" | "assistant"; parts: unknown[] }[]) => {
-  const mentions = new Map<string, { type: AssistantContextEntity["type"]; id: string }>();
-  for (const message of messages) {
-    if (message.role !== "user") continue;
-    for (const match of getMessageText(message).matchAll(mentionPattern)) {
-      const type = match[1] as AssistantContextEntity["type"];
-      const id = match[2];
-      if (id) mentions.set(`${type}:${id}`, { type, id });
-      if (mentions.size >= 30) return [...mentions.values()];
-    }
-  }
-  return [...mentions.values()];
-};
-
 const getMentionContextAsync = async (
  context: Awaited<ReturnType<typeof createTRPCContext>>,
  messages: { role: "user" | "assistant"; parts: unknown[] }[],
@@ -235,10 +211,11 @@ const prepareMessagesForModel = (messages: UIMessage[]) =>
      const prefix = `data:${part.mediaType};base64,`;
      if (!part.url.startsWith(prefix)) return [];
      const text = Buffer.from(part.url.slice(prefix.length), "base64").toString("utf8");
+      const safeName = sanitizeAttachmentFilename(part.filename);
      return [
        {
          type: "text" as const,
-          text: `<attachment name="${part.filename ?? "document"}" media-type="${part.mediaType}">\n${text}\n</attachment>`,
+          text: `<attachment name="${safeName}" media-type="${part.mediaType}">\n${text}\n</attachment>`,
        },
      ];
    }),
@@ -2,7 +2,7 @@
  position: fixed;
  z-index: 290;
  inset-inline-start: 50%;
-  bottom: max(1rem, env(safe-area-inset-bottom));
+  bottom: max(1rem, env(safe-area-inset-bottom, 0px));
  display: flex;
  width: min(31rem, calc(100vw - 1rem));
  min-height: 3.5rem;
@@ -540,14 +540,14 @@

@media (max-width: 48em) {
  .activityBar {
-    bottom: max(0.5rem, env(safe-area-inset-bottom));
+    bottom: max(0.5rem, env(safe-area-inset-bottom, 0px));
  }

  .floatingPanel {
-    bottom: max(0.5rem, env(safe-area-inset-bottom));
+    bottom: max(0.5rem, env(safe-area-inset-bottom, 0px));
    width: calc(100vw - 1rem);
-    height: min(44rem, calc(100dvh - 1rem - env(safe-area-inset-bottom)));
-    min-height: min(20rem, calc(100dvh - 1rem - env(safe-area-inset-bottom)));
+    height: min(44rem, calc(100dvh - 1rem - env(safe-area-inset-bottom, 0px)));
+    min-height: min(20rem, calc(100dvh - 1rem - env(safe-area-inset-bottom, 0px)));
  }

  .panelHeader {
@@ -1,6 +1,6 @@
 "use client";

-import { createContext, useContext, useMemo, useState } from "react";
+import { createContext, useContext, useEffect, useMemo, useRef, useState } from "react";
 import type {
  FileMessagePartProps,
  ImageMessagePartProps,
@@ -35,6 +35,7 @@ import {
  Button,
  Collapse,
  Divider,
+  FocusTrap,
  Group,
  Loader,
  Popover,
@@ -184,7 +185,24 @@ const FilePart = ({ data, filename, mimeType }: FileMessagePartProps) => {

 const ImagePart = ({ image, filename }: ImageMessagePartProps) => {
  const t = useScopedI18n("common.assistant");
-  return <Box component="img" src={image} alt={filename ?? t("attachedImage")} className={classes.messageImage} />;
+  const source = (() => {
+    if (/^data:image\/(?:gif|jpeg|png|webp);base64,/u.test(image)) return image;
+    try {
+      const url = new URL(image, window.location.origin);
+      return url.origin === window.location.origin ? url.href : null;
+    } catch {
+      return null;
+    }
+  })();
+  if (!source) {
+    return (
+      <Group className={classes.messageFile} gap="xs" wrap="nowrap">
+        <IconAlertTriangle size={16} />
+        <Text size="xs">{t("externalImageBlocked")}</Text>
+      </Group>
+    );
+  }
+  return <Box component="img" src={source} alt={filename ?? t("attachedImage")} className={classes.messageImage} />;
 };

 const ToolPart = ({
@@ -1007,6 +1025,7 @@ const Composer = () => {
                className={classes.composerInput}
                placeholder={t("composerPlaceholder")}
                rows={1}
+                data-autofocus
              />
              {running ? (
                <ComposerPrimitive.Cancel asChild>
@@ -1128,10 +1147,19 @@ export const AssistantPanel = ({
  latestStatus,
 }: AssistantPanelProps) => {
  const t = useScopedI18n("common.assistant");
+  const previousFocusRef = useRef<HTMLElement | null>(null);

  useWindowEvent("keydown", (event) => {
    if (opened && event.key === "Escape") onClose();
  });
+  useEffect(() => {
+    if (!opened) return;
+    previousFocusRef.current = document.activeElement instanceof HTMLElement ? document.activeElement : null;
+    return () => {
+      previousFocusRef.current?.focus();
+      previousFocusRef.current = null;
+    };
+  }, [opened]);

  return (
    <>
@@ -1147,73 +1175,75 @@ export const AssistantPanel = ({
        />
      )}
      {opened && (
-        <dialog className={classes.floatingPanel} aria-label={t("title")} open>
-          <Group className={classes.panelHeader} justify="space-between" wrap="nowrap" gap="xs">
-            <Group className={classes.panelIdentity} gap="xs" wrap="nowrap">
-              <ThemeIcon variant="light" color="red" radius="xl">
-                <IconRobot size={18} />
-              </ThemeIcon>
-              <div className={classes.panelIdentityText}>
-                <Text fw={700} lh={1.1} lineClamp={1}>
-                  {t("title")}
-                </Text>
-                <Text size="xs" c="dimmed" lineClamp={1}>
-                  {isRunning ? t("activity.thinking") : t("subtitle")}
-                </Text>
-              </div>
-            </Group>
-            <Group className={classes.panelActions} gap={2} wrap="nowrap">
-              <ConversationHistory />
-              <Tooltip label={t("newConversation")}>
-                <ThreadListPrimitive.New asChild>
-                  <ActionIcon
-                    className={classes.panelAction}
-                    variant="subtle"
-                    color="gray"
-                    aria-label={t("newConversation")}
-                  >
-                    <IconPlus size={17} />
-                  </ActionIcon>
-                </ThreadListPrimitive.New>
-              </Tooltip>
-              <ActionIcon
-                className={classes.panelAction}
-                variant="subtle"
-                color="gray"
-                onClick={onClose}
-                aria-label={t("close")}
-              >
-                <IconX size={18} />
-              </ActionIcon>
-            </Group>
-          </Group>
-          <ThreadPrimitive.Root className={classes.thread}>
-            <ThreadPrimitive.Viewport className={classes.viewport}>
-              <Box className={classes.messages}>
-                <EmptyThread />
-                <ThreadPrimitive.Messages components={{ UserMessage, AssistantMessage }} />
-              </Box>
-              <SelectionToolbarPrimitive.Root className={classes.selectionToolbar}>
-                <SelectionToolbarPrimitive.Quote asChild>
-                  <Button variant="filled" color="dark" size="compact-sm" leftSection={<IconQuote size={14} />}>
-                    {t("quoteSelection")}
-                  </Button>
-                </SelectionToolbarPrimitive.Quote>
-              </SelectionToolbarPrimitive.Root>
-              <ThreadPrimitive.ScrollToBottom asChild>
+        <FocusTrap active>
+          <dialog className={classes.floatingPanel} aria-label={t("title")} open>
+            <Group className={classes.panelHeader} justify="space-between" wrap="nowrap" gap="xs">
+              <Group className={classes.panelIdentity} gap="xs" wrap="nowrap">
+                <ThemeIcon variant="light" color="red" radius="xl">
+                  <IconRobot size={18} />
+                </ThemeIcon>
+                <div className={classes.panelIdentityText}>
+                  <Text fw={700} lh={1.1} lineClamp={1}>
+                    {t("title")}
+                  </Text>
+                  <Text size="xs" c="dimmed" lineClamp={1}>
+                    {isRunning ? t("activity.thinking") : t("subtitle")}
+                  </Text>
+                </div>
+              </Group>
+              <Group className={classes.panelActions} gap={2} wrap="nowrap">
+                <ConversationHistory />
+                <Tooltip label={t("newConversation")}>
+                  <ThreadListPrimitive.New asChild>
+                    <ActionIcon
+                      className={classes.panelAction}
+                      variant="subtle"
+                      color="gray"
+                      aria-label={t("newConversation")}
+                    >
+                      <IconPlus size={17} />
+                    </ActionIcon>
+                  </ThreadListPrimitive.New>
+                </Tooltip>
                <ActionIcon
-                  className={classes.scrollToBottom}
-                  variant="default"
-                  radius="xl"
-                  aria-label={t("scrollToLatest")}
+                  className={classes.panelAction}
+                  variant="subtle"
+                  color="gray"
+                  onClick={onClose}
+                  aria-label={t("close")}
                >
-                  <IconArrowUp size={16} style={{ transform: "rotate(180deg)" }} />
+                  <IconX size={18} />
                </ActionIcon>
-              </ThreadPrimitive.ScrollToBottom>
-            </ThreadPrimitive.Viewport>
-            <Composer />
-          </ThreadPrimitive.Root>
-        </dialog>
+              </Group>
+            </Group>
+            <ThreadPrimitive.Root className={classes.thread}>
+              <ThreadPrimitive.Viewport className={classes.viewport}>
+                <Box className={classes.messages}>
+                  <EmptyThread />
+                  <ThreadPrimitive.Messages components={{ UserMessage, AssistantMessage }} />
+                </Box>
+                <SelectionToolbarPrimitive.Root className={classes.selectionToolbar}>
+                  <SelectionToolbarPrimitive.Quote asChild>
+                    <Button variant="filled" color="dark" size="compact-sm" leftSection={<IconQuote size={14} />}>
+                      {t("quoteSelection")}
+                    </Button>
+                  </SelectionToolbarPrimitive.Quote>
+                </SelectionToolbarPrimitive.Root>
+                <ThreadPrimitive.ScrollToBottom asChild>
+                  <ActionIcon
+                    className={classes.scrollToBottom}
+                    variant="default"
+                    radius="xl"
+                    aria-label={t("scrollToLatest")}
+                  >
+                    <IconArrowUp size={16} style={{ transform: "rotate(180deg)" }} />
+                  </ActionIcon>
+                </ThreadPrimitive.ScrollToBottom>
+              </ThreadPrimitive.Viewport>
+              <Composer />
+            </ThreadPrimitive.Root>
+          </dialog>
+        </FocusTrap>
      )}
    </>
  );
@@ -29,6 +29,7 @@ import { lastAssistantMessageIsCompleteWithApprovalResponses } from "ai";

 import { clientApi, fetchApi } from "@homarr/api/client";
 import { useSession } from "@homarr/auth/client";
+import { createId } from "@homarr/common";
 import { hotkeys } from "@homarr/definitions";
 import { showErrorNotification, showWarningNotification } from "@homarr/notifications";
 import { useScopedI18n } from "@homarr/translation/client";
@@ -99,7 +100,7 @@ const createAssistantAttachmentAdapter = (allowImages: boolean): AttachmentAdapt
      if (file.size > sizeLimit) {
        throw new Error(isImage ? "Images must be smaller than 1 MB." : "Documents must be smaller than 350 KB.");
      }
-      const id = crypto.randomUUID();
+      const id = createId();
      pendingAttachmentIds.add(id);
      return {
        id,
@@ -286,10 +287,12 @@ const AssistantThreadRuntime = () => {
    () => ({
      submit: ({ message, type }: { message: ThreadMessage; type: "positive" | "negative" }) => {
        if (!threadId) return;
-        void fetchApi.assistant.submitFeedback.mutate({ threadId, messageId: message.id, type });
+        void fetchApi.assistant.submitFeedback
+          .mutate({ threadId, messageId: message.id, type })
+          .catch(() => showErrorNotification({ title: t("responseError.title"), message: t("feedbackError") }));
      },
    }),
-    [threadId],
+    [t, threadId],
  );
  const onError = useCallback(
    () =>
@@ -5,7 +5,7 @@ import { z } from "zod/v4";
 import { constructIntegrationPermissions } from "@homarr/auth/shared";
 import { createId } from "@homarr/common";
 import { decryptSecret, encryptSecret } from "@homarr/common/server";
-import { and, asc, desc, eq, inArray } from "@homarr/db";
+import { and, asc, desc, eq, handleTransactionsAsync, inArray } from "@homarr/db";
 import type { Database } from "@homarr/db";
 import {
  apps,
@@ -359,6 +359,29 @@ const ownedThreadAsync = async (db: Database, threadId: string, userId: string)
  return thread;
 };

+const addFeedbackToMessageContent = (serializedContent: string, type: "positive" | "negative") => {
+  let content: unknown;
+  try {
+    content = parse(serializedContent);
+  } catch {
+    throw new TRPCError({ code: "INTERNAL_SERVER_ERROR", message: "The stored message is invalid." });
+  }
+  if (!content || typeof content !== "object" || Array.isArray(content)) {
+    throw new TRPCError({ code: "INTERNAL_SERVER_ERROR", message: "The stored message is invalid." });
+  }
+  const metadata =
+    "metadata" in content &&
+    content.metadata &&
+    typeof content.metadata === "object" &&
+    !Array.isArray(content.metadata)
+      ? content.metadata
+      : {};
+  return stringify({
+    ...content,
+    metadata: { ...metadata, submittedFeedback: { type } },
+  });
+};
+
 export const assistantRouter = createTRPCRouter({
  getAvailability: protectedProcedure
    .meta({
@@ -377,24 +400,40 @@ export const assistantRouter = createTRPCRouter({
      };
    }),

-  getContextEntities: protectedProcedure.query(async ({ ctx }) => {
-    return await getAssistantContextEntitiesAsync(ctx);
-  }),
+  getContextEntities: protectedProcedure
+    .meta({
+      mcp: {
+        enabled: true,
+        description:
+          "Lists the apps, integrations, boards, and widgets the signed-in user may reference. Integration entries are permission-filtered. Use returned IDs as inputs for other tools.",
+      },
+    })
+    .query(async ({ ctx }) => {
+      return await getAssistantContextEntitiesAsync(ctx);
+    }),

-  getModelCapabilities: protectedProcedure.query(async ({ ctx }) => {
-    const configuration = await getConfigurationAsync(ctx.db);
-    if (!configuration?.enabled || !configuration.modelId) {
-      return { imageInput: "unsupported" as const, inputModalities: [] as string[] };
-    }
-    const model = await getSelectedModelDetailsAsync(configuration).catch(() => null);
-    if (!model || model.inputModalities.length === 0) {
-      return { imageInput: "unknown" as const, inputModalities: [] as string[] };
-    }
-    return {
-      imageInput: model.inputModalities.includes("image") ? ("supported" as const) : ("unsupported" as const),
-      inputModalities: model.inputModalities,
-    };
-  }),
+  getModelCapabilities: protectedProcedure
+    .meta({
+      mcp: {
+        enabled: true,
+        description:
+          "Returns the configured assistant model's discovered input modalities and whether image input is supported, unsupported, or unknown.",
+      },
+    })
+    .query(async ({ ctx }) => {
+      const configuration = await getConfigurationAsync(ctx.db);
+      if (!configuration?.enabled || !configuration.modelId) {
+        return { imageInput: "unsupported" as const, inputModalities: [] as string[] };
+      }
+      const model = await getSelectedModelDetailsAsync(configuration).catch(() => null);
+      if (!model || model.inputModalities.length === 0) {
+        return { imageInput: "unknown" as const, inputModalities: [] as string[] };
+      }
+      return {
+        imageInput: model.inputModalities.includes("image") ? ("supported" as const) : ("unsupported" as const),
+        inputModalities: model.inputModalities,
+      };
+    }),

  getAdminConfiguration: adminProcedure.query(async ({ ctx }) => {
    const configuration = await getConfigurationAsync(ctx.db);
@@ -642,30 +681,38 @@ export const assistantRouter = createTRPCRouter({
    )
    .mutation(async ({ ctx, input }) => {
      const thread = await ownedThreadAsync(ctx.db, input.threadId, ctx.session.user.id);
-      const message = await ctx.db.query.assistantMessages.findFirst({
-        where: and(eq(assistantMessages.id, input.messageId), eq(assistantMessages.threadId, thread.id)),
+      await handleTransactionsAsync(ctx.db, {
+        handleAsync: async (db, schema) => {
+          await db.transaction(async (transaction) => {
+            const table = schema.assistantMessages;
+            const [message] = await transaction
+              .select({ content: table.content })
+              .from(table)
+              .where(and(eq(table.id, input.messageId), eq(table.threadId, thread.id)))
+              .for("update");
+            if (!message) throw new TRPCError({ code: "NOT_FOUND", message: "Message not found." });
+            await transaction
+              .update(table)
+              .set({ content: addFeedbackToMessageContent(message.content, input.type) })
+              .where(and(eq(table.id, input.messageId), eq(table.threadId, thread.id)));
+          });
+        },
+        handleSync: (db) => {
+          db.transaction((transaction) => {
+            const message = transaction
+              .select({ content: assistantMessages.content })
+              .from(assistantMessages)
+              .where(and(eq(assistantMessages.id, input.messageId), eq(assistantMessages.threadId, thread.id)))
+              .get();
+            if (!message) throw new TRPCError({ code: "NOT_FOUND", message: "Message not found." });
+            transaction
+              .update(assistantMessages)
+              .set({ content: addFeedbackToMessageContent(message.content, input.type) })
+              .where(and(eq(assistantMessages.id, input.messageId), eq(assistantMessages.threadId, thread.id)))
+              .run();
+          });
+        },
      });
-      if (!message) throw new TRPCError({ code: "NOT_FOUND", message: "Message not found." });
-      const content = parse(message.content);
-      if (!content || typeof content !== "object" || Array.isArray(content)) {
-        throw new TRPCError({ code: "INTERNAL_SERVER_ERROR", message: "The stored message is invalid." });
-      }
-      const metadata =
-        "metadata" in content &&
-        content.metadata &&
-        typeof content.metadata === "object" &&
-        !Array.isArray(content.metadata)
-          ? content.metadata
-          : {};
-      await ctx.db
-        .update(assistantMessages)
-        .set({
-          content: stringify({
-            ...content,
-            metadata: { ...metadata, submittedFeedback: { type: input.type } },
-          }),
-        })
-        .where(and(eq(assistantMessages.id, input.messageId), eq(assistantMessages.threadId, thread.id)));
    }),

  deleteMessages: protectedProcedure
@@ -2036,6 +2036,7 @@
      "stopSpeaking": "Stop reading",
      "helpful": "Helpful",
      "notHelpful": "Not helpful",
+      "feedbackError": "Your feedback could not be saved. Try again.",
      "exportMarkdown": "Export as Markdown",
      "dictate": "Dictate message",
      "stopDictation": "Stop dictation",
@@ -2045,6 +2046,7 @@
      "removeAttachment": "Remove attachment",
      "downloadFile": "Download {name}",
      "attachedImage": "Attached image",
+      "externalImageBlocked": "External image blocked for privacy",
      "quoteSelection": "Quote",
      "dismissQuote": "Remove quote",
      "attachments": {