refactor(app): introduce attachment storage layer and draft lifecycle management

Replace inline data-URL image handling with a persistent attachment storage service, add draft lifecycle states (active/abandoned/sent) with generation tracking, and update fetchAgent to return project placement alongside agent snapshots.
2026-02-22 16:22:21 +07:00
parent 2652bdda2e
commit 0c710b09f5
55 changed files with 2381 additions and 350 deletions
@@ -0,0 +1,48 @@
+import { describe, expect, it } from "vitest";
+import { shouldSkipDraftPersist } from "./agent-input-area.draft-persist-guard";
+
+describe("shouldSkipDraftPersist", () => {
+  it("blocks persist while hydrate for current uncontrolled generation is incomplete", () => {
+    expect(
+      shouldSkipDraftPersist({
+        isControlled: false,
+        currentGeneration: 2,
+        hydratedGeneration: 1,
+        isCurrentGeneration: true,
+      })
+    ).toBe(true);
+  });
+
+  it("allows persist after hydrate completes for current generation", () => {
+    expect(
+      shouldSkipDraftPersist({
+        isControlled: false,
+        currentGeneration: 3,
+        hydratedGeneration: 3,
+        isCurrentGeneration: true,
+      })
+    ).toBe(false);
+  });
+
+  it("blocks persist for stale generations", () => {
+    expect(
+      shouldSkipDraftPersist({
+        isControlled: false,
+        currentGeneration: 4,
+        hydratedGeneration: 4,
+        isCurrentGeneration: false,
+      })
+    ).toBe(true);
+  });
+
+  it("does not block controlled draft persistence", () => {
+    expect(
+      shouldSkipDraftPersist({
+        isControlled: true,
+        currentGeneration: 0,
+        hydratedGeneration: 0,
+        isCurrentGeneration: true,
+      })
+    ).toBe(false);
+  });
+});
@@ -0,0 +1,20 @@
+export function shouldSkipDraftPersist(input: {
+  isControlled: boolean;
+  currentGeneration: number;
+  hydratedGeneration: number;
+  isCurrentGeneration: boolean;
+}): boolean {
+  if (input.isControlled) {
+    return false;
+  }
+
+  if (input.currentGeneration <= 0) {
+    return true;
+  }
+
+  if (!input.isCurrentGeneration) {
+    return true;
+  }
+
+  return input.hydratedGeneration !== input.currentGeneration;
+}
@@ -12,6 +12,7 @@ import { AgentStatusBar } from './agent-status-bar'
 import { useImageAttachmentPicker } from '@/hooks/use-image-attachment-picker'
 import { useSessionStore } from '@/stores/session-store'
 import { useDraftStore } from '@/stores/draft-store'
