refactor(ai-chat): restructure image thumbnail handling and enhance message components

- Moved the `ImageThumbnail` component to a new location, consolidating its functionality with hover preview and error handling. - Updated `UserMessageParts` to utilize `UserRichTextMessage` instead of the deprecated `AIRichTextMessage`, improving clarity in message rendering. - Enhanced `FileAttachmentNode` to utilize the new `useMessageByIdSelector` for better file attachment management. - Refactored `useFileUpload` to omit the nonce from options, streamlining the file upload process. These changes aim to improve the modularity and maintainability of the AI chat components, enhancing the overall user experience. Signed-off-by: Innei <tukon479@gmail.com>
2025-09-24 18:47:27 +08:00 · 2025-09-24 18:47:27 +08:00 · a06a3fe2fc
parent 3f2b7401b8
commit a06a3fe2fc
12 changed files with 344 additions and 330 deletions
--- a/apps/desktop/layer/renderer/src/modules/ai-chat/components/context-bar/blocks/ContextBlock.tsx
+++ b/apps/desktop/layer/renderer/src/modules/ai-chat/components/context-bar/blocks/ContextBlock.tsx
@ -5,7 +5,7 @@ import { memo } from "react"
 import { useTranslation } from "react-i18next"

 import { ROUTE_FEED_IN_FOLDER } from "~/constants"
-import { ImageThumbnail } from "~/modules/ai-chat/components/layouts/ImageThumbnail"
+import { ImageThumbnail } from "~/modules/ai-chat/components/message/ImageThumbnail"
 import { CircularProgress } from "~/modules/ai-chat/components/ui/UploadProgress"
 import { useChatBlockActions } from "~/modules/ai-chat/store/hooks"
 import type { AIChatContextBlock } from "~/modules/ai-chat/store/types"
