Refactor agent input draft persistence

packages/app/src/components/agent-input-area.draft-persist-guard.test.ts

@@ -1,48 +0,0 @@
-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);
-  });
-});

packages/app/src/components/agent-input-area.draft-persist-guard.ts

@@ -1,20 +0,0 @@
-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;
-}

packages/app/src/components/agent-input-area.tsx

@@ -10,8 +10,6 @@ import { generateMessageId, type StreamItem } from '@/types/stream'
 import { AgentStatusBar, DraftAgentStatusBar, type DraftAgentStatusBarProps } 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,
@@ -35,7 +33,6 @@ import {
   persistAttachmentFromFileUri,
 } from '@/attachments/service'
 import { resolveStatusControlMode } from '@/components/agent-input-area.status-controls'
-import { shouldSkipDraftPersist } from '@/components/agent-input-area.draft-persist-guard'
 import { markScrollInvestigationRender } from '@/utils/scroll-jank-investigation'
 import { useKeyboardShiftStyle } from '@/hooks/use-keyboard-shift-style'
 import { useKeyboardActionHandler } from '@/hooks/use-keyboard-action-handler'
@@ -47,17 +44,21 @@ type QueuedMessage = {
   images?: ImageAttachment[]
 }
 
+type ImageListUpdater = ImageAttachment[] | ((prev: ImageAttachment[]) => ImageAttachment[])
+
 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
   /** When true, blurs the input immediately when submitting. */
   blurOnSubmit?: boolean
-  value?: string
-  onChangeText?: (text: string) => void
+  value: string
+  onChangeText: (text: string) => void
+  images: ImageAttachment[]
+  onChangeImages: (updater: ImageListUpdater) => void
+  clearDraft: (lifecycle: 'sent' | 'abandoned') => void
   /** When true, auto-focuses the text input on web. */
   autoFocus?: boolean
   /** Callback to expose the addImages function to parent components */
@@ -80,12 +81,14 @@ const MOBILE_MESSAGE_PLACEHOLDER = 'Message, @files, /commands'
 export function AgentInputArea({
   agentId,
   serverId,
-  draftId,
   onSubmitMessage,
   isSubmitLoading = false,
   blurOnSubmit = false,
   value,
   onChangeText,
+  images,
+  onChangeImages,
+  clearDraft,
   autoFocus = false,
   onAddImages,
   commandDraftConfig,
@@ -114,14 +117,6 @@ 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)
   )
