From 15a81dbf41325e59e0c83128e2388a11381bdb38 Mon Sep 17 00:00:00 2001
From: Anish Sarkar <104695310+AnishSarkar22@users.noreply.github.com>
Date: Tue, 24 Mar 2026 23:03:00 +0530
Subject: [PATCH] feat: add GenerateImageToolUI component for rendering
 generated images with error handling and loading states

---
 .../components/tool-ui/generate-image.tsx     | 138 ++++++++++++++++++
 .../components/tool-ui/image/index.tsx        |  58 ++++----
 2 files changed, 166 insertions(+), 30 deletions(-)
 create mode 100644 surfsense_web/components/tool-ui/generate-image.tsx

diff --git a/surfsense_web/components/tool-ui/generate-image.tsx b/surfsense_web/components/tool-ui/generate-image.tsx
new file mode 100644
index 000000000..fd2e9992e
--- /dev/null
+++ b/surfsense_web/components/tool-ui/generate-image.tsx
@@ -0,0 +1,138 @@
+"use client";
+
+import type { ToolCallMessagePartProps } from "@assistant-ui/react";
+import { AlertCircleIcon, ImageIcon } from "lucide-react";
+import { z } from "zod";
+import {
+	Image,
+	ImageErrorBoundary,
+	ImageLoading,
+	parseSerializableImage,
+} from "@/components/tool-ui/image";
+
+const GenerateImageArgsSchema = z.object({
+	prompt: z.string(),
+	n: z.number().nullish(),
+});
+
+const GenerateImageResultSchema = z.object({
+	id: z.string(),
+	assetId: z.string(),
+	src: z.string(),
+	alt: z.string().nullish(),
+	title: z.string().nullish(),
+	description: z.string().nullish(),
+	domain: z.string().nullish(),
+	ratio: z.string().nullish(),
+	generated: z.boolean().nullish(),
+	prompt: z.string().nullish(),
+	image_count: z.number().nullish(),
+	error: z.string().nullish(),
+});
+
+type GenerateImageArgs = z.infer<typeof GenerateImageArgsSchema>;
+type GenerateImageResult = z.infer<typeof GenerateImageResultSchema>;
+
+function ImageErrorState({ prompt, error }: { prompt: string; error: string }) {
+	return (
+		<div className="my-4 overflow-hidden rounded-xl border border-destructive/20 bg-destructive/5 p-4 max-w-md">
+			<div className="flex items-center gap-4">
+				<div className="flex size-12 shrink-0 items-center justify-center rounded-lg bg-destructive/10">
+					<AlertCircleIcon className="size-6 text-destructive" />
+				</div>
+				<div className="flex-1 min-w-0">
+					<p className="font-medium text-destructive text-sm">Image generation failed</p>
+					<p className="text-muted-foreground text-xs mt-0.5 truncate">{prompt}</p>
+					<p className="text-muted-foreground text-xs mt-1">{error}</p>
+				</div>
+			</div>
+		</div>
+	);
+}
+
+function ImageCancelledState({ prompt }: { prompt: string }) {
+	return (
+		<div className="my-4 rounded-xl border border-muted p-4 text-muted-foreground max-w-md">
+			<p className="flex items-center gap-2">
+				<ImageIcon className="size-4" />
+				<span className="line-through truncate">Generate: {prompt}</span>
+			</p>
+		</div>
+	);
+}
+
+function ParsedImage({ result }: { result: unknown }) {
+	const image = parseSerializableImage(result);
+	return (
+		<Image
+			id={image.id}
+			assetId={image.assetId}
+			src={image.src}
+			alt={image.alt}
+			title={image.title ?? undefined}
+			description={image.description ?? undefined}
+			href={image.href ?? undefined}
+			domain={image.domain ?? undefined}
+			ratio={image.ratio ?? undefined}
+			source={image.source ?? undefined}
+			maxWidth="512px"
+		/>
+	);
+}
+
+/**
+ * Tool UI for generate_image — renders the generated image directly
+ * from the tool result directly.
+ */
+export const GenerateImageToolUI = ({ args, result, status }: ToolCallMessagePartProps<GenerateImageArgs, GenerateImageResult>) => {
+	const prompt = args.prompt || "Generating image...";
+
+	if (status.type === "running" || status.type === "requires-action") {
+		return (
+			<div className="my-4">
+				<ImageLoading title="Generating image" />
+			</div>
+		);
+	}
+
+	if (status.type === "incomplete") {
+		if (status.reason === "cancelled") {
+			return <ImageCancelledState prompt={prompt} />;
+		}
+		if (status.reason === "error") {
+			return (
+				<ImageErrorState
+					prompt={prompt}
+					error={typeof status.error === "string" ? status.error : "An error occurred"}
+				/>
+			);
+		}
+	}
+
+	if (!result) {
+		return (
+			<div className="my-4">
+				<ImageLoading title="Loading" />
+			</div>
+		);
+	}
+
+	if (result.error) {
+		return <ImageErrorState prompt={prompt} error={result.error} />;
+	}
+
+	return (
+		<div className="my-4">
+			<ImageErrorBoundary>
+				<ParsedImage result={result} />
+			</ImageErrorBoundary>
+		</div>
+	);
+};
+
+export {
+	GenerateImageArgsSchema,
+	GenerateImageResultSchema,
+	type GenerateImageArgs,
+	type GenerateImageResult,
+};
diff --git a/surfsense_web/components/tool-ui/image/index.tsx b/surfsense_web/components/tool-ui/image/index.tsx
index ec04b779f..b8e37b620 100644
--- a/surfsense_web/components/tool-ui/image/index.tsx
+++ b/surfsense_web/components/tool-ui/image/index.tsx
@@ -6,7 +6,7 @@ import { Component, type ReactNode, useState } from "react";
 import { z } from "zod";
 import { Badge } from "@/components/ui/badge";
 import { Card } from "@/components/ui/card";
