From b9b29e5a76ceade0a8e12222ef1ab40f840f0b6a Mon Sep 17 00:00:00 2001
From: ephraimduncan <ephraimduncan68@gmail.com>
Date: Mon, 20 Apr 2026 15:48:18 +0000
Subject: [PATCH] fix: bulk download partial failure, abort, and race-safe e2e

- onSuccess now reports successful envelope ids so the parent clears
  only those rows from selection instead of wiping all pages.
- Partial failures no longer auto-close the dialog; failed/unprocessed
  ids stay selected for retry.
- Cancel button turns into Stop while downloading and aborts the batch
  at the next envelope boundary.
- Replace dual waitForEvent('download') with a page.on collector +
  expect.poll so both downloads are captured reliably.
---
 .../envelopes-bulk-download-dialog.tsx        | 61 ++++++++++---------
 .../t.$teamUrl+/documents._index.tsx          | 10 ++-
 .../documents/bulk-document-actions.spec.ts   | 12 ++--
 3 files changed, 46 insertions(+), 37 deletions(-)

diff --git a/apps/remix/app/components/dialogs/envelopes-bulk-download-dialog.tsx b/apps/remix/app/components/dialogs/envelopes-bulk-download-dialog.tsx
index bd217ac03..44f8ad02b 100644
--- a/apps/remix/app/components/dialogs/envelopes-bulk-download-dialog.tsx
+++ b/apps/remix/app/components/dialogs/envelopes-bulk-download-dialog.tsx
@@ -1,4 +1,4 @@
-import { useEffect, useState } from 'react';
+import { useEffect, useRef, useState } from 'react';
 
 import { plural } from '@lingui/core/macro';
 import { Plural, useLingui } from '@lingui/react/macro';
@@ -37,7 +37,7 @@ export type EnvelopesBulkDownloadDialogProps = {
   envelopes: EnvelopeBulkDownloadItem[];
   open: boolean;
   onOpenChange: (open: boolean) => void;
-  onSuccess?: () => void;
+  onSuccess?: (successfulEnvelopeIds: string[]) => void;
 } & Omit<DialogPrimitive.DialogProps, 'children'>;
 
 type DownloadVersion = 'signed' | 'original';