@@ -131,7 +126,6 @@ export function AgentInputArea({
   const setAgentStreamTail = useSessionStore((state) => state.setAgentStreamTail)
   const setAgentStreamHead = useSessionStore((state) => state.setAgentStreamHead)
 
-  const [internalInput, setInternalInput] = useState('')
   const isDesktopWebBreakpoint =
     Platform.OS === 'web' &&
     UnistylesRuntime.breakpoint !== 'xs' &&
@@ -139,17 +133,16 @@ export function AgentInputArea({
   const messagePlaceholder = isDesktopWebBreakpoint
     ? DESKTOP_MESSAGE_PLACEHOLDER
     : MOBILE_MESSAGE_PLACEHOLDER
-  const userInput = value ?? internalInput
-  const setUserInput = onChangeText ?? setInternalInput
+  const userInput = value
+  const setUserInput = onChangeText
+  const selectedImages = images
+  const setSelectedImages = onChangeImages
   const [cursorIndex, setCursorIndex] = useState(0)
   const [isProcessing, setIsProcessing] = useState(false)
-  const [selectedImages, setSelectedImages] = useState<ImageAttachment[]>([])
   const [isCancellingAgent, setIsCancellingAgent] = useState(false)
   const [sendError, setSendError] = useState<string | null>(null)
   const [isMessageInputFocused, setIsMessageInputFocused] = useState(false)
   const messageInputRef = useRef<MessageInputRef>(null)
-  const draftGenerationRef = useRef(0)
-  const hydratedGenerationRef = useRef(0)
   const keyboardHandlerIdRef = useRef(
     `message-input:${serverId}:${agentId}:${Math.random().toString(36).slice(2)}`
   )
@@ -187,7 +180,7 @@ export function AgentInputArea({
   // Expose addImages function to parent for drag-and-drop support
   const addImages = useCallback((images: ImageAttachment[]) => {
     setSelectedImages((prev) => [...prev, ...images])
-  }, [])
+  }, [setSelectedImages])
 
   useEffect(() => {
     onAddImages?.(addImages)
@@ -326,10 +319,7 @@ export function AgentInputArea({
       return next
     })
 
-    const isControlled = value !== undefined
-    if (!isControlled) {
-      setUserInput('')
-    }
+    setUserInput('')
     setSelectedImages([])
   }
 
@@ -354,28 +344,20 @@ export function AgentInputArea({
     // Save values so we can restore on error.
     const savedImages = imageAttachments
     if (!onSubmitMessageRef.current) {
-      if (value !== undefined) {
-        onChangeText?.('')
-      } else {
-        setUserInput('')
-      }
+      setUserInput('')
+      setSelectedImages([])
     }
-    setSelectedImages([])
     setSendError(null)
     setIsProcessing(true)
 
     try {
       await submitMessage(trimmedMessage, imageAttachments)
-      clearDraftInput({ draftKey: draftStoreKey, lifecycle: 'sent' })
+      clearDraft('sent')
     } catch (error) {
       console.error('[AgentInput] Failed to send message:', error)
       // Restore input so the user never loses their message
       if (!onSubmitMessageRef.current) {
-        if (value !== undefined) {
-          onChangeText?.(trimmedMessage)
-        } else {
-          setUserInput(trimmedMessage)
-        }
+        setUserInput(trimmedMessage)
       }
       if (savedImages) {
         setSelectedImages(savedImages)
@@ -434,109 +416,6 @@ export function AgentInputArea({
     }
   }, [isAgentRunning, isConnected])
 
-  // Hydrate draft only when switching agents (uncontrolled mode only)
-  const isControlled = value !== undefined
-  useEffect(() => {
-    // Skip draft hydration for controlled inputs - parent manages state
-    if (isControlled) {
-      return
-    }
-    const generation = beginDraftGeneration(draftStoreKey)
-    draftGenerationRef.current = generation
-    hydratedGenerationRef.current = 0
-    setUserInput('')
-    setSelectedImages([])
-    let cancelled = false
-
-    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 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 ?? []
-    const isSameImages =
-      existingImages.length === selectedImages.length &&
-      existingImages.every((img, idx) => {
-        return (
-          img.id === selectedImages[idx]?.id &&
-          img.mimeType === selectedImages[idx]?.mimeType &&
-          img.storageType === selectedImages[idx]?.storageType &&
-          img.storageKey === selectedImages[idx]?.storageKey
-        )
-      })
-
-    if (isSameText && isSameImages) {
-      return
-    }
-
-    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,
-  ])
-
   const handleKeyboardAction = useCallback((action: KeyboardActionDefinition): boolean => {
     if (!isScreenFocused) {
       return false

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,
   }
 }

packages/app/src/panels/agent-panel.tsx

@@ -20,6 +20,7 @@ import {
 } from "@/hooks/use-agent-screen-state-machine";
 import { useArchiveAgent } from "@/hooks/use-archive-agent";
 import { useDelayedHistoryRefreshToast } from "@/hooks/use-delayed-history-refresh-toast";
+import { useAgentInputDraft } from "@/hooks/use-agent-input-draft";
 import { useKeyboardShiftStyle } from "@/hooks/use-keyboard-shift-style";
 import { usePaneContext } from "@/panels/pane-context";
 import type { PanelDescriptor, PanelRegistration } from "@/panels/panel-registry";
@@ -42,6 +43,7 @@ import {
 import { mergePendingCreateImages } from "@/utils/pending-create-images";
 import { deriveSidebarStateBucket } from "@/utils/sidebar-agent-state";
 import { useCreateFlowStore } from "@/stores/create-flow-store";
+import { buildDraftStoreKey } from "@/stores/draft-keys";
 import { useSessionStore, type Agent } from "@/stores/session-store";
 import type { PendingPermission } from "@/types/shared";
 import type { StreamItem } from "@/types/stream";
@@ -234,6 +236,12 @@ function AgentPanelBody({
     routeKey: string;
     reason: "initial-entry" | "resume";
   } | null>(null);
+  const agentInputDraft = useAgentInputDraft(
+    buildDraftStoreKey({
+      serverId,
+      agentId: agentId ?? "__pending__",
+    })
+  );
 
   const handleFilesDropped = useCallback((files: ImageAttachment[]) => {
     addImagesRef.current?.(files);
@@ -851,6 +859,11 @@ function AgentPanelBody({
             <AgentInputArea
               agentId={agentId}
               serverId={serverId}
+              value={agentInputDraft.text}
+              onChangeText={agentInputDraft.setText}
+              images={agentInputDraft.images}
+              onChangeImages={agentInputDraft.setImages}
+              clearDraft={agentInputDraft.clear}
               autoFocus={isPaneFocused}
               isSubmitLoading={showPendingCreateSubmitLoading}
               onAttentionInputFocus={attentionController.clearOnInputFocus}

packages/app/src/screens/agent/draft-agent-screen.tsx

@@ -32,7 +32,7 @@ import { collectAgentWorkingDirectorySuggestions } from '@/utils/agent-working-d
 import { buildWorkingDirectorySuggestions } from '@/utils/working-directory-suggestions'
 import { useExplorerOpenGesture } from '@/hooks/use-explorer-open-gesture'
 import { useSessionStore } from '@/stores/session-store'
-import { generateDraftId } from '@/stores/draft-keys'
+import { buildDraftStoreKey, generateDraftId } from '@/stores/draft-keys'
 import { getHostRuntimeStore, useHostRuntimeClient, useHostRuntimeIsConnected } from '@/runtime/host-runtime'
 import { ExplorerSidebarAnimationProvider } from '@/contexts/explorer-sidebar-animation-context'
 import { usePanelStore, type ExplorerCheckoutContext } from '@/stores/panel-store'
@@ -50,6 +50,7 @@ import { prepareWorkspaceTab } from '@/utils/workspace-navigation'
 import { useDesktopDragHandlers } from '@/utils/desktop-window'
 import { useKeyboardShiftStyle } from '@/hooks/use-keyboard-shift-style'
 import { normalizeAgentSnapshot } from '@/utils/agent-snapshots'
+import { useAgentInputDraft } from '@/hooks/use-agent-input-draft'
 import { useDraftAgentCreateFlow } from '@/hooks/use-draft-agent-create-flow'
 
 const EMPTY_PENDING_PERMISSIONS = new Map()
@@ -242,6 +243,13 @@ function DraftAgentScreenContent({
   const isExplorerOpen = isMobile ? mobileView === 'file-explorer' : desktopFileExplorerOpen
   const draftIdRef = useRef(generateDraftId())
   const draftAgentIdRef = useRef(generateDraftId())
+  const draftInput = useAgentInputDraft(
+    buildDraftStoreKey({
+      serverId: selectedServerId ?? '',
+      agentId: draftAgentIdRef.current,
+      draftId: draftIdRef.current,
+    })
+  )
 
   const [worktreeMode, setWorktreeMode] = useState<'none' | 'create' | 'attach'>(
     initialWorktreeMode
@@ -766,12 +774,10 @@ function DraftAgentScreenContent({
   ])
 
   const {
-    promptValue,
     formErrorMessage,
     isSubmitting,
     optimisticStreamItems,
     draftAgent,
-    setPromptText,
     handleCreateFromInput,
   } = useDraftAgentCreateFlow<Agent, { id: string; cwd: string }>({
     draftId: draftIdRef.current,
@@ -1184,12 +1190,14 @@ function DraftAgentScreenContent({
               onSubmitMessage={handleCreateFromInput}
               isSubmitLoading={isSubmitting}
               blurOnSubmit={true}
-              value={promptValue}
-              onChangeText={setPromptText}
+              value={draftInput.text}
+              onChangeText={draftInput.setText}
+              images={draftInput.images}
+              onChangeImages={draftInput.setImages}
+              clearDraft={draftInput.clear}
               autoFocus={!isSubmitting}
               onAddImages={handleAddImagesCallback}
               commandDraftConfig={draftCommandConfig}
-              draftId={draftIdRef.current}
               statusControls={{
                 providerDefinitions,
                 selectedProvider,

packages/app/src/screens/workspace/workspace-draft-agent-tab.tsx

@@ -6,8 +6,10 @@ import { FileDropZone } from '@/components/file-drop-zone'
 import { AgentStreamView } from '@/components/agent-stream-view'
 import type { ImageAttachment } from '@/components/message-input'
 import { useAgentFormState } from '@/hooks/use-agent-form-state'
+import { useAgentInputDraft } from '@/hooks/use-agent-input-draft'
 import { useDraftAgentCreateFlow } from '@/hooks/use-draft-agent-create-flow'
 import { useHostRuntimeClient, useHostRuntimeIsConnected } from '@/runtime/host-runtime'
+import { buildDraftStoreKey } from '@/stores/draft-keys'
 import type { Agent } from '@/stores/session-store'
 import { encodeImages } from '@/utils/encode-images'
 import type { AgentCapabilityFlags, AgentSessionConfig } from '@server/server/agent/agent-sdk-types'
@@ -43,6 +45,13 @@ export function WorkspaceDraftAgentTab({
   const client = useHostRuntimeClient(serverId)
   const isConnected = useHostRuntimeIsConnected(serverId)
   const addImagesRef = useRef<((images: ImageAttachment[]) => void) | null>(null)
+  const draftInput = useAgentInputDraft(
+    buildDraftStoreKey({
+      serverId,
+      agentId: tabId,
+      draftId,
+    })
+  )
 
   const {
     selectedProvider,
@@ -81,12 +90,10 @@ export function WorkspaceDraftAgentTab({
   }, [setWorkingDir, workingDir, workspaceId])
 
   const {
-    promptValue,
     formErrorMessage,
     isSubmitting,
     optimisticStreamItems,
     draftAgent,
-    setPromptText,
     handleCreateFromInput,
   } = useDraftAgentCreateFlow<Agent, AgentSnapshotPayload>({
     draftId,
@@ -233,12 +240,14 @@ export function WorkspaceDraftAgentTab({
             onSubmitMessage={handleCreateFromInput}
             isSubmitLoading={isSubmitting}
             blurOnSubmit={true}
-            value={promptValue}
-            onChangeText={setPromptText}
+            value={draftInput.text}
+            onChangeText={draftInput.setText}
+            images={draftInput.images}
+            onChangeImages={draftInput.setImages}
+            clearDraft={draftInput.clear}
             autoFocus={!isSubmitting}
             onAddImages={handleAddImagesCallback}
             commandDraftConfig={draftCommandConfig}
-            draftId={draftId}
             statusControls={{
               providerDefinitions,
               selectedProvider,