optimistic thread state for thread context menu (#1300)

# Implement Optimistic UI for Label Management in Email Threads

This PR adds optimistic UI updates for label management in email threads context menu.

## Type of Change

- ✨ New feature (non-breaking change which adds functionality)

## Areas Affected

- [x] User Interface/Experience
- [x] Data Storage/Management

## Testing Done

- [x] Manual testing performed

<!-- This is an auto-generated comment: release notes by coderabbit.ai -->
## Summary by CodeRabbit

- **New Features**
  - Added support for bulk label toggling on multiple mail threads.
  - Labels now update instantly in the UI with optimistic updates, reflecting changes before server confirmation.

- **Improvements**
  - Checkbox states for labels now accurately reflect pending changes for a smoother user experience.
  - Enhanced feedback with toast messages indicating label changes, including support for multiple threads.

- **Style**
  - Minor formatting improvements in context menu components.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

apps/mail/components/context/thread-context.tsx

@@ -27,8 +27,6 @@ import { useOptimisticActions } from '@/hooks/use-optimistic-actions';
 import { type ThreadDestination } from '@/lib/thread-actions';
 import { useThread, useThreads } from '@/hooks/use-threads';
 import { ExclamationCircle, Mail } from '../icons/icons';
-import { useTRPC } from '@/providers/query-provider';
-import { useMutation } from '@tanstack/react-query';
 import { useMemo, type ReactNode } from 'react';
 import { useLabels } from '@/hooks/use-labels';
 import { FOLDERS, LABELS } from '@/lib/utils';
@@ -59,52 +57,69 @@ interface EmailContextMenuProps {
   refreshCallback?: () => void;
 }
 
-const LabelsList = ({ threadId }: { threadId: string }) => {
+const LabelsList = ({ threadId, bulkSelected }: { threadId: string; bulkSelected: string[] }) => {
   const { data: labels } = useLabels();
-  const { data: thread, refetch } = useThread(threadId);
-  const t = useTranslations();
-  const trpc = useTRPC();
-  const { mutateAsync: modifyLabels } = useMutation(trpc.mail.modifyLabels.mutationOptions());
+  const { optimisticToggleLabel } = useOptimisticActions();
+  const targetThreadIds = bulkSelected.length > 0 ? bulkSelected : [threadId];
+
+  const { data: thread } = useThread(threadId);
+  const rightClickedThreadOptimisticState = useOptimisticThreadState(threadId);
 
   if (!labels || !thread) return null;
 
   const handleToggleLabel = async (labelId: string) => {
     if (!labelId) return;
-    const hasLabel = thread.labels?.map((label) => label.id).includes(labelId);
-    const promise = modifyLabels({
-      threadId: [threadId],
-      addLabels: hasLabel ? [] : [labelId],
-      removeLabels: hasLabel ? [labelId] : [],
-    });
-    toast.promise(promise, {
-      error: hasLabel ? 'Failed to remove label' : 'Failed to add label',
-      finally: async () => {
-        await refetch();
-      },
-    });
+
+    let shouldAddLabel = false;
+
+    let hasLabel = thread.labels?.map((label) => label.id).includes(labelId) || false;
+
+    if (rightClickedThreadOptimisticState.optimisticLabels) {
+      if (rightClickedThreadOptimisticState.optimisticLabels.addedLabelIds.includes(labelId)) {
+        hasLabel = true;
+      } else if (
+        rightClickedThreadOptimisticState.optimisticLabels.removedLabelIds.includes(labelId)
+      ) {
+        hasLabel = false;
+      }
+    }
+
+    shouldAddLabel = !hasLabel;
+
+    optimisticToggleLabel(targetThreadIds, labelId, shouldAddLabel);
   };
 
   return (
     <>
       {labels
         .filter((label) => label.id)
-        .map((label) => (
-          <ContextMenuItem
-            key={label.id}
-            onClick={() => label.id && handleToggleLabel(label.id)}
-            className="font-normal"
-          >
-            <div className="flex items-center">
-              <Checkbox
-                checked={
-                  label.id ? thread.labels?.map((label) => label.id).includes(label.id) : false
-                }
-                className="mr-2 h-4 w-4"
-              />
-              {label.name}
-            </div>
-          </ContextMenuItem>
-        ))}
+        .map((label) => {
+          let isChecked = label.id ? thread.labels?.map((l) => l.id).includes(label.id) : false;
+
+          const checkboxOptimisticState = useOptimisticThreadState(threadId);
+          if (label.id && checkboxOptimisticState.optimisticLabels) {
+            if (checkboxOptimisticState.optimisticLabels.addedLabelIds.includes(label.id)) {
+              isChecked = true;
+            } else if (
+              checkboxOptimisticState.optimisticLabels.removedLabelIds.includes(label.id)
+            ) {
+              isChecked = false;
+            }
+          }
+
+          return (
+            <ContextMenuItem
+              key={label.id}
+              onClick={() => label.id && handleToggleLabel(label.id)}
+              className="font-normal"
+            >
+              <div className="flex items-center">
+                <Checkbox checked={isChecked} className="mr-2 h-4 w-4" />
+                {label.name}
+              </div>
+            </ContextMenuItem>
+          );
+        })}
     </>
   );
 };
@@ -410,7 +425,7 @@ export function ThreadContextMenu({
     {
       id: 'toggle-important',
       label: isImportant ? t('common.mail.removeFromImportant') : t('common.mail.markAsImportant'),
-      icon: <ExclamationCircle className='mr-2.5 h-4 w-4 opacity-60' />,
+      icon: <ExclamationCircle className="mr-2.5 h-4 w-4 opacity-60" />,
       action: handleToggleImportant,
     },
     {
@@ -446,7 +461,7 @@ export function ThreadContextMenu({
         {children}
       </ContextMenuTrigger>
       <ContextMenuContent
-        className="dark:bg-panelDark w-56 overflow-y-auto bg-white "
+        className="dark:bg-panelDark w-56 overflow-y-auto bg-white"
         onContextMenu={(e) => e.preventDefault()}
       >
         {primaryActions.map(renderAction)}
@@ -459,7 +474,7 @@ export function ThreadContextMenu({
             {t('common.mail.labels')}
           </ContextMenuSubTrigger>
           <ContextMenuSubContent className="dark:bg-panelDark max-h-[520px] w-48 overflow-y-auto bg-white">
-            <LabelsList threadId={threadId} />
+            <LabelsList threadId={threadId} bulkSelected={mail.bulkSelected} />
           </ContextMenuSubContent>
         </ContextMenuSub>
 

apps/mail/components/mail/mail-list.tsx

@@ -124,19 +124,35 @@ const Thread = memo(
     const optimisticLabels = useMemo(() => {
       if (!getThreadData?.labels) return [];
 
-      const labels = [...getThreadData.labels];
+      let labels = [...getThreadData.labels];
       const hasStarredLabel = labels.some((label) => label.name === 'STARRED');
 
       if (optimisticState.optimisticStarred !== null) {
         if (optimisticState.optimisticStarred && !hasStarredLabel) {
           labels.push({ id: 'starred-optimistic', name: 'STARRED' });
         } else if (!optimisticState.optimisticStarred && hasStarredLabel) {
-          return labels.filter((label) => label.name !== 'STARRED');
+          labels = labels.filter((label) => label.name !== 'STARRED');
         }
       }
 
+      if (optimisticState.optimisticLabels) {
+        labels = labels.filter(
+          (label) => !optimisticState.optimisticLabels.removedLabelIds.includes(label.id),
+        );
+
+        optimisticState.optimisticLabels.addedLabelIds.forEach((labelId) => {
+          if (!labels.some((label) => label.id === labelId)) {
+            labels.push({ id: labelId, name: labelId });
+          }
+        });
+      }
+
       return labels;
-    }, [getThreadData?.labels, optimisticState.optimisticStarred]);
+    }, [
+      getThreadData?.labels,
+      optimisticState.optimisticStarred,
+      optimisticState.optimisticLabels,
+    ]);
 
     const { optimisticToggleStar, optimisticToggleImportant, optimisticMoveThreadsTo } =
       useOptimisticActions();
@@ -216,7 +232,7 @@ const Thread = memo(
     }, [latestMessage?.body, latestMessage?.sender?.email, settingsData?.settings, queryClient]);
 
     const { labels: threadLabels } = useThreadLabels(
-      getThreadData?.labels ? getThreadData.labels.map((l) => l.id) : [],
+      optimisticLabels ? optimisticLabels.map((l) => l.id) : [],
     );
 
     const mainSearchTerm = useMemo(() => {

apps/mail/components/mail/optimistic-thread-state.tsx

@@ -31,6 +31,10 @@ export function useOptimisticThreadState(threadId: string) {
       optimisticRead: null as boolean | null,
       optimisticDestination: null as string | null,
       optimisticImportant: null as boolean | null,
+      optimisticLabels: {
+        addedLabelIds: [] as string[],
+        removedLabelIds: [] as string[],
+      },
     };
 
     if (!isAffectedByOptimisticAction || !optimisticActions || optimisticActions.length === 0) {
@@ -57,6 +61,11 @@ export function useOptimisticThreadState(threadId: string) {
 
         case 'LABEL':
           states.isAddingLabel = action.add;
+          if (action.add) {
+            states.optimisticLabels.addedLabelIds.push(...action.labelIds);
+          } else {
+            states.optimisticLabels.removedLabelIds.push(...action.labelIds);
+          }
           break;
 
         case 'IMPORTANT':

apps/mail/hooks/use-optimistic-actions.ts

@@ -32,6 +32,7 @@ export function useOptimisticActions() {
   const { mutateAsync: bulkArchive } = useMutation(trpc.mail.bulkArchive.mutationOptions());
   const { mutateAsync: bulkStar } = useMutation(trpc.mail.bulkStar.mutationOptions());
   const { mutateAsync: bulkDeleteThread } = useMutation(trpc.mail.bulkDelete.mutationOptions());
+  const { mutateAsync: modifyLabels } = useMutation(trpc.mail.modifyLabels.mutationOptions());
 
   const generatePendingActionId = () =>
     `pending_${Date.now()}_${Math.random().toString(36).substring(2, 9)}`;
@@ -50,9 +51,10 @@ export function useOptimisticActions() {
             queryKey: trpc.mail.get.queryKey({ id }),
           }),
         ),
+        queryClient.refetchQueries({ queryKey: trpc.labels.list.queryKey() }),
       ]);
     },
-    [queryClient, trpc.mail.get],
+    [queryClient, trpc.mail.get, trpc.labels.list],
   );
 
   function createPendingAction({
@@ -376,6 +378,41 @@ export function useOptimisticActions() {
     });
   }
 
+  function optimisticToggleLabel(threadIds: string[], labelId: string, add: boolean) {
+    if (!threadIds.length || !labelId) return;
+
+    const optimisticId = addOptimisticAction({
+      type: 'LABEL',
+      threadIds,
+      labelIds: [labelId],
+      add,
+    });
+
+    createPendingAction({
+      type: 'LABEL',
+      threadIds,
+      params: { labelId, add },
+      optimisticId,
+      execute: async () => {
+        await modifyLabels({
+          threadId: threadIds,
+          addLabels: add ? [labelId] : [],
+          removeLabels: add ? [] : [labelId],
+        });
+
+        if (mail.bulkSelected.length > 0) {
+          setMail({ ...mail, bulkSelected: [] });
+        }
+      },
+      undo: () => {
+        removeOptimisticAction(optimisticId);
+      },
+      toastMessage: add
+        ? `Label added${threadIds.length > 1 ? ` to ${threadIds.length} threads` : ''}`
+        : `Label removed${threadIds.length > 1 ? ` from ${threadIds.length} threads` : ''}`,
+    });
+  }
+
   function undoLastAction() {
     if (!optimisticActionsManager.lastActionId) return;
 
@@ -405,6 +442,7 @@ export function useOptimisticActions() {
     optimisticMoveThreadsTo,
     optimisticDeleteThreads,
     optimisticToggleImportant,
+    optimisticToggleLabel,
     undoLastAction,
   };
 }