@ -93,16 +93,12 @@ export const ContextBlock: FC<{ block: AIChatContextBlock }> = memo(({ block })
        const fileCategory = getFileCategoryFromMimeType(type)

        if (fileCategory === "image" && (dataUrl || previewUrl)) {
-          const validPreviewUrl = (dataUrl || previewUrl)!
          return (
            <div className="flex items-center gap-1.5">
              <div className="relative">
                <ImageThumbnail
-                  previewUrl={validPreviewUrl}
-                  originalUrl={dataUrl || validPreviewUrl}
-                  alt={name}
-                  filename={name}
                  className={"m-0.5 size-5 rounded-md"}
+                  attachment={block.attachment}
                />
                {uploadStatus === "uploading" && uploadProgress !== undefined && (
                  <div className="absolute inset-0 flex items-center justify-center rounded-md bg-black/50">
--- a/apps/desktop/layer/renderer/src/modules/ai-chat/components/layouts/ImageThumbnail.tsx
+++ b/apps/desktop/layer/renderer/src/modules/ai-chat/components/layouts/ImageThumbnail.tsx
@ -1,120 +0,0 @@
-import { cn } from "@follow/utils"
-import * as HoverCard from "@radix-ui/react-hover-card"
-import { m } from "motion/react"
-import type { FC } from "react"
-import * as React from "react"
-
-import { useModalStack } from "~/components/ui/modal/stacked/hooks"
-
-export const ImageThumbnail: FC<{
-  previewUrl: string
-  originalUrl: string
-  alt: string
-  filename: string
-  className?: string
-}> = ({ previewUrl, originalUrl, alt, filename, className }) => {
-  const [imageError, setImageError] = React.useState(false)
-  const { present } = useModalStack()
-
-  return (
-    <HoverCard.Root openDelay={300} closeDelay={100}>
-      <HoverCard.Trigger
-        className="cursor-pointer"
-        onClick={() =>
-          present({
-            max: true,
-            title: "Preview Image",
-            content: () => (
-              <div className="flex max-h-full max-w-full items-center justify-center">
-                <img src={originalUrl} className="max-h-full max-w-full" />
-              </div>
-            ),
-          })
-        }
-      >
-        <ImageThumbnailInner className={className} src={previewUrl} alt={alt} />
-      </HoverCard.Trigger>
-      <HoverCard.Portal>
-        <HoverCard.Content
-          className="bg-material-thick border-border w-fit rounded-md border shadow-lg"
-          sideOffset={8}
-        >
-          <div className="relative overflow-hidden rounded-md">
-            {imageError ? (
-              <div className="bg-fill-secondary border-border flex h-32 w-40 items-center justify-center rounded-md border">
-                <div className="text-text-tertiary flex flex-col items-center gap-2">
-                  <i className="i-mgc-photo-album-cute-fi size-6" />
-                  <span className="text-xs">{filename}</span>
-                </div>
-              </div>
-            ) : (
-              <m.img
-                src={previewUrl}
-                alt={alt}
-                onError={() => setImageError(true)}
-                initial={{ opacity: 0, scale: 0.95 }}
-                animate={{ opacity: 1, scale: 1 }}
-                transition={{ duration: 0.2 }}
-                className="max-h-[300px] max-w-[400px] rounded-md"
-              />
-            )}
-          </div>
-        </HoverCard.Content>
-      </HoverCard.Portal>
-    </HoverCard.Root>
-  )
-}
-
-export const ImageThumbnailInner: React.FC<{
-  src: string
-  alt: string
-  className?: string
-}> = ({ src, alt, className }) => {
-  const [imageError, setImageError] = React.useState(false)
-  const [imageLoaded, setImageLoaded] = React.useState(false)
-
-  const handleError = React.useCallback(() => {
-    setImageError(true)
-  }, [])
-
-  const handleLoad = React.useCallback(() => {
-    setImageLoaded(true)
-  }, [])
-
-  if (imageError) {
-    return (
-      <div
-        className={cn(
-          "bg-fill-secondary border-border flex items-center justify-center border",
-          className,
-        )}
-      >
-        <i className="i-mgc-photo-album-cute-re text-text-tertiary size-3" />
-      </div>
-    )
-  }
-
-  return (
-    <div className={cn("relative overflow-hidden", className)}>
-      {!imageLoaded && (
-        <div
-          className={
-            "bg-fill-secondary border-border absolute inset-0 flex items-center justify-center border"
-          }
-        >
-          <i className="i-mgc-loading-3-cute-re text-text-tertiary size-3 animate-spin" />
-        </div>
-      )}
-      <m.img
-        src={src}
-        alt={alt}
-        onError={handleError}
-        onLoad={handleLoad}
-        initial={{ opacity: 0 }}
-        animate={{ opacity: imageLoaded ? 1 : 0 }}
-        transition={{ duration: 0.2 }}
-        className={"size-full object-cover"}
-      />
-    </div>
-  )
-}
--- a/apps/desktop/layer/renderer/src/modules/ai-chat/components/message/AIMessageIdContext.tsx
+++ b/apps/desktop/layer/renderer/src/modules/ai-chat/components/message/AIMessageIdContext.tsx
@ -0,0 +1,11 @@
+import { createContext, use } from "react"
+
+export const AIMessageIdContext = createContext<string | null>(null)
+
+export const useAIMessageId = () => {
+  const ctx = use(AIMessageIdContext)
+  if (!ctx && import.meta.env.DEV) {
+    throw new Error("useAIMessageId must be used within a AIMessageIdContext")
+  }
+  return ctx
+}
--- a/apps/desktop/layer/renderer/src/modules/ai-chat/components/message/ImageThumbnail.tsx
+++ b/apps/desktop/layer/renderer/src/modules/ai-chat/components/message/ImageThumbnail.tsx
@ -1,70 +1,160 @@
 import { cn } from "@follow/utils/utils"
+import * as HoverCard from "@radix-ui/react-hover-card"
+import { m } from "motion/react"
 import * as React from "react"

+import { useModalStack } from "~/components/ui/modal/stacked/hooks"
 import type { FileAttachment } from "~/modules/ai-chat/store/types"

 import { getImageUrl } from "./ai-block-constants"

-interface ImageThumbnailProps {
+type AttachmentImageProps = {
  attachment: FileAttachment
  className?: string
  fallbackIcon?: string
 }

+type ImageThumbnailProps = AttachmentImageProps
+
 /**
- * A robust image thumbnail component with proper error handling and fallback states
+ * Unified ImageThumbnail with hover preview, click-to-modal, loading and error fallbacks
 */
-export const ImageThumbnail: React.FC<ImageThumbnailProps> = React.memo(
-  ({
-    attachment,
-    className = "size-3 rounded object-cover",
-    fallbackIcon = "i-mgc-pic-cute-re",
-  }) => {
-    const [hasError, setHasError] = React.useState(false)
-    const [isLoading, setIsLoading] = React.useState(true)
+export const ImageThumbnail: React.FC<ImageThumbnailProps> = React.memo((props) => {
+  const { present } = useModalStack()

-    const imageUrl = React.useMemo(() => getImageUrl(attachment), [attachment])
+  const [contentImageError, setContentImageError] = React.useState(false)

-    const handleError = React.useCallback(() => {
-      setHasError(true)
-      setIsLoading(false)
-    }, [])
-
-    const handleLoad = React.useCallback(() => {
-      setIsLoading(false)
-    }, [])
-
-    // Reset error state when imageUrl changes
-    React.useEffect(() => {
-      if (imageUrl) {
-        setHasError(false)
-        setIsLoading(true)
-      }
-    }, [imageUrl])
-
-    // If no image URL or there was an error, show fallback icon
-    if (!imageUrl || hasError) {
-      return <i className={cn("size-3", fallbackIcon)} />
+  const computed = React.useMemo(() => {
+    const { attachment, className, fallbackIcon } = props
+    const imageUrl = getImageUrl(attachment)
+    const { name, serverUrl } = attachment
+    const original = serverUrl || imageUrl || null
+    return {
+      previewUrl: imageUrl,
+      originalUrl: original,
+      alt: name,
+      filename: name,
+      className: className ?? "size-3 rounded object-cover",
+      fallbackIcon: fallbackIcon ?? "i-mgc-pic-cute-re",
    }
+  }, [props])

-    return (
-      <div className="relative">
-        <img
-          src={imageUrl}
-          alt={attachment.name}
-          className={cn(className, isLoading && "opacity-50")}
-          onError={handleError}
-          onLoad={handleLoad}
-          loading="lazy"
+  React.useEffect(() => {
+    setContentImageError(false)
+  }, [computed.previewUrl])
+
+  // If no preview URL available, show fallback icon
+  if (!computed.previewUrl) {
+    return <i className={cn("size-3", computed.fallbackIcon)} />
+  }
+
+  return (
+    <HoverCard.Root openDelay={300} closeDelay={100}>
+      <HoverCard.Trigger
+        className="cursor-pointer"
+        onClick={() =>
+          computed.originalUrl &&
+          present({
+            max: true,
+            title: "Preview Image",
+            clickOutsideToDismiss: true,
+            content: () => (
+              <div className="flex max-h-full max-w-full items-center justify-center">
+                <img src={computed.originalUrl!} className="max-h-full max-w-full" />
+              </div>
+            ),
+          })
+        }
+      >
+        <ImageThumbnailInner
+          src={computed.previewUrl}
+          alt={computed.alt}
+          className={computed.className}
        />
-        {isLoading && (
-          <div className="absolute inset-0 flex items-center justify-center">
-            <i className="i-mgc-loading-3-cute-re size-2 animate-spin opacity-50" />
+      </HoverCard.Trigger>
+      <HoverCard.Portal>
+        <HoverCard.Content
+          className="bg-material-thick border-border w-fit rounded-md border shadow-lg"
+          sideOffset={8}
+        >
+          <div className="relative overflow-hidden rounded-md">
+            {contentImageError ? (
+              <div className="bg-fill-secondary border-border flex h-32 w-40 items-center justify-center rounded-md border">
+                <div className="text-text-tertiary flex flex-col items-center gap-2">
+                  <i className="i-mgc-photo-album-cute-fi size-6" />
+                  <span className="text-xs">{computed.filename}</span>
+                </div>
+              </div>
+            ) : (
+              <m.img
+                src={computed.previewUrl}
+                alt={computed.alt}
+                onError={() => setContentImageError(true)}
+                initial={{ opacity: 0, scale: 0.95 }}
+                animate={{ opacity: 1, scale: 1 }}
+                transition={{ duration: 0.2 }}
+                className="max-h-[300px] max-w-[400px] rounded-md"
+              />
+            )}
          </div>
-        )}
-      </div>
-    )
-  },
-)
+        </HoverCard.Content>
+      </HoverCard.Portal>
+    </HoverCard.Root>
+  )
+})

 ImageThumbnail.displayName = "ImageThumbnail"
+
+const ImageThumbnailInner: React.FC<{ src: string; alt: string; className?: string }> = ({
+  src,
+  alt,
+  className,
+}) => {
+  const [imageError, setImageError] = React.useState(false)
+  const [imageLoaded, setImageLoaded] = React.useState(false)
+
+  const handleError = React.useCallback(() => {
+    setImageError(true)
+  }, [])
+
+  const handleLoad = React.useCallback(() => {
+    setImageLoaded(true)
+  }, [])
+
+  if (imageError) {
+    return (
+      <div
+        className={cn(
+          "bg-fill-secondary border-border flex items-center justify-center border",
+          className,
+        )}
+      >
+        <i className="i-mgc-photo-album-cute-re text-text-tertiary size-3" />
+      </div>
+    )
+  }
+
+  return (
+    <div className={cn("relative overflow-hidden", className)}>
+      {!imageLoaded && (
+        <div
+          className={
+            "bg-fill-secondary border-border absolute inset-0 flex items-center justify-center border"
+          }
+        >
+          <i className="i-mgc-loading-3-cute-re text-text-tertiary size-3 animate-spin" />
+        </div>
+      )}
+      <m.img
+        src={src}
+        alt={alt}
+        onError={handleError}
+        onLoad={handleLoad}
+        initial={{ opacity: 0 }}
+        animate={{ opacity: imageLoaded ? 1 : 0 }}
+        transition={{ duration: 0.2 }}
+        className={"size-full object-cover"}
+      />
+    </div>
+  )
+}
--- a/apps/desktop/layer/renderer/src/modules/ai-chat/components/message/UserChatMessage.tsx
+++ b/apps/desktop/layer/renderer/src/modules/ai-chat/components/message/UserChatMessage.tsx
@ -11,6 +11,7 @@ import type { AIChatContextBlock, BizUIMessage } from "~/modules/ai-chat/store/t
 import type { RichTextPart } from "../../types/ChatSession"
 import { convertLexicalToMarkdown } from "../../utils/lexical-markdown"
 import { AIDataBlockPart } from "./AIDataBlockPart"
+import { AIMessageIdContext } from "./AIMessageIdContext"
 import { EditableMessage } from "./EditableMessage"
 import { UserMessageParts } from "./UserMessageParts"

@ -94,117 +95,119 @@ export const UserChatMessage: React.FC<UserChatMessageProps> = React.memo(({ mes
  }, [chatActions, messageId])

  return (
-    <div className="relative flex flex-col gap-3">
-      {/* Render data-block parts separately, outside the chat bubble */}
-      {dataBlockParts.length > 0 && (
-        <div ref={dataBlockRef} className="flex justify-end">
-          <div className="max-w-[calc(100%-1rem)]">
-            {dataBlockParts.map((part) => {
-              if (part.type === "data-block" && "data" in part) {
-                const blocks = part.data as AIChatContextBlock[]
-                return (
-                  <AIDataBlockPart
-                    key={`${messageId}-datablock-${blocks.map((b) => b.id).join("-")}`}
-                    blocks={blocks}
-                  />
-                )
-              }
-              return null
-            })}
-          </div>
-        </div>
-      )}
-
-      {/* Main chat message */}
-      <m.div
-        initial={
-          isStreaming
-            ? {
-                opacity: 0,
-                y: 20,
-                scale: 0.95,
-              }
-            : true
-        }
-        animate={{
-          opacity: 1,
-          y: 0,
-          scale: 1,
-        }}
-        transition={Spring.presets.smooth}
-        onContextMenu={stopPropagation}
-        className="group flex justify-end"
-        onMouseEnter={() => setIsHovered(true)}
-        onMouseLeave={() => setIsHovered(false)}
-      >
-        <div className="text-text relative flex max-w-[calc(100%-1rem)] flex-col gap-2">
-          {/* Normal message display - always rendered to maintain layout */}
-          <div className="text-text bg-mix-accent/background-1/4 rounded-2xl px-4 py-2.5 backdrop-blur-sm">
-            <div className="flex select-text flex-col gap-2 text-sm">
-              <UserMessageParts message={message} />
+    <AIMessageIdContext value={messageId}>
+      <div className="relative flex flex-col gap-3">
+        {/* Render data-block parts separately, outside the chat bubble */}
+        {dataBlockParts.length > 0 && (
+          <div ref={dataBlockRef} className="flex justify-end">
+            <div className="max-w-[calc(100%-1rem)]">
+              {dataBlockParts.map((part) => {
+                if (part.type === "data-block" && "data" in part) {
+                  const blocks = part.data as AIChatContextBlock[]
+                  return (
+                    <AIDataBlockPart
+                      key={`${messageId}-datablock-${blocks.map((b) => b.id).join("-")}`}
+                      blocks={blocks}
+                    />
+                  )
+                }
+                return null
+              })}
            </div>
          </div>
+        )}

-          {/* Action buttons - only show when not editing */}
-          {!isEditing && (
+        {/* Main chat message */}
+        <m.div
+          initial={
+            isStreaming
+              ? {
+                  opacity: 0,
+                  y: 20,
+                  scale: 0.95,
+                }
+              : true
+          }
+          animate={{
+            opacity: 1,
+            y: 0,
+            scale: 1,
+          }}
+          transition={Spring.presets.smooth}
+          onContextMenu={stopPropagation}
+          className="group flex justify-end"
+          onMouseEnter={() => setIsHovered(true)}
+          onMouseLeave={() => setIsHovered(false)}
+        >
+          <div className="text-text relative flex max-w-[calc(100%-1rem)] flex-col gap-2">
+            {/* Normal message display - always rendered to maintain layout */}
+            <div className="text-text bg-mix-accent/background-1/4 rounded-2xl px-4 py-2.5 backdrop-blur-sm">
+              <div className="flex select-text flex-col gap-2 text-sm">
+                <UserMessageParts message={message} />
+              </div>
+            </div>
+
+            {/* Action buttons - only show when not editing */}
+            {!isEditing && (
+              <m.div
+                className="absolute bottom-1 right-0 flex gap-1"
+                initial={{ opacity: 0 }}
+                animate={{
+                  opacity: isHovered ? 1 : 0,
+                }}
+                transition={{ duration: 0.2, ease: "easeOut" }}
+              >
+                <button
+                  type="button"
+                  onClick={handleEdit}
+                  className="text-text-secondary hover:bg-fill-secondary flex items-center gap-1 rounded-md px-2 py-1 text-xs transition-colors"
+                  title="Edit message"
+                >
+                  <i className="i-mgc-edit-cute-re size-3" />
+                  <span>Edit</span>
+                </button>
+                <button
+                  type="button"
+                  onClick={handleRetry}
+                  className="text-text-secondary hover:bg-fill-secondary flex items-center gap-1 rounded-md px-2 py-1 text-xs transition-colors"
+                  title="Retry"
+                >
+                  <i className="i-mgc-refresh-2-cute-re size-3" />
+                  <span>Retry</span>
+                </button>
+              </m.div>
+            )}
+
+            <div className="h-6" />
+          </div>
+        </m.div>
+
+        {/* Full-width edit overlay - positioned at the top level to span entire container */}
+        <AnimatePresence>
+          {isEditing && (
            <m.div
-              className="absolute bottom-1 right-0 flex gap-1"
-              initial={{ opacity: 0 }}
-              animate={{
-                opacity: isHovered ? 1 : 0,
+              className="absolute inset-x-0 bottom-0 z-[1] flex"
+              style={{
+                top: dataBlockHeight > 0 ? `${dataBlockHeight}px` : 0,
              }}
-              transition={{ duration: 0.2, ease: "easeOut" }}
+              initial={{ opacity: 0, scale: 0.98 }}
+              animate={{ opacity: 1, scale: 1 }}
+              exit={{ opacity: 0, scale: 0.98 }}
+              transition={{ duration: 0.15, ease: "easeOut" }}
            >
-              <button
-                type="button"
-                onClick={handleEdit}
-                className="text-text-secondary hover:bg-fill-secondary flex items-center gap-1 rounded-md px-2 py-1 text-xs transition-colors"
-                title="Edit message"
-              >
-                <i className="i-mgc-edit-cute-re size-3" />
-                <span>Edit</span>
-              </button>
-              <button
-                type="button"
-                onClick={handleRetry}
-                className="text-text-secondary hover:bg-fill-secondary flex items-center gap-1 rounded-md px-2 py-1 text-xs transition-colors"
-                title="Retry"
-              >
-                <i className="i-mgc-refresh-2-cute-re size-3" />
-                <span>Retry</span>
-              </button>
+              <div className="w-full max-w-[var(--ai-chat-message-container-width,65ch)]">
+                <EditableMessage
+                  messageId={messageId}
+                  parts={message.parts}
+                  onSave={handleSaveEdit}
+                  onCancel={handleCancelEdit}
+                  className="min-h-full w-full"
+                />
+              </div>
            </m.div>
          )}
-
-          <div className="h-6" />
-        </div>
-      </m.div>
-
-      {/* Full-width edit overlay - positioned at the top level to span entire container */}
-      <AnimatePresence>
-        {isEditing && (
-          <m.div
-            className="absolute inset-x-0 bottom-0 z-[1] flex"
-            style={{
-              top: dataBlockHeight > 0 ? `${dataBlockHeight}px` : 0,
-            }}
-            initial={{ opacity: 0, scale: 0.98 }}
-            animate={{ opacity: 1, scale: 1 }}
-            exit={{ opacity: 0, scale: 0.98 }}
-            transition={{ duration: 0.15, ease: "easeOut" }}
-          >
-            <div className="w-full max-w-[var(--ai-chat-message-container-width,65ch)]">
-              <EditableMessage
-                messageId={messageId}
-                parts={message.parts}
-                onSave={handleSaveEdit}
-                onCancel={handleCancelEdit}
-                className="min-h-full w-full"
-              />
-            </div>
-          </m.div>
-        )}
-      </AnimatePresence>
-    </div>
+        </AnimatePresence>
+      </div>
+    </AIMessageIdContext>
  )
 })
--- a/apps/desktop/layer/renderer/src/modules/ai-chat/components/message/UserMessageParts.tsx
+++ b/apps/desktop/layer/renderer/src/modules/ai-chat/components/message/UserMessageParts.tsx
@ -3,7 +3,7 @@ import * as React from "react"
 import type { BizUIMessage } from "~/modules/ai-chat/store/types"

 import { AIMarkdownStreamingMessage } from "./AIMarkdownMessage"
-import { AIRichTextMessage } from "./AIRichTextMessage"
+import { UserRichTextMessage } from "./UserRichTextMessage"

 interface UserMessagePartsProps {
  message: BizUIMessage
@ -25,7 +25,7 @@ export const UserMessageParts: React.FC<UserMessagePartsProps> = React.memo(({ m

      case "data-rich-text": {
        return (
-          <AIRichTextMessage
+          <UserRichTextMessage
            key={partKey}
            data={part.data as { state: string; text: string }}
            className="text-text"
--- a/apps/desktop/layer/renderer/src/modules/ai-chat/components/message/UserRichTextMessage.tsx
+++ b/apps/desktop/layer/renderer/src/modules/ai-chat/components/message/UserRichTextMessage.tsx
@ -16,7 +16,7 @@ function onError(error: Error) {
  console.error("Lexical Read-Only Editor Error:", error)
 }

-interface AIRichTextMessageProps {
+interface UserRichTextMessageProps {
  data: {
    state: SerializedEditorState | string // Serialized editor state as a JSON string
    text: string
@ -24,7 +24,7 @@ interface AIRichTextMessageProps {
  className?: string
 }

-export const AIRichTextMessage: React.FC<AIRichTextMessageProps> = React.memo(
+export const UserRichTextMessage: React.FC<UserRichTextMessageProps> = React.memo(
  ({ data, className }) => {
    let initialConfig: InitialConfigType = null!
    if (!initialConfig) {
--- a/apps/desktop/layer/renderer/src/modules/ai-chat/editor/plugins/file-upload/FileAttachmentNode.tsx
+++ b/apps/desktop/layer/renderer/src/modules/ai-chat/editor/plugins/file-upload/FileAttachmentNode.tsx
@ -9,9 +9,10 @@ import type {
 } from "lexical"
 import { DecoratorNode } from "lexical"
 import * as React from "react"
-import { useShallow } from "zustand/shallow"

-import { useAIChatStore } from "~/modules/ai-chat/store/AIChatContext"
+import { useAIMessageId } from "~/modules/ai-chat/components/message/AIMessageIdContext"
+import { useMessageByIdSelector } from "~/modules/ai-chat/store/hooks"
+import { findFileAttachmentBlock, isDataBlockPart } from "~/modules/ai-chat/utils/extractor"

 export type SerializedFileAttachmentNode = Spread<
  {
@ -112,14 +113,18 @@ interface FileAttachmentComponentProps {
 function FileAttachmentComponent({ node }: FileAttachmentComponentProps) {
  const attachmentId = node.getAttachmentId()

-  const fileAttachment = useAIChatStore()(
-    useShallow((state) => {
-      const block = state.blocks.find(
-        (block) => block.type === "fileAttachment" && block.attachment.id === attachmentId,
-      )
-      return block?.type === "fileAttachment" ? block.attachment : null
-    }),
-  )
+  const messageId = useAIMessageId()!
+
+  const fileAttachment = useMessageByIdSelector(messageId, (message) => {
+    for (const part of message.parts) {
+      if (!isDataBlockPart(part)) continue
+
+      const block = findFileAttachmentBlock(part, attachmentId)
+      if (block) {
+        return block.attachment
+      }
+    }
+  })

  if (!fileAttachment) {
    return (
--- a/apps/desktop/layer/renderer/src/modules/ai-chat/hooks/useFileUpload.ts
+++ b/apps/desktop/layer/renderer/src/modules/ai-chat/hooks/useFileUpload.ts
@ -48,7 +48,9 @@ export interface FileUploadHandlers {
 /**
 * Hook for handling file uploads with progress tracking and block management
 */
-export function useFileUpload(options: UseFileUploadOptions = {}): FileUploadHandlers {
+export function useFileUpload(
+  options: Omit<UseFileUploadOptions, "nonce"> = {},
+): FileUploadHandlers {
  const {
    showSuccessToast = false,
    showErrorToast = true,
@ -62,6 +64,7 @@ export function useFileUpload(options: UseFileUploadOptions = {}): FileUploadHan
  const uploadFile = useCallback(
    async (file: File, id?: string): Promise<ProcessFileResult> => {
      // Create initial file attachment for immediate UI feedback
+
      const initialFileAttachment: FileAttachment = {
        id: id || nanoid(),
        name: file.name,
@ -76,14 +79,18 @@ export function useFileUpload(options: UseFileUploadOptions = {}): FileUploadHan
      blockActions.addFileAttachment(initialFileAttachment)

      try {
-        const result = await processAndUploadFile(file, processOptions, (updatedAttachment) => {
-          // Update the attachment with the same ID to maintain consistency
-          const syncedAttachment = {
-            ...updatedAttachment,
-            id: initialFileAttachment.id, // Keep the same ID
-          }
-          blockActions.updateFileAttachment(initialFileAttachment.id, syncedAttachment)
-        })
+        const result = await processAndUploadFile(
+          file,
+          { ...processOptions, nonce: initialFileAttachment.id },
+          (updatedAttachment) => {
+            // Update the attachment with the same ID to maintain consistency
+            const syncedAttachment = {
+              ...updatedAttachment,
+              id: initialFileAttachment.id, // Keep the same ID
+            }
+            blockActions.updateFileAttachment(initialFileAttachment.id, syncedAttachment)
+          },
+        )

        if (result.success && result.fileAttachment) {
          // Update the final completed state with the same ID
--- a/apps/desktop/layer/renderer/src/modules/ai-chat/store/hooks.ts
+++ b/apps/desktop/layer/renderer/src/modules/ai-chat/store/hooks.ts
@ -1,3 +1,6 @@
+import type { BizUIMessage } from "@folo-services/ai-tools"
+import { useShallow } from "zustand/shallow"
+
 import { useAIChatStore } from "./AIChatContext"

 /**
@ -40,6 +43,19 @@ export const useMessages = () => {
  return store((state) => state.messages)
 }

+export const useMessageByIdSelector = <T>(
+  messageId: string,
+  selector: (message: BizUIMessage) => T,
+): T | undefined => {
+  const store = useAIChatStore()
+  return store(
+    useShallow((state) => {
+      const message = state.messages.find((message) => message.id === messageId)
+      return message ? selector(message) : undefined
+    }),
+  )
+}
+
 /**
 * Hook to check if the chat has messages
 */
--- a/apps/desktop/layer/renderer/src/modules/ai-chat/utils/extractor.ts
+++ b/apps/desktop/layer/renderer/src/modules/ai-chat/utils/extractor.ts
@ -0,0 +1,30 @@
+import type { AIChatContextBlock, FileAttachmentContextBlock } from "../store/types"
+
+type AIMessageDataBlockPart = {
+  type: "data-block"
+  data: AIChatContextBlock[]
+}
+export const isDataBlockPart = (part: unknown): part is AIMessageDataBlockPart => {
+  return !!part && typeof part === "object" && "type" in part && part.type === "data-block"
+}
+
+// Narrow a context block to the file attachment block
+export const isFileAttachmentBlock = (
+  block: AIChatContextBlock,
+): block is FileAttachmentContextBlock => {
+  return block.type === "fileAttachment"
+}
+
+export const findFileAttachmentBlock = (
+  part: AIMessageDataBlockPart,
+  attachmentId: string,
+): FileAttachmentContextBlock | undefined => {
+  if (!isDataBlockPart(part)) return
+
+  for (const block of part.data) {
+    if (isFileAttachmentBlock(block) && block.attachment.id === attachmentId) {
+      return block
+    }
+  }
+  return
+}
--- a/apps/desktop/layer/renderer/src/modules/ai-chat/utils/file-processing.ts
+++ b/apps/desktop/layer/renderer/src/modules/ai-chat/utils/file-processing.ts
@ -1,5 +1,3 @@
-import { nanoid } from "nanoid"
-
 import { followApi } from "~/lib/api-client"

 import type { FileAttachment } from "../store/types"
@ -10,6 +8,8 @@ export interface ProcessFileOptions {
  maxImageWidth?: number
  maxImageHeight?: number
  imageQuality?: number
+
+  nonce: string
 }

 export interface ProcessFileResult {
@ -20,7 +20,7 @@ export interface ProcessFileResult {

 export async function processFile(
  file: File,
-  options: ProcessFileOptions = {},
+  options: ProcessFileOptions,
 ): Promise<ProcessFileResult> {
  const { maxImageWidth = 1920, maxImageHeight = 1080, imageQuality = 0.85 } = options

@ -34,7 +34,7 @@ export async function processFile(
  }

  try {
-    const fileId = nanoid()
+    const { nonce: fileId } = options
    let dataUrl: string
    let previewUrl: string | undefined

@ -157,30 +157,6 @@ function fileToDataUrl(file: File): Promise<string> {
  })
 }

-export async function processFileList(
-  files: FileList,
-  options?: ProcessFileOptions,
-): Promise<ProcessFileResult[]> {
-  const results: ProcessFileResult[] = []
-
-  for (const file of files) {
-    if (file) {
-      const result = await processFile(file, options)
-      results.push(result)
-    }
-  }
-
-  return results
-}
-
-export function createFileAttachmentBlock(fileAttachment: FileAttachment) {
-  return {
-    id: fileAttachment.id,
-    type: "fileAttachment" as const,
-    attachment: fileAttachment,
-  }
-}
-
 // Utility to clean up object URLs to prevent memory leaks
 export function cleanupFileAttachment(fileAttachment: FileAttachment) {
  if (fileAttachment.previewUrl?.startsWith("blob:")) {
@ -267,7 +243,7 @@ export async function uploadFileAttachment(

 export async function processAndUploadFile(
  file: File,
-  options: ProcessFileOptions = {},
+  options: ProcessFileOptions,
  onProgressUpdate?: (attachment: FileAttachment) => void,
 ): Promise<ProcessFileResult> {
  // First process the file locally