From 3fef38743d42e3041fa46f04603468a7017cbe6f Mon Sep 17 00:00:00 2001
From: Charles Vien <charles.v@posthog.com>
Date: Mon, 12 Jan 2026 16:54:30 -0800
Subject: [PATCH] Restore draft on mount or when sessionId/editor changes

---
 .../message-editor/stores/draftStore.ts       | 11 +++++++---
 .../message-editor/tiptap/useDraftSync.ts     | 22 +++++++++++++++++--
 2 files changed, 28 insertions(+), 5 deletions(-)

diff --git a/apps/array/src/renderer/features/message-editor/stores/draftStore.ts b/apps/array/src/renderer/features/message-editor/stores/draftStore.ts
index 243c21713..b994e96fa 100644
--- a/apps/array/src/renderer/features/message-editor/stores/draftStore.ts
+++ b/apps/array/src/renderer/features/message-editor/stores/draftStore.ts
@@ -1,4 +1,5 @@
 import type { AvailableCommand } from "@agentclientprotocol/sdk";
+import { electronStorage } from "@renderer/lib/electronStorage";
 import { create } from "zustand";
 import { persist } from "zustand/middleware";
 import { immer } from "zustand/middleware/immer";
@@ -52,16 +53,19 @@ export const useDraftStore = create<DraftStore>()(
       _hasHydrated: false,
 
       actions: {
-        setHasHydrated: (hydrated) => set({ _hasHydrated: hydrated }),
+        setHasHydrated: (hydrated) => {
+          set({ _hasHydrated: hydrated });
+        },
 
-        setDraft: (sessionId, draft) =>
+        setDraft: (sessionId, draft) => {
           set((state) => {
             if (draft === null) {
               delete state.drafts[sessionId];
             } else {
               state.drafts[sessionId] = draft;
             }
-          }),
+          });
+        },
 
         getDraft: (sessionId) => get().drafts[sessionId] ?? null,
 
@@ -110,6 +114,7 @@ export const useDraftStore = create<DraftStore>()(
     })),
     {
       name: "message-editor-drafts",
+      storage: electronStorage,
       partialize: (state) => ({ drafts: state.drafts }),
       onRehydrateStorage: () => (state) => {
         state?.actions.setHasHydrated(true);
diff --git a/apps/array/src/renderer/features/message-editor/tiptap/useDraftSync.ts b/apps/array/src/renderer/features/message-editor/tiptap/useDraftSync.ts
index 615fb5439..7ae5b7a07 100644
--- a/apps/array/src/renderer/features/message-editor/tiptap/useDraftSync.ts
+++ b/apps/array/src/renderer/features/message-editor/tiptap/useDraftSync.ts
@@ -103,6 +103,8 @@ export function useDraftSync(
   context?: DraftContext,
 ) {
   const hasRestoredRef = useRef(false);
+  const lastSessionIdRef = useRef(sessionId);
+  const lastEditorRef = useRef(editor);
   const editorRef = useRef(editor);
   editorRef.current = editor;
 
@@ -110,6 +112,18 @@ export function useDraftSync(
   const draft = useDraftStore((s) => s.drafts[sessionId] ?? null);
   const hasHydrated = useDraftStore((s) => s._hasHydrated);
 
+  // Reset restoration flag when sessionId changes (e.g., navigating between tasks)
+  if (lastSessionIdRef.current !== sessionId) {
+    lastSessionIdRef.current = sessionId;
+    hasRestoredRef.current = false;
+  }
+
+  // Reset restoration flag when editor instance changes (e.g., when disabled state changes)
+  if (lastEditorRef.current !== editor && editor !== null) {
+    lastEditorRef.current = editor;
+    hasRestoredRef.current = false;
+  }
+
   // Set context for this session
   useLayoutEffect(() => {
     draftActions.setContext(sessionId, {
@@ -121,7 +135,7 @@ export function useDraftSync(
     };
   }, [sessionId, context?.taskId, context?.repoPath, draftActions]);
 
-  // Restore draft on mount
+  // Restore draft on mount or when sessionId/editor changes
   useLayoutEffect(() => {
     if (!hasHydrated || !editor || hasRestoredRef.current) return;
     if (!draft || isContentEmpty(draft)) return;
@@ -137,6 +151,10 @@ export function useDraftSync(
 
   const saveDraft = useCallback(
     (e: Editor) => {
+      // Don't save until store has hydrated from storage
+      // This prevents overwriting stored drafts with empty content before restoration
+      if (!hasHydrated) return;
+
       const json = e.getJSON();
       const content = tiptapJsonToEditorContent(json);
       draftActions.setDraft(
@@ -144,7 +162,7 @@ export function useDraftSync(
         isContentEmpty(content) ? null : content,
       );
     },
-    [sessionId, draftActions],
+    [sessionId, draftActions, hasHydrated],
   );
 
   const clearDraft = useCallback(() => {