refactor: remove display_image tool and update related components to streamline image handling

surfsense_web/components/assistant-ui/assistant-message.tsx

@@ -18,6 +18,7 @@ import { CommentPanelContainer } from "@/components/chat-comments/comment-panel-
 import { CommentSheet } from "@/components/chat-comments/comment-sheet/comment-sheet";
 import { CreateConfluencePageToolUI, DeleteConfluencePageToolUI, UpdateConfluencePageToolUI } from "@/components/tool-ui/confluence";
 import { DisplayImageToolUI } from "@/components/tool-ui/display-image";
+import { GenerateImageToolUI } from "@/components/tool-ui/generate-image";
 import { GeneratePodcastToolUI } from "@/components/tool-ui/generate-podcast";
 import { GenerateReportToolUI } from "@/components/tool-ui/generate-report";
 import { GenerateVideoPresentationToolUI } from "@/components/tool-ui/video-presentation";
@@ -60,6 +61,7 @@ const AssistantMessageInner: FC = () => {
 								link_preview: LinkPreviewToolUI,
 								multi_link_preview: MultiLinkPreviewToolUI,
 								display_image: DisplayImageToolUI,
+								generate_image: GenerateImageToolUI,
 								scrape_webpage: ScrapeWebpageToolUI,
 								save_memory: SaveMemoryToolUI,
 								recall_memory: RecallMemoryToolUI,

surfsense_web/components/assistant-ui/image.tsx

@@ -0,0 +1,268 @@
+"use client";
+
+import {
+  memo,
+  useState,
+  useEffect,
+  useRef,
+  type PropsWithChildren,
+} from "react";
+import { createPortal } from "react-dom";
+import { cva, type VariantProps } from "class-variance-authority";
+import { ImageIcon, ImageOffIcon } from "lucide-react";
+import type { ImageMessagePartComponent } from "@assistant-ui/react";
+import { cn } from "@/lib/utils";
+
+const imageVariants = cva(
+  "aui-image-root relative overflow-hidden rounded-lg",
+  {
+    variants: {
+      variant: {
+        outline: "border border-border",
+        ghost: "",
+        muted: "bg-muted/50",
+      },
+      size: {
+        sm: "max-w-64",
+        default: "max-w-96",
+        lg: "max-w-[512px]",
+        full: "w-full",
+      },
+    },
+    defaultVariants: {
+      variant: "outline",
+      size: "default",
+    },
+  },
+);
+
+export type ImageRootProps = React.ComponentProps<"div"> &
+  VariantProps<typeof imageVariants>;
+
+function ImageRoot({
+  className,
+  variant,
+  size,
+  children,
+  ...props
+}: ImageRootProps) {
+  return (
+    <div
+      data-slot="image-root"
+      data-variant={variant}
+      data-size={size}
+      className={cn(imageVariants({ variant, size, className }))}
+      {...props}
+    >
+      {children}
+    </div>
+  );
+}
+
+type ImagePreviewProps = Omit<React.ComponentProps<"img">, "children"> & {
+  containerClassName?: string;
+};
+
+function ImagePreview({
+  className,
+  containerClassName,
+  onLoad,
+  onError,
+  alt = "Image content",
+  src,
+  ...props
+}: ImagePreviewProps) {
+  const imgRef = useRef<HTMLImageElement>(null);
+  const [loadedSrc, setLoadedSrc] = useState<string | undefined>(undefined);
+  const [errorSrc, setErrorSrc] = useState<string | undefined>(undefined);
+
+  const loaded = loadedSrc === src;
+  const error = errorSrc === src;
+
+  useEffect(() => {
+    if (
+      typeof src === "string" &&
+      imgRef.current?.complete &&
+      imgRef.current.naturalWidth > 0
+    ) {
+      setLoadedSrc(src);
+    }
+  }, [src]);
+
+  return (
+    <div
+      data-slot="image-preview"
+      className={cn("relative min-h-32", containerClassName)}
+    >
+      {!loaded && !error && (
+        <div
+          data-slot="image-preview-loading"
+          className="absolute inset-0 flex items-center justify-center bg-muted/50"
+        >
+          <ImageIcon className="size-8 animate-pulse text-muted-foreground" />
+        </div>
+      )}
+      {error ? (
+        <div
+          data-slot="image-preview-error"
+          className="flex min-h-32 items-center justify-center bg-muted/50 p-4"
+        >
+          <ImageOffIcon className="size-8 text-muted-foreground" />
+        </div>
+      ) : (
+        // biome-ignore lint/performance/noImgElement: intentional for dynamic external URLs
+        <img
+          ref={imgRef}
+          src={src}
+          alt={alt}
+          className={cn(
+            "block h-auto w-full object-contain",
+            !loaded && "invisible",
+            className,
+          )}
+          onLoad={(e) => {
+            if (typeof src === "string") setLoadedSrc(src);
+            onLoad?.(e);
+          }}
+          onError={(e) => {
+            if (typeof src === "string") setErrorSrc(src);
+            onError?.(e);
+          }}
+          {...props}
+        />
+      )}
+    </div>
+  );
+}
+
+function ImageFilename({
+  className,
+  children,
+  ...props
+}: React.ComponentProps<"span">) {
+  if (!children) return null;
+
+  return (
+    <span
+      data-slot="image-filename"
+      className={cn(
+        "block truncate px-2 py-1.5 text-muted-foreground text-xs",
+        className,
+      )}
+      {...props}
+    >
+      {children}
+    </span>
+  );
+}
+
+type ImageZoomProps = PropsWithChildren<{
+  src: string;
+  alt?: string;
+}>;
+
+function ImageZoom({ src, alt = "Image preview", children }: ImageZoomProps) {
+  const [isMounted, setIsMounted] = useState(false);
+  const [isOpen, setIsOpen] = useState(false);
+
+  useEffect(() => {
+    setIsMounted(true);
+  }, []);
+
+  const handleOpen = () => setIsOpen(true);
+  const handleClose = () => setIsOpen(false);
+
+  useEffect(() => {
+    if (!isOpen) return;
+    const handleKeyDown = (e: KeyboardEvent) => {
+      if (e.key === "Escape") setIsOpen(false);
+    };
+    document.addEventListener("keydown", handleKeyDown);
+    return () => document.removeEventListener("keydown", handleKeyDown);
+  }, [isOpen]);
+
+  useEffect(() => {
+    if (!isOpen) return;
+    const originalOverflow = document.body.style.overflow;
+    document.body.style.overflow = "hidden";
+    return () => {
+      document.body.style.overflow = originalOverflow;
+    };
+  }, [isOpen]);
+
+  return (
+    <>
+      <button
+        type="button"
+        onClick={handleOpen}
+        className="aui-image-zoom-trigger cursor-zoom-in border-0 bg-transparent p-0 text-left"
+        aria-label="Click to zoom image"
+      >
+        {children}
+      </button>
+      {isMounted &&
+        isOpen &&
+        createPortal(
+          <button
+            type="button"
+            data-slot="image-zoom-overlay"
+            className="aui-image-zoom-overlay fade-in fixed inset-0 z-50 flex animate-in cursor-zoom-out items-center justify-center border-0 bg-black/80 p-0 duration-200"
+            onClick={handleClose}
+            aria-label="Close zoomed image"
+          >
+            {/** biome-ignore lint/performance/noImgElement: <explanation> */}
+            <img
+              data-slot="image-zoom-content"
+              src={src}
+              alt={alt}
+              className="aui-image-zoom-content fade-in zoom-in-95 max-h-[90vh] max-w-[90vw] animate-in object-contain duration-200"
+              onClick={(e) => {
+                e.stopPropagation();
+                handleClose();
+              }}
+              onKeyDown={(e) => {
+                if (e.key === "Enter") {
+                  e.stopPropagation();
+                  handleClose();
+                }
+              }}
+            />
+          </button>,
+          document.body,
+        )}
+    </>
+  );
+}
+
+const ImageImpl: ImageMessagePartComponent = ({ image, filename }) => {
+  return (
+    <ImageRoot>
+      <ImageZoom src={image} alt={filename || "Image content"}>
+        <ImagePreview src={image} alt={filename || "Image content"} />
+      </ImageZoom>
+      <ImageFilename>{filename}</ImageFilename>
+    </ImageRoot>
+  );
+};
+
+const Image = memo(ImageImpl) as unknown as ImageMessagePartComponent & {
+  Root: typeof ImageRoot;
+  Preview: typeof ImagePreview;
+  Filename: typeof ImageFilename;
+  Zoom: typeof ImageZoom;
+};
+
+Image.displayName = "Image";
+Image.Root = ImageRoot;
+Image.Preview = ImagePreview;
+Image.Filename = ImageFilename;
+Image.Zoom = ImageZoom;
+
+export {
+  Image,
+  ImageRoot,
+  ImagePreview,
+  ImageFilename,
+  ImageZoom,
+  imageVariants,
+};

surfsense_web/components/assistant-ui/markdown-text.tsx

@@ -8,8 +8,9 @@ import {
 	unstable_memoizeMarkdownComponents as memoizeMarkdownComponents,
 	useIsMarkdownCodeBlock,
 } from "@assistant-ui/react-markdown";
-import { CheckIcon, CopyIcon } from "lucide-react";
+import { CheckIcon, CopyIcon, ExternalLinkIcon } from "lucide-react";
 import { type FC, memo, type ReactNode, useState } from "react";
+import { ImagePreview, ImageRoot, ImageZoom } from "@/components/assistant-ui/image";
 import rehypeKatex from "rehype-katex";
 import remarkGfm from "remark-gfm";
 import remarkMath from "remark-math";
@@ -188,17 +189,17 @@ const useCopyToClipboard = ({ copiedDuration = 3000 }: { copiedDuration?: number
 function processChildrenWithCitations(children: ReactNode): ReactNode {
 	if (typeof children === "string") {
 		const parsed = parseTextWithCitations(children);
-		return parsed.length === 1 && typeof parsed[0] === "string" ? children : <>{parsed}</>;
+		return parsed.length === 1 && typeof parsed[0] === "string" ? children : parsed;
 	}
 
 	if (Array.isArray(children)) {
-		return children.map((child, index) => {
+		return children.map((child) => {
 			if (typeof child === "string") {
 				const parsed = parseTextWithCitations(child);
 				return parsed.length === 1 && typeof parsed[0] === "string" ? (
 					child
 				) : (
-					<span key={index}>{parsed}</span>
+					<span key={child}>{parsed}</span>
 				);
 			}
 			return child;
@@ -208,6 +209,56 @@ function processChildrenWithCitations(children: ReactNode): ReactNode {
 	return children;
 }
 
+function extractDomain(url: string): string {
+	try {
+		const parsed = new URL(url);
+		return parsed.hostname.replace(/^www\./, "");
+	} catch {
+		return "";
+	}
+}
+
+function MarkdownImage({ src, alt }: { src?: string; alt?: string }) {
+	if (!src) return null;
+
+	const domain = extractDomain(src);
+
+	return (
+		<div className="my-4 w-fit max-w-lg overflow-hidden rounded-2xl border bg-muted/30 select-none">
+			<ImageRoot variant="ghost" size="full">
+				<ImageZoom src={src} alt={alt || "Image"}>
+					<ImagePreview
+						src={src}
+						alt={alt || "Image"}
+						className="max-h-[20rem] w-auto max-w-full object-contain"
+					/>
+				</ImageZoom>
+			</ImageRoot>
+
+			<div className="flex items-center justify-between px-5 py-3">
+				<div className="min-w-0 flex-1">
+					{alt && alt !== "Image" && (
+						<p className="text-sm font-semibold text-foreground line-clamp-2">{alt}</p>
+					)}
+					{domain && (
+						<p className="text-xs text-muted-foreground mt-0.5 truncate">{domain}</p>
+					)}
+				</div>
+				<a
+					href={src}
+					target="_blank"
+					rel="noopener noreferrer"
+					className="ml-3 shrink-0 inline-flex items-center gap-1.5 rounded-lg border px-3 py-1.5 text-xs font-medium text-foreground transition-colors hover:bg-muted"
+					onClick={(e) => e.stopPropagation()}
+				>
+					Open
+					<ExternalLinkIcon className="size-3" />
+				</a>
+			</div>
+		</div>
+	);
+}
+
 const defaultComponents = memoizeMarkdownComponents({
 	h1: ({ className, children, ...props }) => (
 		<h1
@@ -371,5 +422,6 @@ const defaultComponents = memoizeMarkdownComponents({
 			{processChildrenWithCitations(children)}
 		</em>
 	),
+	img: ({ src, alt }) => <MarkdownImage src={typeof src === "string" ? src : undefined} alt={alt} />,
 	CodeHeader,
 });

surfsense_web/components/assistant-ui/thread.tsx

@@ -1058,7 +1058,7 @@ const TOOL_GROUPS: ToolGroup[] = [
 	},
 	{
 		label: "Generate",
-		tools: ["generate_podcast", "generate_video_presentation", "generate_report", "generate_image", "display_image"],
+		tools: ["generate_podcast", "generate_video_presentation", "generate_report", "generate_image"],
 	},
 	{
 		label: "Memory",