From c66247e4a5ec5c9ad63d48f06c139b414621d898 Mon Sep 17 00:00:00 2001
From: HardMax71 <maxymazatyan@gmail.com>
Date: Wed, 7 Oct 2026 17:39:58 +0200
Subject: [PATCH] fix(attachments): show the fallback icon when an image fails
 to load

---
 .../attachment-preview-image-fallback.md      |  5 ++
 .../components/(chatbot)/attachments.mdx      |  2 +-
 .../elements/__tests__/attachments.test.tsx   | 57 +++++++++++++++++++
 packages/elements/src/attachments.tsx         | 29 ++++++++--
 4 files changed, 88 insertions(+), 5 deletions(-)
 create mode 100644 .changeset/attachment-preview-image-fallback.md
 create mode 100644 packages/elements/__tests__/attachments.test.tsx

diff --git a/.changeset/attachment-preview-image-fallback.md b/.changeset/attachment-preview-image-fallback.md
new file mode 100644
index 00000000..4820230b
--- /dev/null
+++ b/.changeset/attachment-preview-image-fallback.md
@@ -0,0 +1,5 @@
+---
+"ai-elements": patch
+---
+
+Show the fallback icon in AttachmentPreview when an image fails to load
diff --git a/apps/docs/content/components/(chatbot)/attachments.mdx b/apps/docs/content/components/(chatbot)/attachments.mdx
index 0302fdf4..e323af0e 100644
--- a/apps/docs/content/components/(chatbot)/attachments.mdx
+++ b/apps/docs/content/components/(chatbot)/attachments.mdx
@@ -133,7 +133,7 @@ Displays the media preview (image, video, or icon).
 <TypeTable
   type={{
     fallbackIcon: {
-      description: "Custom icon to display when no preview is available.",
+      description: "Custom icon to display when no preview is available or the image fails to load.",
       type: "React.ReactNode",
       optional: true,
     },
diff --git a/packages/elements/__tests__/attachments.test.tsx b/packages/elements/__tests__/attachments.test.tsx
new file mode 100644
index 00000000..a1e98939
--- /dev/null
+++ b/packages/elements/__tests__/attachments.test.tsx
@@ -0,0 +1,57 @@
+import { render, screen } from "@testing-library/react";
+
+import type { AttachmentData } from "../src/attachments";
+import { Attachment, AttachmentPreview, Attachments } from "../src/attachments";
+
+const brokenImage: AttachmentData = {
+  filename: "photo.png",
+  id: "file-1",
+  mediaType: "image/png",
+  type: "file",
+  url: "data:image/png;base64,bm90IGFuIGltYWdl",
+};
+
+const validImage: AttachmentData = {
+  ...brokenImage,
+  url: "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII=",
+};
+
+describe("attachmentPreview", () => {
+  it("shows the fallback icon when the image fails to load", async () => {
+    render(
+      <Attachments>
+        <Attachment data={brokenImage}>
+          <AttachmentPreview fallbackIcon={<span data-testid="fallback" />} />
+        </Attachment>
+      </Attachments>
+    );
+
+    await expect(screen.findByTestId("fallback")).resolves.toBeInTheDocument();
+    expect(screen.queryByAltText("photo.png")).not.toBeInTheDocument();
+  });
+
+  it("tries a new url after a failed one", async () => {
+    const { rerender } = render(
+      <Attachments>
+        <Attachment data={brokenImage}>
+          <AttachmentPreview fallbackIcon={<span data-testid="fallback" />} />
+        </Attachment>
+      </Attachments>
+    );
+    await screen.findByTestId("fallback");
+
+    rerender(
+      <Attachments>
+        <Attachment data={validImage}>
+          <AttachmentPreview fallbackIcon={<span data-testid="fallback" />} />
+        </Attachment>
+      </Attachments>
+    );
+
+    expect(screen.getByAltText("photo.png")).toHaveAttribute(
+      "src",
+      validImage.url
+    );
+    expect(screen.queryByTestId("fallback")).not.toBeInTheDocument();
+  });
+});
diff --git a/packages/elements/src/attachments.tsx b/packages/elements/src/attachments.tsx
index cc7b2495..bec40a59 100644
--- a/packages/elements/src/attachments.tsx
+++ b/packages/elements/src/attachments.tsx
@@ -18,7 +18,13 @@ import {
   XIcon,
 } from "lucide-react";
 import type { ComponentProps, HTMLAttributes, ReactNode } from "react";
-import { createContext, useCallback, useContext, useMemo } from "react";
+import {
+  createContext,
+  useCallback,
+  useContext,
+  useMemo,
+  useState,
+} from "react";
 
 // ============================================================================
 // Types
@@ -88,13 +94,15 @@ export const getAttachmentLabel = (data: AttachmentData): string => {
 const renderAttachmentImage = (
   url: string,
   filename: string | undefined,
-  isGrid: boolean
+  isGrid: boolean,
+  onError: () => void
 ) =>
   isGrid ? (
     <img
       alt={filename || "Image"}
       className="size-full object-cover"
       height={96}
+      onError={onError}
       src={url}
       width={96}
     />
@@ -103,6 +111,7 @@ const renderAttachmentImage = (
       alt={filename || "Image"}
       className="size-full rounded object-cover"
       height={20}
+      onError={onError}
       src={url}
       width={20}
     />
@@ -239,6 +248,13 @@ export const AttachmentPreview = ({
   ...props
 }: AttachmentPreviewProps) => {
   const { data, mediaCategory, variant } = useAttachmentContext();
+  const [failedUrl, setFailedUrl] = useState<string>();
+
+  const imageUrl = data.type === "file" ? data.url : undefined;
+  const handleImageError = useCallback(
+    () => setFailedUrl(imageUrl),
+    [imageUrl]
+  );
 
   const iconSize = variant === "inline" ? "size-3" : "size-4";
 
@@ -247,8 +263,13 @@ export const AttachmentPreview = ({
   );
 
   const renderContent = () => {
-    if (mediaCategory === "image" && data.type === "file" && data.url) {
-      return renderAttachmentImage(data.url, data.filename, variant === "grid");
+    if (mediaCategory === "image" && imageUrl && imageUrl !== failedUrl) {
+      return renderAttachmentImage(
+        imageUrl,
+        data.filename,
+        variant === "grid",
+        handleImageError
+      );
     }
 
     if (mediaCategory === "video" && data.type === "file" && data.url) {