fix: add undo feature to reply mails and fix default mail for replies

apps/mail/components/mail/reply-composer.tsx

@@ -1,4 +1,4 @@
-import { useUndoSend } from '@/hooks/use-undo-send';
+import { useUndoSend, deserializeFiles, type EmailData, type ReplyMode, type UndoContext } from '@/hooks/use-undo-send';
 import { constructReplyBody, constructForwardBody } from '@/lib/utils';
 import { useActiveConnection } from '@/hooks/use-connections';
 import { useEmailAliases } from '@/hooks/use-email-aliases';
@@ -14,7 +14,7 @@ import { useDraft } from '@/hooks/use-drafts';
 import { m } from '@/paraglide/messages';
 import type { Sender } from '@/types';
 import { useQueryState } from 'nuqs';
-import { useEffect } from 'react';
+import { useEffect, useMemo, useSyncExternalStore } from 'react';
 import posthog from 'posthog-js';
 import { toast } from 'sonner';
 
@@ -29,7 +29,7 @@ export default function ReplyCompose({ messageId }: ReplyComposeProps) {
 
   const [draftId, setDraftId] = useQueryState('draftId');
   const [threadId] = useQueryState('threadId');
-  const [, setActiveReplyId] = useQueryState('activeReplyId');
+  const [activeReplyId, setActiveReplyId] = useQueryState('activeReplyId');
   const { data: emailData, refetch, latestDraft } = useThread(threadId);
   const { data: draft } = useDraft(draftId ?? null);
   const trpc = useTRPC();
@@ -43,59 +43,122 @@ export default function ReplyCompose({ messageId }: ReplyComposeProps) {
   const replyToMessage =
     (messageId && emailData?.messages.find((msg) => msg.id === messageId)) || emailData?.latest;
 
-  // Initialize recipients and subject when mode changes
-  useEffect(() => {
-    if (!replyToMessage || !mode || !activeConnection?.email) return;
+  const undoReplyRaw = useSyncExternalStore(
+    (listener) => {
+      if (typeof window === 'undefined') return () => {};
+      const handler = () => listener();
+      window.addEventListener('storage', handler);
+      window.addEventListener('zero:undoReplyEmailData', handler as EventListener);
+      return () => {
+        window.removeEventListener('storage', handler);
+        window.removeEventListener('zero:undoReplyEmailData', handler as EventListener);
+      };
+    },
+    () => {
+      try {
+        return typeof window !== 'undefined'
+          ? window.localStorage.getItem('undoReplyEmailData')
+          : null;
+      } catch {
+        return null;
+      }
+    },
+    () => null,
+  );
+
+  const undoReplyEmailData = useMemo((): EmailData | null => {
+    if (!mode) return null;
+    if (!undoReplyRaw) return null;
+    try {
+      const parsed = JSON.parse(undoReplyRaw);
+      const ctx = parsed?.__replyContext as
+        | { threadId: string; activeReplyId: string; mode: string; draftId?: string | null }
+        | undefined;
+
+      const currentThread = threadId || replyToMessage?.threadId || '';
+      const currentReply = messageId || activeReplyId || replyToMessage?.id || '';
+      const matches = !!ctx && ctx.threadId === currentThread && ctx.activeReplyId === currentReply;
+      if (!matches) return null;
+
+      if (parsed.attachments && Array.isArray(parsed.attachments)) {
+        parsed.attachments = deserializeFiles(parsed.attachments);
+      }
+      return parsed as EmailData;
+    } catch (err) {
+      console.error('Failed to parse undo reply email data:', err);
+      return null;
+    }
+  }, [undoReplyRaw, mode, threadId, messageId, activeReplyId, replyToMessage?.id, replyToMessage?.threadId]);
+
+  const { defaultTo, defaultCc, defaultSubject } = useMemo(() => {
+    const result = { defaultTo: [] as string[], defaultCc: [] as string[], defaultSubject: '' };
+    if (!replyToMessage || !mode || !activeConnection?.email) {
+      return result;
+    }
 
     const userEmail = activeConnection.email.toLowerCase();
     const senderEmail = replyToMessage.sender.email.toLowerCase();
 
-    // Set subject based on mode
+    const baseSubject = replyToMessage.subject || '';
+    const lower = baseSubject.trim().toLowerCase();
+    const hasRePrefix = lower.startsWith('re:');
+    const hasFwdPrefix = lower.startsWith('fwd:') || lower.startsWith('fw:');
+    if (mode === 'forward') {
+      result.defaultSubject = hasFwdPrefix ? baseSubject : `Fwd: ${baseSubject}`.trim();
+    } else {
+      result.defaultSubject = hasRePrefix ? baseSubject : `Re: ${baseSubject}`.trim();
+    }
 
     if (mode === 'reply') {
-      // Reply to sender
-      const to: string[] = [];
-
-      // If the sender is not the current user, add them to the recipients
       if (senderEmail !== userEmail) {
-        to.push(replyToMessage.sender.email);
+        result.defaultTo.push(replyToMessage.sender.email);
       } else if (replyToMessage.to && replyToMessage.to.length > 0 && replyToMessage.to[0]?.email) {
-        // If we're replying to our own email, reply to the first recipient
-        to.push(replyToMessage.to[0].email);
+        result.defaultTo.push(replyToMessage.to[0].email);
       }
+      return result;
+    }
 
-      // Initialize email composer with these recipients
-      // Note: The actual initialization happens in the EmailComposer component
-    } else if (mode === 'replyAll') {
-      const to: string[] = [];
-      const cc: string[] = [];
+    if (mode === 'replyAll') {
+      const seen = new Set<string>();
+
+      const pushIfNew = (email: string, target: 'to' | 'cc') => {
+        const key = email.toLowerCase();
+        if (seen.has(key)) return;
+        seen.add(key);
+        if (target === 'to') {
+          result.defaultTo.push(email);
+        } else {
+          result.defaultCc.push(email);
+        }
+      };
 
       // Add original sender if not current user
       if (senderEmail !== userEmail) {
-        to.push(replyToMessage.sender.email);
+        pushIfNew(replyToMessage.sender.email, 'to');
       }
 
       // Add original recipients from To field
       replyToMessage.to?.forEach((recipient) => {
-        const recipientEmail = recipient.email.toLowerCase();
-        if (recipientEmail !== userEmail && recipientEmail !== senderEmail) {
-          to.push(recipient.email);
+        const recipientEmail = recipient.email;
+        const key = recipientEmail.toLowerCase();
+        if (key !== userEmail && key !== senderEmail) {
+          pushIfNew(recipientEmail, 'to');
         }
       });
 
       // Add CC recipients
       replyToMessage.cc?.forEach((recipient) => {
-        const recipientEmail = recipient.email.toLowerCase();
-        if (recipientEmail !== userEmail && !to.includes(recipient.email)) {
-          cc.push(recipient.email);
+        const recipientEmail = recipient.email;
+        const key = recipientEmail.toLowerCase();
+        if (key !== userEmail) {
+          pushIfNew(recipientEmail, 'cc');
         }
       });
 
-      // Initialize email composer with these recipients
-    } else if (mode === 'forward') {
-      // For forward, we start with empty recipients
-      // Just set the subject and include the original message
+      return result;
     }
+
+    return result;
   }, [mode, replyToMessage, activeConnection?.email]);
 
   const handleSendEmail = async (data: {
@@ -219,7 +282,13 @@ export default function ReplyCompose({ messageId }: ReplyComposeProps) {
         message: data.message,
         attachments: data.attachments,
         scheduleAt: data.scheduleAt,
-      });
+      }, {
+        kind: 'reply',
+        threadId: replyToMessage.threadId || threadId || '',
+        mode: (mode ?? 'reply') as ReplyMode,
+        activeReplyId: replyToMessage.id,
+        draftId,
+      } satisfies UndoContext);
     } catch (error) {
       console.error('Error sending email:', error);
       toast.error(m['pages.createEmail.failedToSendEmail']());
@@ -257,19 +326,30 @@ export default function ReplyCompose({ messageId }: ReplyComposeProps) {
   return (
     <div className="w-full rounded-2xl overflow-visible border">
       <EmailComposer
+        key={draftId || replyToMessage?.id || 'reply-composer'}
         editorClassName="min-h-[50px]"
         className="w-full max-w-none! pb-1 overflow-visible"
         onSendEmail={handleSendEmail}
         onClose={async () => {
           setMode(null);
           setDraftId(null);
           setActiveReplyId(null);
+          if (typeof window !== 'undefined') {
+            localStorage.removeItem('undoReplyEmailData');
+          }
         }}
-        initialMessage={draft?.content ?? latestDraft?.decodedBody}
-        initialTo={ensureEmailArray(draft?.to)}
-        initialCc={ensureEmailArray(draft?.cc)}
-        initialBcc={ensureEmailArray(draft?.bcc)}
-        initialSubject={draft?.subject}
+        initialMessage={undoReplyEmailData?.message ?? draft?.content ?? latestDraft?.decodedBody}
+        initialTo={
+          undoReplyEmailData?.to ??
+          (ensureEmailArray(draft?.to).length ? ensureEmailArray(draft?.to) : defaultTo)
+        }
+        initialCc={
+          undoReplyEmailData?.cc ??
+          (ensureEmailArray(draft?.cc).length ? ensureEmailArray(draft?.cc) : defaultCc)
+        }
+        initialBcc={undoReplyEmailData?.bcc ?? ensureEmailArray(draft?.bcc)}
+        initialSubject={undoReplyEmailData?.subject ?? draft?.subject ?? defaultSubject}
+        initialAttachments={undoReplyEmailData?.attachments}
         autofocus={true}
         settingsLoading={settingsLoading}
         replyingTo={replyToMessage?.sender.email}

apps/mail/hooks/use-undo-send.ts

@@ -1,6 +1,6 @@
 import { useMutation } from '@tanstack/react-query';
 import { toast } from 'sonner';
-
+import { useQueryState } from 'nuqs';
 import { useTRPC } from '@/providers/query-provider';
 import { isSendResult } from '@/lib/email-utils';
 import type { UserSettings } from '@zero/server/schemas';
@@ -28,6 +28,18 @@ type SerializableEmailData = Omit<EmailData, 'attachments'> & {
   attachments: SerializedFile[];
 };
 
+export type ReplyMode = 'reply' | 'replyAll' | 'forward';
+
+export type UndoContext =
+  | { kind: 'compose' }
+  | {
+      kind: 'reply';
+      threadId: string;
+      mode: ReplyMode;
+      activeReplyId: string;
+      draftId?: string | null;
+    };
+
 const serializeFiles = async (files: File[]): Promise<SerializedFile[]> => {
   return Promise.all(
     files.map(async (file) => ({
@@ -60,10 +72,17 @@ export const useUndoSend = () => {
   const trpc = useTRPC();
   const { mutateAsync: unsendEmail } = useMutation(trpc.mail.unsend.mutationOptions());
 
+  const [, setIsComposeOpen] = useQueryState('isComposeOpen');
+  const [, setThreadId] = useQueryState('threadId');
+  const [, setActiveReplyId] = useQueryState('activeReplyId');
+  const [, setMode] = useQueryState('mode');
+  const [, setDraftId] = useQueryState('draftId');
+
   const handleUndoSend = (
-    result: unknown, 
+    result: unknown,
     settings: { settings: UserSettings } | undefined,
-    emailData?: EmailData
+    emailData?: EmailData,
+    context: UndoContext = { kind: 'compose' },
   ) => {
     if (isSendResult(result) && settings?.settings?.undoSendEnabled) {
       const { messageId, sendAt } = result;
@@ -84,15 +103,34 @@ export const useUndoSend = () => {
                     ...emailData,
                     attachments: serializedAttachments,
                   };
-                  localStorage.setItem('undoEmailData', JSON.stringify(serializableData));
+                  if (context.kind === 'reply') {
+                    const withContext = {
+                      ...serializableData,
+                      __replyContext: {
+                        threadId: context.threadId,
+                        activeReplyId: context.activeReplyId,
+                        mode: context.mode,
+                        draftId: context.draftId ?? null,
+                      },
+                    } as const;
+                    localStorage.setItem('undoReplyEmailData', JSON.stringify(withContext));
+                  } else {
+                    localStorage.setItem('undoEmailData', JSON.stringify(serializableData));
+                  }
                 }
 
-                const url = new URL(window.location.href);
-                url.searchParams.delete('activeReplyId');
-                url.searchParams.delete('mode');
-                url.searchParams.delete('draftId');
-                url.searchParams.set('isComposeOpen', 'true');
-                window.history.replaceState({}, '', url.toString());
+                if (context.kind === 'reply') {
+                  setIsComposeOpen(null);
+                  setThreadId(context.threadId);
+                  setActiveReplyId(context.activeReplyId);
+                  setMode(context.mode);
+                  setDraftId(context.draftId ?? null);
+                } else {
+                  setActiveReplyId(null);
+                  setMode(null);
+                  setDraftId(null);
+                  setIsComposeOpen('true');
+                }
 
                 toast.info('Send cancelled');
               } catch {