Refactor agent input draft persistence

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,