cli: harden ask-user state and preserve pending images on followups

Author: brandonkachen

cli/src/chat.tsx

@@ -625,8 +625,9 @@ export const Chat = ({
     ) => {
       ensureQueueActiveBeforeSubmit()
 
+      const preserveInput = options?.preserveInputValue === true
       const previousInputValue =
-        options?.preserveInputValue === true
+        preserveInput
           ? (() => {
               const {
                 inputValue: text,
@@ -636,39 +637,58 @@ export const Chat = ({
               return { text, cursorPosition, lastEditDueToNav }
             })()
           : null
+      const preservedPendingImages =
+        preserveInput && useChatStore.getState().pendingImages.length > 0
+          ? [...useChatStore.getState().pendingImages]
+          : null
 
-      const result = await routeUserPrompt({
-        abortControllerRef,
-        agentMode: mode,
-        inputRef,
-        inputValue: content,
-        isChainInProgressRef,
-        isStreaming,
-        logoutMutation,
-        streamMessageIdRef,
-        addToQueue,
-        clearMessages,
-        saveToHistory,
-        scrollToLatest,
-        sendMessage,
-        setCanProcessQueue,
-        setInputFocused,
-        setInputValue,
-        setIsAuthenticated,
-        setMessages,
-        setUser,
-        stopStreaming,
-      })
+      if (preserveInput && preservedPendingImages) {
+        useChatStore.getState().clearPendingImages()
+      }
 
-      if (previousInputValue) {
-        setInputValue({
-          text: previousInputValue.text,
-          cursorPosition: previousInputValue.cursorPosition,
-          lastEditDueToNav: previousInputValue.lastEditDueToNav,
+      try {
+        const result = await routeUserPrompt({
+          abortControllerRef,
+          agentMode: mode,
+          inputRef,
+          inputValue: content,
+          isChainInProgressRef,
+          isStreaming,
+          logoutMutation,
+          streamMessageIdRef,
+          addToQueue,
+          clearMessages,
+          saveToHistory,
+          scrollToLatest,
+          sendMessage,
+          setCanProcessQueue,
+          setInputFocused,
+          setInputValue,
+          setIsAuthenticated,
+          setMessages,
+          setUser,
+          stopStreaming,
         })
-      }
 
-      return result
+        return result
+      } finally {
+        if (previousInputValue) {
+          setInputValue({
+            text: previousInputValue.text,
+            cursorPosition: previousInputValue.cursorPosition,
+            lastEditDueToNav: previousInputValue.lastEditDueToNav,
+          })
+        }
+
+        if (preserveInput && preservedPendingImages) {
+          const currentPending = useChatStore.getState().pendingImages
+          if (currentPending.length === 0) {
+            useChatStore.setState((state) => {
+              state.pendingImages = preservedPendingImages
+            })
+          }
+        }
+      }
     },
   )
 

cli/src/components/ask-user/index.tsx

@@ -124,7 +124,7 @@ export const MultipleChoiceForm: React.FC<MultipleChoiceFormProps> = ({
     ) => {
       setAnswers((prev) => {
         const nextAnswers = new Map(prev)
-        const previousAnswer = prev.get(questionIndex)
+        const previousAnswer = prev.get(questionIndex) ?? {}
         nextAnswers.set(questionIndex, updater(previousAnswer))
         return nextAnswers
       })
@@ -226,7 +226,7 @@ export const MultipleChoiceForm: React.FC<MultipleChoiceFormProps> = ({
 
       setAnswers((prev) => {
         const newAnswers = new Map(prev)
-        const currentAnswer = prev.get(questionIndex)
+        const currentAnswer: AccordionAnswer = prev.get(questionIndex) ?? {}
 
         if (optionIndex === OTHER_OPTION_INDEX) {
           toggledOtherOn = !(currentAnswer?.isOther ?? false)