feat: preserve input text draft across navigation

gui/src/components/mainInput/TipTapEditor/TipTapEditor.tsx

@@ -95,12 +95,6 @@ function TipTapEditorInner(props: TipTapEditorProps) {
     }
   }, [editor, props.placeholder, historyLength]);
 
-  useEffect(() => {
-    if (props.isMainInput) {
-      editor?.commands.clearContent(true);
-    }
-  }, [editor, props.isMainInput]);
-
   useEffect(() => {
     if (isInEdit) {
       setShouldHideToolbar(false);

gui/src/components/mainInput/TipTapEditor/utils/editorConfig.ts

@@ -22,6 +22,7 @@ import { selectSelectedChatModel } from "../../../../redux/slices/configSlice";
 import { AppDispatch } from "../../../../redux/store";
 import { exitEdit } from "../../../../redux/thunks/edit";
 import { getFontSize, isJetBrains } from "../../../../util";
+import { setLocalStorage } from "../../../../util/localStorage";
 import { CodeBlock, Mention, PromptBlock, SlashCommand } from "../extensions";
 import { TipTapEditorProps } from "../TipTapEditor";
 import {
@@ -392,6 +393,33 @@ export function createEditorConfig(options: {
     },
     content: props.editorState,
     editable: !isStreaming || props.isMainInput,
+    onUpdate: ({ editor }) => {
+      const content = editor.getJSON();
+      if (props.isMainInput) {
+        if (hasValidEditorContent(content)) {
+          setLocalStorage(`inputDraft_${props.historyKey}`, content);
+          localStorage.removeItem(`editingDraft_${props.historyKey}`);
+        } else {
+          // clear draft if content is empty
+          localStorage.removeItem(`inputDraft_${props.historyKey}`);
+        }
+      } else {
+        if (hasValidEditorContent(content)) {
+          const scrollContainer = document.getElementById(
+            "chat-scroll-container",
+          );
+          const scrollTop = scrollContainer?.scrollTop ?? 0;
+          setLocalStorage(`editingDraft_${props.historyKey}`, {
+            content,
+            messageId: props.inputId,
+            scrollTop,
+          });
+          localStorage.removeItem(`inputDraft_${props.historyKey}`);
+        } else {
+          localStorage.removeItem(`editingDraft_${props.historyKey}`);
+        }
+      }
+    },
   });
 
   const onEnter = (modifiers: InputModifiers) => {
@@ -409,9 +437,12 @@ export function createEditorConfig(options: {
       return;
     }
 
+    // clear draft from localStorage after successful submission
     if (props.isMainInput) {
       addRef.current(json);
     }
+    localStorage.removeItem(`inputDraft_${props.historyKey}`);
+    localStorage.removeItem(`editingDraft_${props.historyKey}`);
 
     props.onEnter(json, modifiers, editor);
   };

gui/src/hooks/ParallelListeners.tsx

@@ -267,6 +267,13 @@ function ParallelListeners() {
     migrateLocalStorage(dispatch);
   }, []);
 
+  useEffect(() => {
+    return () => {
+      localStorage.removeItem("editingDraft_edit");
+      localStorage.removeItem("editingDraft_chat");
+    };
+  });
+
   return <></>;
 }
 

gui/src/pages/gui/Chat.tsx

@@ -55,7 +55,11 @@ import { resolveEditorContent } from "../../components/mainInput/TipTapEditor/ut
 import { setDialogMessage, setShowDialog } from "../../redux/slices/uiSlice";
 import { RootState } from "../../redux/store";
 import { cancelStream } from "../../redux/thunks/cancelStream";
-import { getLocalStorage, setLocalStorage } from "../../util/localStorage";
+import {
+  getLocalStorage,
+  InputDraftWithPosition,
+  setLocalStorage,
+} from "../../util/localStorage";
 import { EmptyChatBody } from "./EmptyChatBody";
 import { ExploreDialogWatcher } from "./ExploreDialogWatcher";
 import { useAutoScroll } from "./useAutoScroll";
@@ -119,12 +123,13 @@ export function Chat() {
   const mainTextInputRef = useRef<HTMLInputElement>(null);
   const stepsDivRef = useRef<HTMLDivElement>(null);
   const tabsRef = useRef<HTMLDivElement>(null);
+  const timerRef = useRef<NodeJS.Timeout | undefined>(undefined);
   const history = useAppSelector((state) => state.session.history);
   const showChatScrollbar = useAppSelector(
     (state) => state.config.config.ui?.showChatScrollbar,
   );
-  const codeToEdit = useAppSelector((state) => state.editModeState.codeToEdit);
   const isInEdit = useAppSelector((store) => store.session.isInEdit);
+  const sessionId = useAppSelector((state) => state.session.id);
 
   const lastSessionId = useAppSelector((state) => state.session.lastSessionId);
   const allSessionMetadata = useAppSelector(
@@ -134,13 +139,38 @@ export function Chat() {
     (state) => state.ui.hasDismissedExploreDialog,
   );
   const mode = useAppSelector((state) => state.session.mode);
-  const currentOrg = useAppSelector(selectCurrentOrg);
   const jetbrains = useMemo(() => {
     return isJetBrains();
   }, []);
 
   useAutoScroll(stepsDivRef, history);
 
+  useEffect(() => {
+    const historyKey = isInEdit ? "edit" : "chat";
+    const savedDraft = getLocalStorage(`editingDraft_${historyKey}`) as
+      | InputDraftWithPosition
+      | undefined;
+    if (savedDraft && savedDraft.messageId && stepsDivRef.current) {
+      timerRef.current = setTimeout(() => {
+        // scroll to and focus on the message being edited
+        requestAnimationFrame(() => {
+          if (stepsDivRef.current && savedDraft.scrollTop !== undefined) {
+            stepsDivRef.current.scrollTop = savedDraft.scrollTop;
+          }
+          const editorElement = document.querySelector(
+            `[data-testid="editor-input-${savedDraft.messageId}"]`,
+          ) as HTMLElement;
+          if (editorElement) {
+            editorElement.focus();
+          }
+        });
+      }, 100);
+    }
+    return () => {
+      clearTimeout(timerRef.current);
+    };
+  }, [sessionId, isInEdit]);
+
   useEffect(() => {
     // Cmd + Backspace to delete current step
     const listener = (e: KeyboardEvent) => {
@@ -337,14 +367,21 @@ export function Chat() {
         latestSummaryIndex !== -1 && index < latestSummaryIndex;
 
       if (message.role === "user") {
+        const historyKey = isInEdit ? "edit" : "chat";
+        const savedDraft = getLocalStorage(`editingDraft_${historyKey}`) as
+          | InputDraftWithPosition
+          | undefined;
+        const draftContent =
+          savedDraft?.messageId === message.id ? savedDraft.content : undefined;
+
         return (
           <ContinueInputBox
             onEnter={(editorState, modifiers) =>
               sendInput(editorState, modifiers, index)
             }
             isLastUserInput={isLastUserInput(index)}
             isMainInput={false}
-            editorState={editorState ?? item.message.content}
+            editorState={draftContent ?? editorState ?? item.message.content}
             contextItems={contextItems}
             appliedRules={appliedRules}
             inputId={message.id}
@@ -439,6 +476,7 @@ export function Chat() {
       {widget}
 
       <StepsDiv
+        id="chat-scroll-container"
         ref={stepsDivRef}
         className={`overflow-y-scroll pt-[8px] ${showScrollbar ? "thin-scrollbar" : "no-scrollbar"} ${history.length > 0 ? "flex-1" : ""}`}
       >
@@ -471,6 +509,9 @@ export function Chat() {
           onEnter={(editorState, modifiers, editor) =>
             sendInput(editorState, modifiers, undefined, editor)
           }
+          editorState={getLocalStorage(
+            `inputDraft_${isInEdit ? "edit" : "chat"}`,
+          )}
           inputId={MAIN_EDITOR_INPUT_ID}
         />
 

gui/src/util/localStorage.ts

@@ -1,6 +1,12 @@
 import { JSONContent } from "@tiptap/react";
 import { OnboardingStatus } from "../components/OnboardingCard";
 
+export type InputDraftWithPosition = {
+  content: JSONContent;
+  messageId: string;
+  scrollTop: number;
+};
+
 type LocalStorageTypes = {
   isExploreDialogOpen: boolean;
   hasDismissedExploreDialog: boolean;
@@ -11,6 +17,8 @@ type LocalStorageTypes = {
   vsCodeUriScheme: string;
   fontSize: number;
   [key: `inputHistory_${string}`]: JSONContent[];
+  [key: `inputDraft_${string}`]: JSONContent;
+  [key: `editingDraft_${string}`]: InputDraftWithPosition;
   extensionVersion: string;
   showTutorialCard: boolean;
   shownProfilesIntroduction: boolean;