Refactor agent input draft persistence

packages/app/src/hooks/use-agent-input-draft.ts

@@ -0,0 +1,161 @@
+import { useCallback, useEffect, useRef, useState } from 'react'
+import type { AttachmentMetadata } from '@/attachments/types'
+import { useDraftStore } from '@/stores/draft-store'
+
+type ImageUpdater =
+  | AttachmentMetadata[]
+  | ((prev: AttachmentMetadata[]) => AttachmentMetadata[])
+
+interface AgentInputDraft {
+  text: string
+  setText: (text: string) => void
+  images: AttachmentMetadata[]
+  setImages: (updater: ImageUpdater) => void
+  clear: (lifecycle: 'sent' | 'abandoned') => void
+  isHydrated: boolean
+}
+
+function hasDraftContent(input: { text: string; images: AttachmentMetadata[] }): boolean {
+  return input.text.trim().length > 0 || input.images.length > 0
+}
+
+function areImagesEqual(input: {
+  left: AttachmentMetadata[]
+  right: AttachmentMetadata[]
+}): boolean {
+  if (input.left.length !== input.right.length) {
+    return false
+  }
+
+  return input.left.every((image, index) => {
+    const other = input.right[index]
+    return (
+      image.id === other?.id &&
+      image.mimeType === other?.mimeType &&
+      image.storageType === other?.storageType &&
+      image.storageKey === other?.storageKey
+    )
+  })
+}
+
+export function useAgentInputDraft(draftKey: string): AgentInputDraft {
+  const [text, setText] = useState('')
+  const [images, setImagesState] = useState<AttachmentMetadata[]>([])
+  const [isHydrated, setIsHydrated] = useState(false)
+  const draftGenerationRef = useRef(0)
+  const hydratedGenerationRef = useRef(0)
+
+  const setImages = useCallback((updater: ImageUpdater) => {
+    setImagesState((previousImages) => {
+      if (typeof updater === 'function') {
+        return updater(previousImages)
+      }
+      return updater
+    })
+  }, [])
+
+  const clear = useCallback(
+    (lifecycle: 'sent' | 'abandoned') => {
+      const store = useDraftStore.getState()
+      store.clearDraftInput({ draftKey, lifecycle })
+
+      const generation = store.beginDraftGeneration(draftKey)
+      draftGenerationRef.current = generation
+      hydratedGenerationRef.current = generation
+
+      setText('')
+      setImagesState([])
+      setIsHydrated(true)
+    },
+    [draftKey]
+  )
+
+  useEffect(() => {
+    const store = useDraftStore.getState()
+    const generation = store.beginDraftGeneration(draftKey)
+    draftGenerationRef.current = generation
+    hydratedGenerationRef.current = 0
+
+    setText('')
+    setImagesState([])
+    setIsHydrated(false)
+
+    let cancelled = false
+
+    void (async () => {
+      const draft = await store.hydrateDraftInput(draftKey)
+      if (cancelled) {
+        return
+      }
+      if (!useDraftStore.getState().isDraftGenerationCurrent({ draftKey, generation })) {
+        return
+      }
+
+      if (draft) {
+        setText(draft.text)
+        setImagesState(draft.images)
+      }
+
+      hydratedGenerationRef.current = generation
+      setIsHydrated(true)
+    })()
+
+    return () => {
+      cancelled = true
+    }
+  }, [draftKey])
+
+  useEffect(() => {
+    const currentGeneration = draftGenerationRef.current
+    if (currentGeneration <= 0) {
+      return
+    }
+
+    const store = useDraftStore.getState()
+    const isCurrentGeneration = store.isDraftGenerationCurrent({
+      draftKey,
+      generation: currentGeneration,
+    })
+    if (!isCurrentGeneration) {
+      return
+    }
+    if (hydratedGenerationRef.current !== currentGeneration) {
+      return
+    }
+
+    const existing = store.getDraftInput(draftKey)
+    const isSameDraft =
+      existing?.text === text &&
+      areImagesEqual({
+        left: existing?.images ?? [],
+        right: images,
+      })
+    if (isSameDraft) {
+      return
+    }
+
+    if (!hasDraftContent({ text, images })) {
+      if (existing) {
+        store.clearDraftInput({ draftKey, lifecycle: 'abandoned' })
+      }
+      return
+    }
+
+    store.saveDraftInput({
+      draftKey,
+      draft: {
+        text,
+        images,
+      },
+    })
+  }, [draftKey, images, text])
+
+  return {
+    text,
+    setText,
+    images,
+    setImages,
+    clear,
+    isHydrated,
+  }
+}