+import { buildDraftStoreKey } from '@/stores/draft-keys'
 import {
  MessageInput,
  type MessagePayload,
@@ -30,6 +31,11 @@ import { Shortcut } from '@/components/ui/shortcut'
 import { Autocomplete } from '@/components/ui/autocomplete'
 import { useAgentAutocomplete } from '@/hooks/use-agent-autocomplete'
 import { useHostRuntimeSession } from '@/runtime/host-runtime'
+import {
+  deleteAttachments,
+  persistAttachmentFromFileUri,
+} from '@/attachments/service'
+import { shouldSkipDraftPersist } from '@/components/agent-input-area.draft-persist-guard'

 type QueuedMessage = {
  id: string
@@ -40,6 +46,7 @@ type QueuedMessage = {
 interface AgentInputAreaProps {
  agentId: string
  serverId: string
+  draftId?: string
  onSubmitMessage?: (payload: MessagePayload) => Promise<void>
  /** Externally controlled loading state. When true, disables the submit button. */
  isSubmitLoading?: boolean
@@ -60,6 +67,7 @@ const EMPTY_ARRAY: readonly QueuedMessage[] = []
 export function AgentInputArea({
  agentId,
  serverId,
+  draftId,
  onSubmitMessage,
  isSubmitLoading = false,
  blurOnSubmit = false,
@@ -89,8 +97,13 @@ export function AgentInputArea({

  const agent = useSessionStore((state) => state.sessions[serverId]?.agents?.get(agentId))

+  const draftStoreKey = buildDraftStoreKey({ serverId, agentId, draftId })
  const getDraftInput = useDraftStore((state) => state.getDraftInput)
+  const hydrateDraftInput = useDraftStore((state) => state.hydrateDraftInput)
  const saveDraftInput = useDraftStore((state) => state.saveDraftInput)
+  const clearDraftInput = useDraftStore((state) => state.clearDraftInput)
+  const beginDraftGeneration = useDraftStore((state) => state.beginDraftGeneration)
+  const isDraftGenerationCurrent = useDraftStore((state) => state.isDraftGenerationCurrent)

  const queuedMessagesRaw = useSessionStore((state) =>
    state.sessions[serverId]?.queuedMessages?.get(agentId)
@@ -111,6 +124,8 @@ export function AgentInputArea({
  const [sendError, setSendError] = useState<string | null>(null)
  const lastHandledMessageInputActionRequestIdRef = useRef<number | null>(null)
  const messageInputRef = useRef<MessageInputRef>(null)
+  const draftGenerationRef = useRef(0)
+  const hydratedGenerationRef = useRef(0)

  const autocomplete = useAgentAutocomplete({
    userInput,
@@ -176,10 +191,10 @@ export function AgentInputArea({
        throw new Error('Host is not connected')
      }

-      const messageId = generateMessageId()
+      const clientMessageId = generateMessageId()
      const userMessage: StreamItem = {
        kind: 'user_message',
-        id: messageId,
+        id: clientMessageId,
        text,
        timestamp: new Date(),
        ...(images && images.length > 0 ? { images } : {}),
@@ -209,7 +224,7 @@ export function AgentInputArea({

      const imagesData = await encodeImages(images)
      await client.sendAgentMessage(agentId, text, {
-        messageId,
+        messageId: clientMessageId,
        ...(imagesData && imagesData.length > 0 ? { images: imagesData } : {}),
      })
    }
@@ -295,7 +310,7 @@ export function AgentInputArea({
    forceSend?: boolean
  ) {
    const trimmedMessage = message.trim()
-    if (!trimmedMessage) return
+    if (!trimmedMessage && !imageAttachments?.length) return
    // When the parent controls submission (e.g. draft agent creation), let it
    // decide what to do even if the socket is currently disconnected (so we
    // don't no-op and lose deterministic error handling in the UI/tests).
@@ -322,6 +337,7 @@ export function AgentInputArea({

    try {
      await submitMessage(trimmedMessage, imageAttachments)
+      clearDraftInput({ draftKey: draftStoreKey, lifecycle: 'sent' })
    } catch (error) {
      console.error('[AgentInput] Failed to send message:', error)
      // Restore input so the user never loses their message
@@ -353,15 +369,26 @@ export function AgentInputArea({
      return
    }

-    const newImages = result.assets.map((asset) => ({
-      uri: asset.uri,
-      mimeType: asset.mimeType || 'image/jpeg',
-    }))
+    const newImages = await Promise.all(
+      result.assets.map(async (asset) => {
+        return await persistAttachmentFromFileUri({
+          uri: asset.uri,
+          mimeType: asset.mimeType || 'image/jpeg',
+          fileName: asset.fileName ?? null,
+        })
+      })
+    )
    setSelectedImages((prev) => [...prev, ...newImages])
  }

  function handleRemoveImage(index: number) {
-    setSelectedImages((prev) => prev.filter((_, i) => i !== index))
+    setSelectedImages((prev) => {
+      const removed = prev[index]
+      if (removed) {
+        void deleteAttachments([removed])
+      }
+      return prev.filter((_, i) => i !== index)
+    })
  }

  useEffect(() => {
@@ -377,27 +404,72 @@ export function AgentInputArea({
    if (isControlled) {
      return
    }
-    const draft = getDraftInput(agentId)
-    if (!draft) {
-      setUserInput('')
-      setSelectedImages([])
-      return
-    }
+    const generation = beginDraftGeneration(draftStoreKey)
+    draftGenerationRef.current = generation
+    hydratedGenerationRef.current = 0
+    setUserInput('')
+    setSelectedImages([])
+    let cancelled = false

-    setUserInput(draft.text)
-    setSelectedImages(draft.images as ImageAttachment[])
-  }, [agentId, getDraftInput, isControlled])
+    void (async () => {
+      const draft = await hydrateDraftInput(draftStoreKey)
+      if (cancelled) {
+        return
+      }
+      if (!isDraftGenerationCurrent({ draftKey: draftStoreKey, generation })) {
+        return
+      }
+      if (!draft) {
+        hydratedGenerationRef.current = generation
+        return
+      }
+
+      setUserInput(draft.text)
+      setSelectedImages(draft.images)
+      hydratedGenerationRef.current = generation
+    })()
+
+    return () => {
+      cancelled = true
+    }
+  }, [
+    beginDraftGeneration,
+    draftStoreKey,
+    hydrateDraftInput,
+    isControlled,
+    isDraftGenerationCurrent,
+    setUserInput,
+  ])

  // Persist drafts into the shared session store with change detection to avoid redundant work
  useEffect(() => {
-    const existing = getDraftInput(agentId)
+    const currentGeneration = draftGenerationRef.current
+    const isCurrentGeneration =
+      currentGeneration > 0
+        ? isDraftGenerationCurrent({ draftKey: draftStoreKey, generation: currentGeneration })
+        : true
+    if (
+      shouldSkipDraftPersist({
+        isControlled,
+        currentGeneration,
+        hydratedGeneration: hydratedGenerationRef.current,
+        isCurrentGeneration,
+      })
+    ) {
+      return
+    }
+
+    const existing = getDraftInput(draftStoreKey)
    const isSameText = existing?.text === userInput
-    const existingImages: ImageAttachment[] = (existing?.images ?? []) as ImageAttachment[]
+    const existingImages: ImageAttachment[] = existing?.images ?? []
    const isSameImages =
      existingImages.length === selectedImages.length &&
      existingImages.every((img, idx) => {
        return (
-          img.uri === selectedImages[idx]?.uri && img.mimeType === selectedImages[idx]?.mimeType
+          img.id === selectedImages[idx]?.id &&
+          img.mimeType === selectedImages[idx]?.mimeType &&
+          img.storageType === selectedImages[idx]?.storageType &&
+          img.storageKey === selectedImages[idx]?.storageKey
        )
      })

@@ -405,8 +477,28 @@ export function AgentInputArea({
      return
    }

-    saveDraftInput(agentId, { text: userInput, images: selectedImages })
-  }, [agentId, userInput, selectedImages, getDraftInput, saveDraftInput])
+    const hasContent = userInput.trim().length > 0 || selectedImages.length > 0
+    if (!hasContent) {
+      if (existing) {
+        clearDraftInput({ draftKey: draftStoreKey, lifecycle: 'abandoned' })
+      }
+      return
+    }
+
+    saveDraftInput({
+      draftKey: draftStoreKey,
+      draft: { text: userInput, images: selectedImages },
+    })
+  }, [
+    clearDraftInput,
+    draftStoreKey,
+    getDraftInput,
+    isControlled,
+    isDraftGenerationCurrent,
+    saveDraftInput,
+    selectedImages,
+    userInput,
+  ])

  // Keyboard-dispatched message-input actions are routed through store requests.
  useEffect(() => {
@@ -134,7 +134,6 @@ export function LeftSidebar({ selectedAgentId }: LeftSidebarProps) {
    isRevalidating,
    refreshAll,
  } = useSidebarAgentsList({
-    isOpen,
    serverId: activeServerId,
    selectedProjectFilterKeys,
  });
@@ -28,14 +28,13 @@ import {
  collectImageFilesFromClipboardData,
  filesToImageAttachments,
 } from '@/utils/image-attachments-from-files'
+import type { AttachmentMetadata } from '@/attachments/types'
+import { useAttachmentPreviewUrl } from '@/attachments/use-attachment-preview-url'
 import { Tooltip, TooltipContent, TooltipTrigger } from '@/components/ui/tooltip'
 import { Shortcut } from '@/components/ui/shortcut'
 import type { MessageInputKeyboardActionKind } from '@/keyboard/actions'

-export interface ImageAttachment {
-  uri: string
-  mimeType: string
-}
+export type ImageAttachment = AttachmentMetadata

 export interface MessagePayload {
  text: string
@@ -111,6 +110,14 @@ type TextAreaHandle = {
  } & Record<string, unknown>
 }

+function ImageAttachmentThumbnail({ image }: { image: ImageAttachment }) {
+  const uri = useAttachmentPreviewUrl(image)
+  if (!uri) {
+    return <View style={styles.imageThumbnailPlaceholder} />
+  }
+  return <Image source={{ uri }} style={styles.imageThumbnail} />
+}
+
 export const MessageInput = forwardRef<MessageInputRef, MessageInputProps>(function MessageInput(
  {
    value,
@@ -636,14 +643,14 @@ export const MessageInput = forwardRef<MessageInputRef, MessageInputProps>(funct
          <View style={styles.imagePreviewContainer} testID="message-input-image-preview">
            {images.map((image, index) => (
              <Pressable
-                key={`${image.uri}-${index}`}
+                key={`${image.id}-${index}`}
                testID="message-input-image-pill"
                style={styles.imagePill}
                onPress={onRemoveImage ? () => onRemoveImage(index) : undefined}
              >
                {({ hovered }) => (
                  <>
-                    <Image source={{ uri: image.uri }} style={styles.imageThumbnail} />
+                    <ImageAttachmentThumbnail image={image} />
                    {onRemoveImage && (
                      <View
                        style={[
@@ -894,6 +901,11 @@ const styles = StyleSheet.create(((theme: any) => ({
    width: 48,
    height: 48,
  },
+  imageThumbnailPlaceholder: {
+    width: 48,
+    height: 48,
+    backgroundColor: theme.colors.surface2,
+  },
  removeImageButton: {
    position: 'absolute',
    top: 0,
@@ -75,6 +75,7 @@ import { openExternalUrl } from "@/utils/open-external-url";
 export type { InlinePathTarget } from "@/utils/inline-path";
 import { useToolCallSheet } from "./tool-call-sheet";
 import { ToolCallDetailsContent } from "./tool-call-details";
+import { useAttachmentPreviewUrl } from "@/attachments/use-attachment-preview-url";

 interface UserMessageProps {
  message: string;
@@ -306,6 +307,11 @@ const userMessageStylesheet = StyleSheet.create((theme) => ({
    width: 48,
    height: 48,
  },
+  imageThumbnailPlaceholder: {
+    width: 48,
+    height: 48,
+    backgroundColor: theme.colors.surface1,
+  },
  copyButton: {
    alignSelf: "flex-end",
    padding: theme.spacing[1],
@@ -319,6 +325,18 @@ const userMessageStylesheet = StyleSheet.create((theme) => ({
  },
 }));

+function UserMessageAttachmentThumbnail({
+  image,
+}: {
+  image: UserMessageImageAttachment;
+}) {
+  const uri = useAttachmentPreviewUrl(image);
+  if (!uri) {
+    return <View style={userMessageStylesheet.imageThumbnailPlaceholder} />;
+  }
+  return <Image source={{ uri }} style={userMessageStylesheet.imageThumbnail} />;
+}
+
 export const UserMessage = memo(function UserMessage({
  message,
  images = [],
@@ -368,13 +386,10 @@ export const UserMessage = memo(function UserMessage({
            >
              {images.map((image, index) => (
                <View
-                  key={`${image.uri}-${index}`}
+                  key={`${image.id}-${index}`}
                  style={userMessageStylesheet.imagePill}
                >
-                  <Image
-                    source={{ uri: image.uri }}
-                    style={userMessageStylesheet.imageThumbnail}
-                  />
+                  <UserMessageAttachmentThumbnail image={image} />
                </View>
              ))}
            </View>