@@ -56,6 +56,8 @@ export const EnvelopesBulkDownloadDialog = ({
   const [progress, setProgress] = useState(0);
   const [isDownloading, setIsDownloading] = useState(false);
 
+  const abortRef = useRef(false);
+
   const trpcUtils = trpc.useUtils();
 
   useEffect(() => {
@@ -89,15 +91,18 @@ export const EnvelopesBulkDownloadDialog = ({
       return;
     }
 
+    abortRef.current = false;
     setIsDownloading(true);
     setProgress(0);
 
-    let successfulDownloads = 0;
+    const successfulEnvelopeIds: string[] = [];
     let failedDownloads = 0;
 
     try {
-      for (let index = 0; index < envelopes.length; index++) {
-        const envelope = envelopes[index];
+      for (const envelope of envelopes) {
+        if (abortRef.current) {
+          break;
+        }
 
         try {
           const downloadVersion: DownloadVersion =
@@ -121,29 +126,28 @@ export const EnvelopesBulkDownloadDialog = ({
             });
           }
 
-          successfulDownloads++;
+          successfulEnvelopeIds.push(envelope.id);
         } catch (error) {
           console.error(error);
           failedDownloads++;
         }
 
-        setProgress(index + 1);
+        setProgress((p) => p + 1);
       }
 
-      if (successfulDownloads === 0) {
+      if (successfulEnvelopeIds.length === 0) {
         toast({
           title: t`Error`,
           description: t`An error occurred while downloading the documents.`,
           variant: 'destructive',
         });
-
         return;
       }
 
       if (failedDownloads > 0) {
         toast({
           title: t`Documents partially downloaded`,
-          description: t`${plural(successfulDownloads, {
+          description: t`${plural(successfulEnvelopeIds.length, {
             one: '# document downloaded.',
             other: '# documents downloaded.',
           })} ${plural(failedDownloads, {
@@ -152,26 +156,20 @@ export const EnvelopesBulkDownloadDialog = ({
           })}`,
           variant: 'destructive',
         });
-      } else {
-        toast({
-          title: t`Documents downloaded`,
-          description: plural(successfulDownloads, {
-            one: '# document has been downloaded.',
-            other: '# documents have been downloaded.',
-          }),
-        });
+        onSuccess?.(successfulEnvelopeIds);
+        return;
       }
 
-      onSuccess?.();
-      onOpenChange(false);
-    } catch (error) {
-      console.error(error);
-
       toast({
-        title: t`Error`,
-        description: t`An error occurred while downloading the documents.`,
-        variant: 'destructive',
+        title: t`Documents downloaded`,
+        description: plural(successfulEnvelopeIds.length, {
+          one: '# document has been downloaded.',
+          other: '# documents have been downloaded.',
+        }),
       });
+
+      onSuccess?.(successfulEnvelopeIds);
+      onOpenChange(false);
     } finally {
       setIsDownloading(false);
     }
@@ -264,10 +262,15 @@ export const EnvelopesBulkDownloadDialog = ({
             <Button
               type="button"
               variant="secondary"
-              onClick={() => onOpenChange(false)}
-              disabled={isDownloading}
+              onClick={() => {
+                if (isDownloading) {
+                  abortRef.current = true;
+                } else {
+                  onOpenChange(false);
+                }
+              }}
             >
-              <Trans>Cancel</Trans>
+              {isDownloading ? <Trans>Stop</Trans> : <Trans>Cancel</Trans>}
             </Button>
 
             <Button
diff --git a/apps/remix/app/routes/_authenticated+/t.$teamUrl+/documents._index.tsx b/apps/remix/app/routes/_authenticated+/t.$teamUrl+/documents._index.tsx
index 3040bce2a..827f95b8c 100644
--- a/apps/remix/app/routes/_authenticated+/t.$teamUrl+/documents._index.tsx
+++ b/apps/remix/app/routes/_authenticated+/t.$teamUrl+/documents._index.tsx
@@ -267,7 +267,15 @@ export default function DocumentsPage() {
           envelopes={selectedEnvelopesForDownload}
           open={isBulkDownloadDialogOpen}
           onOpenChange={setIsBulkDownloadDialogOpen}
-          onSuccess={() => setRowSelection({})}
+          onSuccess={(successfulEnvelopeIds) => {
+            setRowSelection((prev) => {
+              const next = { ...prev };
+              for (const id of successfulEnvelopeIds) {
+                delete next[id];
+              }
+              return next;
+            });
+          }}
         />
 
         <EnvelopesBulkMoveDialog
diff --git a/packages/app-tests/e2e/documents/bulk-document-actions.spec.ts b/packages/app-tests/e2e/documents/bulk-document-actions.spec.ts
index 878fc00df..1398d21c5 100644
--- a/packages/app-tests/e2e/documents/bulk-document-actions.spec.ts
+++ b/packages/app-tests/e2e/documents/bulk-document-actions.spec.ts
@@ -1,3 +1,4 @@
+import type { Download } from '@playwright/test';
 import { expect, test } from '@playwright/test';
 
 import { seedDraftDocument } from '@documenso/prisma/seed/documents';
@@ -133,17 +134,14 @@ test('[BULK_ACTIONS]: can bulk download multiple documents', async ({ page }) =>
   await expect(dialog.getByText('Bulk Test Doc 2')).toBeVisible();
   await expect(dialog.getByText('Draft').first()).toBeVisible();
 
-  const firstDownloadPromise = page.waitForEvent('download');
-  const secondDownloadPromise = page.waitForEvent('download');
+  const downloads: Download[] = [];
+  page.on('download', (d) => downloads.push(d));
 
   await dialog.getByRole('button', { name: 'Download' }).click();
 
-  const [firstDownload, secondDownload] = await Promise.all([
-    firstDownloadPromise,
-    secondDownloadPromise,
-  ]);
+  await expect.poll(() => downloads.length, { timeout: 10_000 }).toBe(2);
 
-  expect([firstDownload.suggestedFilename(), secondDownload.suggestedFilename()]).toEqual(
+  expect(downloads.map((d) => d.suggestedFilename())).toEqual(
     expect.arrayContaining(['Bulk Test Doc 1.pdf', 'Bulk Test Doc 2.pdf']),
   );