packages/app/src/hooks/use-draft-agent-create-flow.ts

@@ -12,11 +12,10 @@ type CreateAttempt = {
 }
 
 type DraftAgentMachineState =
-  | { tag: 'draft'; promptText: string; errorMessage: string }
+  | { tag: 'draft'; errorMessage: string }
   | { tag: 'creating'; attempt: CreateAttempt }
 
 type DraftAgentMachineEvent =
-  | { type: 'DRAFT_SET_PROMPT'; text: string }
   | { type: 'DRAFT_SET_ERROR'; message: string }
   | { type: 'SUBMIT'; attempt: CreateAttempt }
   | { type: 'CREATE_FAILED'; message: string }
@@ -27,12 +26,6 @@ function assertNever(value: never): never {
 
 function reducer(state: DraftAgentMachineState, event: DraftAgentMachineEvent): DraftAgentMachineState {
   switch (event.type) {
-    case 'DRAFT_SET_PROMPT': {
-      if (state.tag !== 'draft') {
-        return state
-      }
-      return { ...state, promptText: event.text }
-    }
     case 'DRAFT_SET_ERROR': {
       if (state.tag !== 'draft') {
         return state
@@ -46,7 +39,7 @@ function reducer(state: DraftAgentMachineState, event: DraftAgentMachineEvent):
       if (state.tag !== 'creating') {
         return state
       }
-      return { tag: 'draft', promptText: state.attempt.text, errorMessage: event.message }
+      return { tag: 'draft', errorMessage: event.message }
     }
     default:
       return assertNever(event)
@@ -94,7 +87,6 @@ export function useDraftAgentCreateFlow<TDraftAgent, TCreateResult>({
 }: UseDraftAgentCreateFlowOptions<TDraftAgent, TCreateResult>) {
   const [machine, dispatch] = useReducer(reducer, {
     tag: 'draft',
-    promptText: '',
     errorMessage: '',
   } as DraftAgentMachineState)
 
@@ -103,7 +95,6 @@ export function useDraftAgentCreateFlow<TDraftAgent, TCreateResult>({
   const markPendingCreateLifecycle = useCreateFlowStore((state) => state.markLifecycle)
   const clearPendingCreateAttempt = useCreateFlowStore((state) => state.clear)
 
-  const promptValue = machine.tag === 'draft' ? machine.promptText : ''
   const formErrorMessage = machine.tag === 'draft' ? machine.errorMessage : ''
   const isSubmitting = machine.tag === 'creating'
 
@@ -132,10 +123,6 @@ export function useDraftAgentCreateFlow<TDraftAgent, TCreateResult>({
     return buildDraftAgent(machine.attempt)
   }, [buildDraftAgent, machine])
 
-  const setPromptText = useCallback((text: string) => {
-    dispatch({ type: 'DRAFT_SET_PROMPT', text })
-  }, [])
-
   const handleCreateFromInput = useCallback(
     async ({ text, images }: SubmitContext) => {
       if (isSubmitting) {
@@ -226,12 +213,10 @@ export function useDraftAgentCreateFlow<TDraftAgent, TCreateResult>({
 
   return {
     machine,
-    promptValue,
     formErrorMessage,
     isSubmitting,
     optimisticStreamItems,
     draftAgent,
-    setPromptText,
     handleCreateFromInput,
   }
 }