-import { Spinner } from "@/components/ui/spinner";
+import { TextShimmerLoader } from "@/components/prompt-kit/loader";
 import { cn } from "@/lib/utils";
 
 /**
@@ -145,14 +145,14 @@ export class ImageErrorBoundary extends Component<
 	render() {
 		if (this.state.hasError) {
 			return (
-				<Card className="w-full max-w-md overflow-hidden">
-					<div className="aspect-square bg-muted flex items-center justify-center">
-						<div className="flex flex-col items-center gap-2 text-muted-foreground">
-							<ImageIcon className="size-8" />
-							<p className="text-sm">Failed to load image</p>
-						</div>
+			<Card className="w-full max-w-md overflow-hidden rounded-2xl border-0 shadow-none select-none">
+				<div className="aspect-square bg-muted flex items-center justify-center">
+					<div className="flex flex-col items-center gap-2 text-muted-foreground">
+						<ImageIcon className="size-8" />
+						<p className="text-sm">Failed to load image</p>
 					</div>
-				</Card>
+				</div>
+			</Card>
 			);
 		}
 
@@ -165,7 +165,7 @@ export class ImageErrorBoundary extends Component<
  */
 export function ImageSkeleton({ maxWidth = "512px" }: { maxWidth?: string }) {
 	return (
-		<Card className="w-full overflow-hidden animate-pulse" style={{ maxWidth }}>
+		<Card className="w-full overflow-hidden rounded-2xl border-0 shadow-none select-none animate-pulse" style={{ maxWidth }}>
 			<div className="aspect-square bg-muted flex items-center justify-center">
 				<ImageIcon className="size-12 text-muted-foreground/30" />
 			</div>
@@ -176,14 +176,11 @@ export function ImageSkeleton({ maxWidth = "512px" }: { maxWidth?: string }) {
 /**
  * Image Loading State
  */
-export function ImageLoading({ title = "Loading image..." }: { title?: string }) {
+export function ImageLoading({ title = "Loading", maxWidth = "512px" }: { title?: string; maxWidth?: string }) {
 	return (
-		<Card className="w-full max-w-md overflow-hidden">
+		<Card className="w-full overflow-hidden rounded-2xl border-0 shadow-none select-none" style={{ maxWidth }}>
 			<div className="aspect-square bg-muted flex items-center justify-center">
-				<div className="flex flex-col items-center gap-3">
-					<Spinner size="lg" className="text-muted-foreground" />
-					<p className="text-muted-foreground text-sm">{title}</p>
-				</div>
+				<TextShimmerLoader text={title} size="md" />
 			</div>
 		</Card>
 	);
@@ -214,8 +211,8 @@ export function Image({
 	const [isHovered, setIsHovered] = useState(false);
 	const [imageError, setImageError] = useState(false);
 	const [imageLoaded, setImageLoaded] = useState(false);
-	const displayDomain = domain || source?.label;
 	const isGenerated = domain === "ai-generated";
+	const displayDomain = isGenerated ? "AI Generated" : (domain || source?.label);
 	const isAutoRatio = !ratio || ratio === "auto";
 
 	const handleClick = () => {
@@ -227,7 +224,7 @@ export function Image({
 
 	if (imageError) {
 		return (
-			<Card id={id} className={cn("w-full overflow-hidden", className)} style={{ maxWidth }}>
+			<Card id={id} className={cn("w-full overflow-hidden rounded-2xl border-0 shadow-none select-none", className)} style={{ maxWidth }}>
 				<div className="aspect-square bg-muted flex items-center justify-center">
 					<div className="flex flex-col items-center gap-2 text-muted-foreground">
 						<ImageIcon className="size-8" />
@@ -242,8 +239,7 @@ export function Image({
 		<Card
 			id={id}
 			className={cn(
-				"group w-full overflow-hidden cursor-pointer transition-shadow duration-200 hover:shadow-lg",
-				isGenerated && "ring-1 ring-primary/10",
+				"group w-full overflow-hidden rounded-2xl border-0 shadow-none select-none cursor-pointer transition-shadow duration-200 hover:shadow-lg",
 				className
 			)}
 			style={{ maxWidth }}
@@ -263,20 +259,24 @@ export function Image({
 				{isAutoRatio ? (
 					/* Auto ratio: image renders at natural dimensions, no cropping */
 					<>
-						{!imageLoaded && (
-							<div className="aspect-square flex items-center justify-center">
-								<Spinner size="lg" className="text-muted-foreground" />
-							</div>
-						)}
-						{/* eslint-disable-next-line @next/next/no-img-element */}
-						<img
+					{!imageLoaded && (
+						<div className="aspect-square flex items-center justify-center">
+							<TextShimmerLoader text="Loading" size="md" />
+						</div>
+					)}
+						<NextImage
 							src={src}
 							alt={alt}
+							width={0}
+							height={0}
+							sizes="100vw"
+							loading="eager"
 							className={cn(
 								"w-full h-auto transition-transform duration-300",
 								isHovered && "scale-[1.02]",
 								!imageLoaded && "hidden"
 							)}
+							unoptimized
 							onLoad={() => setImageLoaded(true)}
 							onError={() => setImageError(true)}
 						/>
@@ -316,11 +316,9 @@ export function Image({
 						{description && (
 							<p className="text-white/80 text-xs line-clamp-2 mb-1.5">{description}</p>
 						)}
-						{displayDomain && (
+						{displayDomain && !isGenerated && (
 							<div className="flex items-center gap-1.5">
-								{isGenerated ? (
-									<SparklesIcon className="size-3.5 text-white/70" />
-								) : source?.iconUrl ? (
+								{source?.iconUrl ? (
 									<NextImage
 										src={source.iconUrl}
 										alt={source.label}