Queue fixes

apps/web/app/(app)/automation/BulkRunRules.tsx

@@ -37,7 +37,7 @@ export function BulkRunRules() {
   const [startDate, setStartDate] = useState<Date | undefined>();
   const [endDate, setEndDate] = useState<Date | undefined>();
 
-  const abortRef = useRef<() => void>();
+  const abortRef = useRef<() => void>(undefined);
 
   return (
     <div>

apps/web/app/(app)/bulk-unsubscribe/ArchiveProgress.tsx

@@ -1,25 +1,23 @@
-import { useEffect } from "react";
+import { memo, useEffect } from "react";
 import { AnimatePresence, motion } from "framer-motion";
 import { useAtomValue } from "jotai";
 import { ProgressBar } from "@tremor/react";
-import { queueAtoms, resetTotalThreads } from "@/store/archive-queue";
+import { queueAtom, resetTotalThreads } from "@/store/archive-queue";
 import { cn } from "@/utils";
 
-export const ArchiveProgress = () => {
-  const { totalThreads, activeThreadIds } = useAtomValue(queueAtoms.archive);
+export const ArchiveProgress = memo(() => {
+  const { totalThreads, activeThreads } = useAtomValue(queueAtom);
 
-  // Make sure activeThreadIds is an object as this was causing an error.
-  const threadsRemaining = Object.values(activeThreadIds || {}).filter(
-    Boolean,
-  ).length;
-  const totalArchived = totalThreads - threadsRemaining;
-  const progress = (totalArchived / totalThreads) * 100;
+  // Make sure activeThreads is an object as this was causing an error.
+  const threadsRemaining = Object.values(activeThreads || {}).length;
+  const totalProcessed = totalThreads - threadsRemaining;
+  const progress = (totalProcessed / totalThreads) * 100;
   const isCompleted = progress === 100;
 
   useEffect(() => {
     if (isCompleted) {
       setTimeout(() => {
-        resetTotalThreads("archive");
+        resetTotalThreads();
       }, 5_000);
     }
   }, [isCompleted]);
@@ -50,11 +48,11 @@ export const ArchiveProgress = () => {
               {isCompleted ? "Archiving complete!" : "Archiving emails..."}
             </span>
             <span>
-              {totalArchived} of {totalThreads} emails archived
+              {totalProcessed} of {totalThreads} emails archived
             </span>
           </p>
         </motion.div>
       </AnimatePresence>
     </div>
   );
-};
+});

apps/web/app/(app)/bulk-unsubscribe/BulkUnsubscribeSection.tsx

@@ -40,6 +40,7 @@ import { SearchBar } from "@/app/(app)/bulk-unsubscribe/SearchBar";
 import { useToggleSelect } from "@/hooks/useToggleSelect";
 import { BulkActions } from "@/app/(app)/bulk-unsubscribe/BulkActions";
 import { ArchiveProgress } from "@/app/(app)/bulk-unsubscribe/ArchiveProgress";
+import { ClientOnly } from "@/components/ClientOnly";
 
 type Newsletter = NewsletterStatsResponse["newsletters"][number];
 
@@ -227,7 +228,9 @@ export function BulkUnsubscribeSection({
           </div>
         </div>
 
-        <ArchiveProgress />
+        <ClientOnly>
+          <ArchiveProgress />
+        </ClientOnly>
 
         {isStatsLoading && !isLoading && !data?.newsletters.length ? (
           <div className="p-4">

apps/web/app/(app)/bulk-unsubscribe/hooks.ts

@@ -143,7 +143,7 @@ async function autoArchive(
   await mutate();
   await decrementUnsubscribeCreditAction();
   await refetchPremium();
-  await archiveAllSenderEmails(name, () => {});
+  await archiveAllSenderEmails(name, () => {}, labelId);
 }
 
 export function useAutoArchive<T extends Row>({
@@ -182,6 +182,8 @@ export function useAutoArchive<T extends Row>({
       status: null,
     });
     await mutate();
+
+    setAutoArchiveLoading(false);
   }, [item.name, mutate, posthog, refetchPremium]);
 
   const onAutoArchiveAndLabel = useCallback(

apps/web/package.json

@@ -95,6 +95,7 @@
     "novel": "0.3.1",
     "openai": "^4.67.3",
     "p-queue": "^8.0.1",
+    "p-retry": "^6.2.0",
     "posthog-js": "^1.167.0",
     "posthog-node": "^4.2.0",
     "prettier": "^3.3.3",

apps/web/store/QueueInitializer.tsx

@@ -1,28 +1,22 @@
 "use client";
 
-import { useAtom } from "jotai";
+import { useAtomValue } from "jotai";
 import { useEffect } from "react";
-import { processQueue, queueAtoms } from "@/store/archive-queue";
+import { processQueue, queueAtom } from "@/store/archive-queue";
 
-function useInitializeQueues() {
-  const [archiveQueue] = useAtom(queueAtoms.archive);
-  const [deleteQueue] = useAtom(queueAtoms.delete);
-  const [markReadQueue] = useAtom(queueAtoms.markRead);
-
-  useEffect(() => {
-    const threadIds = Object.keys(archiveQueue.activeThreadIds || {});
-    if (threadIds.length) processQueue("archive", threadIds);
-  }, [archiveQueue]);
+let isInitialized = false;
 
-  useEffect(() => {
-    const threadIds = Object.keys(deleteQueue.activeThreadIds || {});
-    if (threadIds.length) processQueue("delete", threadIds);
-  }, [deleteQueue]);
+function useInitializeQueues() {
+  const queueState = useAtomValue(queueAtom);
 
   useEffect(() => {
-    const threadIds = Object.keys(markReadQueue.activeThreadIds || {});
-    if (threadIds.length) processQueue("markRead", threadIds);
-  }, [markReadQueue]);
+    if (!isInitialized) {
+      isInitialized = true;
+      if (queueState.activeThreads) {
+        processQueue({ threads: queueState.activeThreads });
+      }
+    }
+  }, [queueState.activeThreads]);
 }
 
 export function QueueInitializer() {

apps/web/store/archive-queue.ts

@@ -1,103 +1,146 @@
 import { atomWithStorage, createJSONStorage } from "jotai/utils";
+import pRetry from "p-retry";
 import { jotaiStore } from "@/store";
 import { emailActionQueue } from "@/utils/queue/email-action-queue";
 import {
   archiveThreadAction,
   trashThreadAction,
   markReadThreadAction,
 } from "@/utils/actions/mail";
+import { isActionError, ServerActionResponse } from "@/utils/error";
+import { exponentialBackoff, sleep } from "@/utils/sleep";
 
-type QueueType = "archive" | "delete" | "markRead";
+type ActionType = "archive" | "delete" | "markRead";
+
+type QueueItem = {
+  threadId: string;
+  actionType: ActionType;
+  labelId?: string;
+};
 
 type QueueState = {
-  activeThreadIds: Record<string, boolean>;
+  activeThreads: Record<`${ActionType}-${string}`, QueueItem>;
   totalThreads: number;
 };
 
-function getInitialState(): QueueState {
-  return { activeThreadIds: {}, totalThreads: 0 };
-}
-
-// some users were somehow getting null for activeThreadIds, this should fix it
+// some users were somehow getting null for activeThreads, this should fix it
 const createStorage = () => {
+  if (typeof window === "undefined") return;
   const storage = createJSONStorage<QueueState>(() => localStorage);
   return {
     ...storage,
     getItem: (key: string, initialValue: QueueState) => {
       const storedValue = storage.getItem(key, initialValue);
       return {
-        activeThreadIds: storedValue.activeThreadIds || {},
+        activeThreads: storedValue.activeThreads || {},
         totalThreads: storedValue.totalThreads || 0,
       };
     },
   };
 };
 
-// Create atoms with localStorage persistence for each queue type
-export const queueAtoms = {
-  archive: atomWithStorage("archiveQueue", getInitialState(), createStorage()),
-  delete: atomWithStorage("deleteQueue", getInitialState(), createStorage()),
-  markRead: atomWithStorage(
-    "markReadQueue",
-    getInitialState(),
-    createStorage(),
-  ),
-};
+// Create atoms with localStorage persistence
+export const queueAtom = atomWithStorage(
+  "gmailActionQueue",
+  { activeThreads: {}, totalThreads: 0 },
+  createStorage(),
+  { getOnInit: true },
+);
 
-type ActionFunction = (threadId: string, ...args: any[]) => Promise<any>;
+type ActionFunction = (
+  threadId: string,
+  labelId?: string,
+) => Promise<ServerActionResponse<{}>>;
 
-const actionMap: Record<QueueType, ActionFunction> = {
-  archive: archiveThreadAction,
+const actionMap: Record<ActionType, ActionFunction> = {
+  archive: (threadId: string, labelId?: string) =>
+    archiveThreadAction(threadId, labelId),
   delete: trashThreadAction,
   markRead: (threadId: string) => markReadThreadAction(threadId, true),
 };
 
-export const addThreadsToQueue = (
-  queueType: QueueType,
-  threadIds: string[],
-  refetch?: () => void,
-) => {
-  const queueAtom = queueAtoms[queueType];
+export const addThreadsToQueue = ({
+  actionType,
+  threadIds,
+  labelId,
+  refetch,
+}: {
+  actionType: ActionType;
+  threadIds: string[];
+  labelId?: string;
+  refetch?: () => void;
+}) => {
+  const threads = Object.fromEntries(
+    threadIds.map((threadId) => [
+      `${actionType}-${threadId}`,
+      { threadId, actionType, labelId },
+    ]),
+  );
 
   jotaiStore.set(queueAtom, (prev) => ({
-    activeThreadIds: {
-      ...prev.activeThreadIds,
-      ...Object.fromEntries(threadIds.map((id) => [id, true])),
+    activeThreads: {
+      ...prev.activeThreads,
+      ...threads,
     },
-    totalThreads: prev.totalThreads + threadIds.length,
+    totalThreads: prev.totalThreads + Object.keys(threads).length,
   }));
 
-  processQueue(queueType, threadIds, refetch);
+  processQueue({ threads, refetch });
 };
 
-export function processQueue(
-  queueType: QueueType,
-  threadIds: string[],
-  refetch?: () => void,
-) {
-  const queueAtom = queueAtoms[queueType];
-  const action = actionMap[queueType];
-
+export function processQueue({
+  threads,
+  refetch,
+}: {
+  threads: Record<string, QueueItem>;
+  refetch?: () => void;
+}) {
   emailActionQueue.addAll(
-    threadIds.map((threadId) => async () => {
-      await action(threadId);
-
-      // remove completed thread from activeThreadIds
-      jotaiStore.set(queueAtom, (prev) => {
-        const { [threadId]: _, ...remainingThreads } = prev.activeThreadIds;
-        return {
-          ...prev,
-          activeThreadIds: remainingThreads,
-        };
-      });
-
-      refetch?.();
-    }),
+    Object.entries(threads).map(
+      ([_key, { threadId, actionType, labelId }]) =>
+        async () => {
+          await pRetry(
+            async (attemptCount) => {
+              console.log(
+                `Queue: ${actionType}. Processing ${threadId}` +
+                  (attemptCount > 1 ? ` (attempt ${attemptCount})` : ""),
+              );
+
+              const result = await actionMap[actionType](threadId, labelId);
+
+              // when Gmail API returns a rate limit error, throw an error so it can be retried
+              if (isActionError(result)) {
+                await sleep(exponentialBackoff(attemptCount, 1_000));
+                throw new Error(result.error);
+              }
+              refetch?.();
+            },
+            { retries: 3 },
+          );
+
+          // remove completed thread from activeThreads
+          jotaiStore.set(queueAtom, (prev) => {
+            const remainingThreads = Object.fromEntries(
+              Object.entries(prev.activeThreads).filter(
+                ([_key, value]) =>
+                  !(
+                    value.threadId === threadId &&
+                    value.actionType === actionType
+                  ),
+              ),
+            );
+
+            return {
+              ...prev,
+              activeThreads: remainingThreads,
+            };
+          });
+        },
+    ),
   );
 }
 
-export const resetTotalThreads = (queueType: QueueType) => {
-  const queueAtom = queueAtoms[queueType];
+export const resetTotalThreads = () => {
   jotaiStore.set(queueAtom, (prev) => ({
     ...prev,
     totalThreads: 0,

apps/web/utils/actions/mail.ts

@@ -25,7 +25,7 @@ const isStatusOk = (status: number) => status >= 200 && status < 300;
 
 export const archiveThreadAction = withActionInstrumentation(
   "archiveThread",
-  async (threadId: string) => {
+  async (threadId: string, labelId?: string) => {
     const { gmail, user, error } = await getSessionAndGmailClient();
     if (error) return { error };
     if (!gmail) return { error: "Could not load Gmail" };
@@ -35,6 +35,7 @@ export const archiveThreadAction = withActionInstrumentation(
       threadId,
       ownerEmail: user.email,
       actionSource: "user",
+      labelId,
     });
 
     if (!isStatusOk(res.status)) return { error: "Failed to archive thread" };

apps/web/utils/gmail/label.ts

@@ -28,19 +28,25 @@ export async function labelThread(options: {
   });
 }
 
-export async function archiveThread(options: {
+export async function archiveThread({
+  gmail,
+  threadId,
+  ownerEmail,
+  actionSource,
+  labelId,
+}: {
   gmail: gmail_v1.Gmail;
   threadId: string;
   ownerEmail: string;
   actionSource: TinybirdEmailAction["actionSource"];
+  labelId?: string;
 }) {
-  const { gmail, threadId, ownerEmail, actionSource } = options;
-
   const archivePromise = gmail.users.threads.modify({
     userId: "me",
     id: threadId,
     requestBody: {
       removeLabelIds: [INBOX_LABEL_ID],
+      ...(labelId ? { addLabelIds: [labelId] } : {}),
     },
   });
 

apps/web/utils/queue/email-action-queue.ts

@@ -3,4 +3,4 @@
 import PQueue from "p-queue";
 
 // Avoid overwhelming Gmail API
-export const emailActionQueue = new PQueue({ concurrency: 3 });
+export const emailActionQueue = new PQueue({ concurrency: 1 });