From 15b40aa17ef674f5760f5caf715d573281abcb35 Mon Sep 17 00:00:00 2001 From: Mariano Fuentes Date: Wed, 27 May 2026 13:45:36 -0400 Subject: [PATCH 1/4] fix(evidence-export): stream runs through PDF/JSON generation to prevent OOM The previous OOM fix loaded automations one at a time but still accumulated all runs for a single automation in memory. For orgs with large cloud security check histories, a single automation's runs could exceed the 6GB heap limit. Now uses async generators to stream run batches (50 at a time) through PDF and JSON generation. Peak memory is bounded by one batch of runs + the jsPDF document, regardless of total automation size. - evidence-data-loader: add streamAutomationRuns async generator - evidence-pdf-generator: extract renderRunToPDF, add generateAutomationPDFFromStream - evidence-json-builder: add buildAutomationJsonStream using Readable.from() - evidence-export.service: wire streaming into ZIP export path Co-Authored-By: Claude Opus 4.6 (1M context) --- .../evidence-export/evidence-data-loader.ts | 55 ++-- .../evidence-export.service.spec.ts | 12 + .../evidence-export.service.ts | 82 +++-- .../evidence-export/evidence-json-builder.ts | 53 ++++ .../evidence-export/evidence-pdf-generator.ts | 293 +++++++++--------- 5 files changed, 298 insertions(+), 197 deletions(-) diff --git a/apps/api/src/tasks/evidence-export/evidence-data-loader.ts b/apps/api/src/tasks/evidence-export/evidence-data-loader.ts index 6d0da40c56..45fa235eac 100644 --- a/apps/api/src/tasks/evidence-export/evidence-data-loader.ts +++ b/apps/api/src/tasks/evidence-export/evidence-data-loader.ts @@ -197,22 +197,38 @@ export async function loadFullAutomation({ taskId: string; header: NormalizedAutomation; }): Promise { + const runs: NormalizedEvidenceRun[] = []; + for await (const batch of streamAutomationRuns({ taskId, header })) { + runs.push(...batch); + } + return { ...header, runs }; +} + +// Yields runs in batches so callers can process incrementally without +// accumulating all runs in memory. Each batch is GC-eligible after processing. +export async function* streamAutomationRuns({ + taskId, + header, +}: { + taskId: string; + header: NormalizedAutomation; +}): AsyncGenerator { if (header.type === 'app_automation' && header.checkId) { - return loadAppAutomationRuns(taskId, header); + yield* streamAppRuns(taskId, header.checkId); + } else { + yield* streamCustomRuns(taskId, header.id); } - return loadCustomAutomationRuns(taskId, header); } -async function loadAppAutomationRuns( +async function* streamAppRuns( taskId: string, - header: NormalizedAutomation, -): Promise { - const runs: NormalizedEvidenceRun[] = []; + checkId: string, +): AsyncGenerator { let cursor: { id: string } | undefined; for (;;) { const batch = await db.integrationCheckRun.findMany({ - where: { taskId, checkId: header.checkId }, + where: { taskId, checkId }, include: { results: true, connection: { include: { provider: true } }, @@ -224,28 +240,25 @@ async function loadAppAutomationRuns( if (batch.length === 0) break; - for (const run of batch) { - runs.push(normalizeAppAutomationRun(toAppAutomationRun(run))); - } + yield batch.map((run) => + normalizeAppAutomationRun(toAppAutomationRun(run)), + ); if (batch.length < RUN_BATCH_SIZE) break; cursor = { id: batch[batch.length - 1].id }; } - - return { ...header, runs }; } -async function loadCustomAutomationRuns( +async function* streamCustomRuns( taskId: string, - header: NormalizedAutomation, -): Promise { - const runs: NormalizedEvidenceRun[] = []; + automationId: string, +): AsyncGenerator { let cursor: { id: string } | undefined; for (;;) { const batch = await db.evidenceAutomationRun.findMany({ where: { - evidenceAutomation: { id: header.id, taskId }, + evidenceAutomation: { id: automationId, taskId }, version: { not: null }, }, include: { @@ -258,15 +271,13 @@ async function loadCustomAutomationRuns( if (batch.length === 0) break; - for (const run of batch) { - runs.push(normalizeCustomAutomationRun(toCustomAutomationRun(run))); - } + yield batch.map((run) => + normalizeCustomAutomationRun(toCustomAutomationRun(run)), + ); if (batch.length < RUN_BATCH_SIZE) break; cursor = { id: batch[batch.length - 1].id }; } - - return { ...header, runs }; } // Prisma result → normalizer interface mappers (single source of truth). diff --git a/apps/api/src/tasks/evidence-export/evidence-export.service.spec.ts b/apps/api/src/tasks/evidence-export/evidence-export.service.spec.ts index ab6de156d8..8b47882611 100644 --- a/apps/api/src/tasks/evidence-export/evidence-export.service.spec.ts +++ b/apps/api/src/tasks/evidence-export/evidence-export.service.spec.ts @@ -87,6 +87,18 @@ jest.mock('@db', () => ({ jest.mock('./evidence-pdf-generator', () => ({ generateTaskSummaryPDF: jest.fn(() => Buffer.from('SUMMARY-PDF')), generateAutomationPDF: jest.fn(() => Buffer.from('AUTOMATION-PDF')), + generateAutomationPDFFromStream: jest.fn( + async ( + _header: unknown, + _context: unknown, + runBatches: AsyncIterable, + ) => { + for await (const _batch of runBatches) { + /* drain so underlying DB queries execute */ + } + return Buffer.from('AUTOMATION-PDF'); + }, + ), sanitizeFilename: (name: string) => name .toLowerCase() diff --git a/apps/api/src/tasks/evidence-export/evidence-export.service.ts b/apps/api/src/tasks/evidence-export/evidence-export.service.ts index cc77678d57..9781f06296 100644 --- a/apps/api/src/tasks/evidence-export/evidence-export.service.ts +++ b/apps/api/src/tasks/evidence-export/evidence-export.service.ts @@ -11,10 +11,11 @@ import type { } from './evidence-export.types'; import { generateAutomationPDF, + generateAutomationPDFFromStream, generateTaskSummaryPDF, sanitizeFilename, } from './evidence-pdf-generator'; -import { buildAutomationJson } from './evidence-json-builder'; +import { buildAutomationJson, buildAutomationJsonStream } from './evidence-json-builder'; import { appendAttachmentToArchive, createFilenameTracker, @@ -24,6 +25,7 @@ import { import { getAutomationHeaders, loadFullAutomation, + streamAutomationRuns, findTasksWithEvidence, } from './evidence-data-loader'; @@ -221,7 +223,8 @@ export class EvidenceExportService { }); } - // Loads each automation's runs individually so peak memory ≈ one automation, not all combined. + // Streams each automation's runs through PDF/JSON generation so peak memory + // is bounded by one batch of runs (~50) instead of the full automation. private async appendTaskContents(params: { archive: Archiver; headers: TaskEvidenceSummary; @@ -257,15 +260,10 @@ export class EvidenceExportService { } for (const automationHeader of headers.automations) { - const automation = await loadFullAutomation({ - taskId: headers.taskId, - header: automationHeader, - }); - - this.appendAutomationToArchive({ + await this.appendAutomationStreaming({ archive, headers, - automation, + automationHeader, folderName, options, perAutomationSubfolders, @@ -273,50 +271,70 @@ export class EvidenceExportService { } } - private appendAutomationToArchive(params: { + private async appendAutomationStreaming(params: { archive: Archiver; headers: TaskEvidenceSummary; - automation: NormalizedAutomation; + automationHeader: NormalizedAutomation; folderName: string; options: { includeRawJson?: boolean }; perAutomationSubfolders: boolean; - }): void { + }): Promise { const { archive, headers, - automation, + automationHeader, folderName, options, perAutomationSubfolders, } = params; const typePrefix = - automation.type === 'app_automation' ? 'app' : 'custom'; - const automationName = sanitizeFilename(automation.name); - const idSuffix = automation.id.slice(-8); - - const pdfBuffer = generateAutomationPDF(automation, { - organizationName: headers.organizationName, - taskTitle: headers.taskTitle, - }); - + automationHeader.type === 'app_automation' ? 'app' : 'custom'; + const automationName = sanitizeFilename(automationHeader.name); + const idSuffix = automationHeader.id.slice(-8); const basePath = perAutomationSubfolders ? `${folderName}/${typePrefix}-${automationName}-${idSuffix}` : folderName; - const pdfName = perAutomationSubfolders - ? `${basePath}/evidence.pdf` - : `${basePath}/${typePrefix}-${automationName}-${idSuffix}.pdf`; + const filePrefix = perAutomationSubfolders + ? `${basePath}/evidence` + : `${basePath}/${typePrefix}-${automationName}-${idSuffix}`; - archive.append(pdfBuffer, { name: pdfName }); + const context = { + organizationName: headers.organizationName, + taskTitle: headers.taskTitle, + }; if (options.includeRawJson) { - const jsonName = perAutomationSubfolders - ? `${basePath}/evidence.json` - : `${basePath}/${typePrefix}-${automationName}-${idSuffix}.json`; - archive.append( - Buffer.from(buildAutomationJson(headers, automation), 'utf-8'), - { name: jsonName }, + // Two independent DB cursors so neither PDF nor JSON buffers the full run set. + const pdfBuffer = await generateAutomationPDFFromStream( + automationHeader, + context, + streamAutomationRuns({ + taskId: headers.taskId, + header: automationHeader, + }), + ); + archive.append(pdfBuffer, { name: `${filePrefix}.pdf` }); + + const jsonStream = buildAutomationJsonStream({ + summary: headers, + header: automationHeader, + runBatches: streamAutomationRuns({ + taskId: headers.taskId, + header: automationHeader, + }), + }); + archive.append(jsonStream, { name: `${filePrefix}.json` }); + } else { + const pdfBuffer = await generateAutomationPDFFromStream( + automationHeader, + context, + streamAutomationRuns({ + taskId: headers.taskId, + header: automationHeader, + }), ); + archive.append(pdfBuffer, { name: `${filePrefix}.pdf` }); } } diff --git a/apps/api/src/tasks/evidence-export/evidence-json-builder.ts b/apps/api/src/tasks/evidence-export/evidence-json-builder.ts index 1f3c317d5f..e0d2914abf 100644 --- a/apps/api/src/tasks/evidence-export/evidence-json-builder.ts +++ b/apps/api/src/tasks/evidence-export/evidence-json-builder.ts @@ -1,8 +1,10 @@ +import { Readable } from 'node:stream'; import { configure as configureStringify } from 'safe-stable-stringify'; import { redactSensitiveData } from './evidence-redaction'; import type { TaskEvidenceSummary, NormalizedAutomation, + NormalizedEvidenceRun, } from './evidence-export.types'; const safeStringify = configureStringify({ @@ -11,6 +13,30 @@ const safeStringify = configureStringify({ deterministic: false, }); +function stringifyAutomationMeta(header: NormalizedAutomation): string { + return ( + safeStringify( + redactSensitiveData({ + id: header.id, + name: header.name, + type: header.type, + integrationName: header.integrationName, + totalRuns: header.totalRuns, + successfulRuns: header.successfulRuns, + failedRuns: header.failedRuns, + latestRunAt: header.latestRunAt, + }), + null, + 2, + ) ?? '{}' + ); +} + +function stringifyRun(run: NormalizedEvidenceRun): string { + const { type, automationName, automationId, ...rest } = run; + return safeStringify(redactSensitiveData(rest), null, 2) ?? '{}'; +} + export function buildAutomationJson( summary: TaskEvidenceSummary, automation: NormalizedAutomation, @@ -38,3 +64,30 @@ export function buildAutomationJson( ) ?? '{}' ); } + +export function buildAutomationJsonStream({ + summary, + header, + runBatches, +}: { + summary: TaskEvidenceSummary; + header: NormalizedAutomation; + runBatches: AsyncIterable; +}): Readable { + async function* chunks(): AsyncGenerator { + yield `{\n "automation": ${stringifyAutomationMeta(header)},\n "runs": [\n`; + + let first = true; + for await (const batch of runBatches) { + for (const run of batch) { + if (!first) yield ',\n'; + first = false; + yield ` ${stringifyRun(run)}`; + } + } + + yield `\n ],\n "exportedAt": ${JSON.stringify(summary.exportedAt.toISOString())}\n}\n`; + } + + return Readable.from(chunks(), { encoding: 'utf-8' }); +} diff --git a/apps/api/src/tasks/evidence-export/evidence-pdf-generator.ts b/apps/api/src/tasks/evidence-export/evidence-pdf-generator.ts index cde5bb2122..4196ed9e29 100644 --- a/apps/api/src/tasks/evidence-export/evidence-pdf-generator.ts +++ b/apps/api/src/tasks/evidence-export/evidence-pdf-generator.ts @@ -10,6 +10,7 @@ import stringify from 'safe-stable-stringify'; import { redactSensitiveData } from './evidence-redaction'; import type { NormalizedAutomation, + NormalizedEvidenceRun, TaskEvidenceSummary, } from './evidence-export.types'; @@ -233,18 +234,8 @@ function addPageNumbers(config: PDFConfig): void { } } -/** - * Generate PDF for a single automation - */ -export function generateAutomationPDF( - automation: NormalizedAutomation, - context: { - organizationName: string; - taskTitle: string; - }, -): Buffer { - const doc = new jsPDF(); - const config: PDFConfig = { +function createPDFConfig(doc: jsPDF): PDFConfig { + return { doc, pageWidth: doc.internal.pageSize.getWidth(), pageHeight: doc.internal.pageSize.getHeight(), @@ -254,25 +245,27 @@ export function generateAutomationPDF( defaultFontSize: 10, yPosition: 20, }; +} - // Header +function renderAutomationHeader( + config: PDFConfig, + header: NormalizedAutomation, + context: { organizationName: string; taskTitle: string }, +): void { addText(config, context.organizationName, { fontSize: 14, bold: true }); config.yPosition += config.lineHeight * 0.5; addText(config, `Task: ${context.taskTitle}`, { fontSize: 11 }); config.yPosition += config.lineHeight; - // Automation title const typeLabel = - automation.type === 'app_automation' - ? 'App Automation' - : 'Custom Automation'; - addText(config, `${typeLabel}: ${automation.name}`, { + header.type === 'app_automation' ? 'App Automation' : 'Custom Automation'; + addText(config, `${typeLabel}: ${header.name}`, { fontSize: 13, bold: true, }); - if (automation.integrationName) { - addText(config, `Integration: ${automation.integrationName}`, { + if (header.integrationName) { + addText(config, `Integration: ${header.integrationName}`, { fontSize: 10, color: [100, 100, 100], }); @@ -280,7 +273,6 @@ export function generateAutomationPDF( config.yPosition += config.lineHeight; - // Export timestamp addText(config, `Exported: ${format(new Date(), 'PPpp')}`, { fontSize: 9, color: [128, 128, 128], @@ -288,157 +280,172 @@ export function generateAutomationPDF( addSeparator(config); - // Summary section addSectionHeader(config, 'Summary'); - addText(config, `Total Runs: ${automation.totalRuns}`); - addText(config, `Successful: ${automation.successfulRuns}`); - addText(config, `Failed: ${automation.failedRuns}`); + addText(config, `Total Runs: ${header.totalRuns}`); + addText(config, `Successful: ${header.successfulRuns}`); + addText(config, `Failed: ${header.failedRuns}`); - if (automation.latestRunAt) { - addText(config, `Latest Run: ${format(automation.latestRunAt, 'PPpp')}`); + if (header.latestRunAt) { + addText(config, `Latest Run: ${format(header.latestRunAt, 'PPpp')}`); } addSeparator(config); - - // Runs section addSectionHeader(config, 'Run History'); +} - for (const run of automation.runs) { - checkPageBreak(config, config.lineHeight * 8); - - // Run header - const statusColor = getStatusColor(run.status, run.failedCount > 0); - addText(config, `Run: ${format(run.createdAt, 'PPpp')}`, { - fontSize: 10, - bold: true, - }); - addText(config, `Status: ${run.status.toUpperCase()}`, { - fontSize: 9, - color: statusColor, - }); +function renderRunToPDF(config: PDFConfig, run: NormalizedEvidenceRun): void { + checkPageBreak(config, config.lineHeight * 8); - if (run.durationMs) { - addText(config, `Duration: ${run.durationMs}ms`, { fontSize: 9 }); - } + const statusColor = getStatusColor(run.status, run.failedCount > 0); + addText(config, `Run: ${format(run.createdAt, 'PPpp')}`, { + fontSize: 10, + bold: true, + }); + addText(config, `Status: ${run.status.toUpperCase()}`, { + fontSize: 9, + color: statusColor, + }); - // Metrics - if (run.type === 'app_automation') { - addText( - config, - `Checked: ${run.totalChecked} | Passed: ${run.passedCount} | Failed: ${run.failedCount}`, - { fontSize: 9 }, - ); - } else if (run.evaluationStatus) { - addText(config, `Evaluation: ${run.evaluationStatus.toUpperCase()}`, { - fontSize: 9, - color: run.evaluationStatus === 'pass' ? [0, 128, 0] : [200, 0, 0], - }); - if (run.evaluationReason) { - addText(config, `Reason: ${run.evaluationReason}`, { - fontSize: 9, - indent: 10, - }); - } - } + if (run.durationMs) { + addText(config, `Duration: ${run.durationMs}ms`, { fontSize: 9 }); + } - // Error - if (run.error) { - config.yPosition += config.lineHeight * 0.5; - addText(config, 'Error:', { + if (run.type === 'app_automation') { + addText( + config, + `Checked: ${run.totalChecked} | Passed: ${run.passedCount} | Failed: ${run.failedCount}`, + { fontSize: 9 }, + ); + } else if (run.evaluationStatus) { + addText(config, `Evaluation: ${run.evaluationStatus.toUpperCase()}`, { + fontSize: 9, + color: run.evaluationStatus === 'pass' ? [0, 128, 0] : [200, 0, 0], + }); + if (run.evaluationReason) { + addText(config, `Reason: ${run.evaluationReason}`, { fontSize: 9, - bold: true, - color: [200, 0, 0], - }); - addText(config, run.error, { - fontSize: 8, - color: [150, 0, 0], indent: 10, }); } + } - // Output (for custom automations) - if (run.output) { - config.yPosition += config.lineHeight * 0.5; - addText(config, 'Output:', { fontSize: 9, bold: true }); - const outputText = formatJsonForPDF(run.output); - addMonospaceText(config, outputText); - } + if (run.error) { + config.yPosition += config.lineHeight * 0.5; + addText(config, 'Error:', { + fontSize: 9, + bold: true, + color: [200, 0, 0], + }); + addText(config, run.error, { + fontSize: 8, + color: [150, 0, 0], + indent: 10, + }); + } - // Logs - if (run.logs) { - config.yPosition += config.lineHeight * 0.5; - addText(config, 'Logs:', { fontSize: 9, bold: true }); - const logsText = formatJsonForPDF(run.logs); - addMonospaceText(config, logsText); - } + if (run.output) { + config.yPosition += config.lineHeight * 0.5; + addText(config, 'Output:', { fontSize: 9, bold: true }); + addMonospaceText(config, formatJsonForPDF(run.output)); + } - // Results (for app automations) - if (run.results.length > 0) { - config.yPosition += config.lineHeight * 0.5; - addText(config, `Results (${run.results.length}):`, { + if (run.logs) { + config.yPosition += config.lineHeight * 0.5; + addText(config, 'Logs:', { fontSize: 9, bold: true }); + addMonospaceText(config, formatJsonForPDF(run.logs)); + } + + if (run.results.length > 0) { + config.yPosition += config.lineHeight * 0.5; + addText(config, `Results (${run.results.length}):`, { + fontSize: 9, + bold: true, + }); + + for (const result of run.results) { + checkPageBreak(config, config.lineHeight * 5); + + const resultIcon = result.passed ? '[PASS]' : '[FAIL]'; + const resultColor: [number, number, number] = result.passed + ? [0, 100, 0] + : [180, 0, 0]; + + addText(config, `${resultIcon} ${result.title}`, { fontSize: 9, bold: true, + color: resultColor, + indent: 10, }); - for (const result of run.results) { - checkPageBreak(config, config.lineHeight * 5); - - const resultIcon = result.passed ? '[PASS]' : '[FAIL]'; - const resultColor: [number, number, number] = result.passed - ? [0, 100, 0] - : [180, 0, 0]; + addText( + config, + `Resource: ${result.resourceType}/${result.resourceId}`, + { fontSize: 8, indent: 15 }, + ); - addText(config, `${resultIcon} ${result.title}`, { - fontSize: 9, - bold: true, - color: resultColor, - indent: 10, + if (result.description) { + addText(config, result.description, { fontSize: 8, indent: 15 }); + } + if (result.severity) { + addText(config, `Severity: ${result.severity}`, { + fontSize: 8, + indent: 15, }); - - addText( - config, - `Resource: ${result.resourceType}/${result.resourceId}`, - { - fontSize: 8, - indent: 15, - }, - ); - - if (result.description) { - addText(config, result.description, { fontSize: 8, indent: 15 }); - } - - if (result.severity) { - addText(config, `Severity: ${result.severity}`, { - fontSize: 8, - indent: 15, - }); - } - - if (result.remediation) { - addText(config, `Remediation: ${result.remediation}`, { - fontSize: 8, - indent: 15, - }); - } - - if (result.evidence) { - addText(config, 'Evidence:', { fontSize: 8, bold: true, indent: 15 }); - const evidenceText = formatJsonForPDF(result.evidence); - addMonospaceText(config, evidenceText, 20); - } - - config.yPosition += config.lineHeight * 0.5; } - } + if (result.remediation) { + addText(config, `Remediation: ${result.remediation}`, { + fontSize: 8, + indent: 15, + }); + } + if (result.evidence) { + addText(config, 'Evidence:', { fontSize: 8, bold: true, indent: 15 }); + addMonospaceText(config, formatJsonForPDF(result.evidence), 20); + } - config.yPosition += config.lineHeight; - addSeparator(config); + config.yPosition += config.lineHeight * 0.5; + } } + config.yPosition += config.lineHeight; + addSeparator(config); +} + +/** + * Generate PDF for a single automation (loads all runs into memory). + * Use generateAutomationPDFFromStream for large automations. + */ +export function generateAutomationPDF( + automation: NormalizedAutomation, + context: { organizationName: string; taskTitle: string }, +): Buffer { + const config = createPDFConfig(new jsPDF()); + renderAutomationHeader(config, automation, context); + for (const run of automation.runs) { + renderRunToPDF(config, run); + } addPageNumbers(config); + return Buffer.from(config.doc.output('arraybuffer')); +} - return Buffer.from(doc.output('arraybuffer')); +/** + * Build a PDF incrementally from an async stream of run batches. + * Peak memory = one batch of runs + the jsPDF document (lightweight page objects). + */ +export async function generateAutomationPDFFromStream( + header: NormalizedAutomation, + context: { organizationName: string; taskTitle: string }, + runBatches: AsyncIterable, +): Promise { + const config = createPDFConfig(new jsPDF()); + renderAutomationHeader(config, header, context); + for await (const batch of runBatches) { + for (const run of batch) { + renderRunToPDF(config, run); + } + } + addPageNumbers(config); + return Buffer.from(config.doc.output('arraybuffer')); } /** From dd5a2f059d2d1121548d3c0fdd783550818bbab1 Mon Sep 17 00:00:00 2001 From: Mariano Fuentes Date: Thu, 28 May 2026 11:26:30 -0400 Subject: [PATCH 2/4] feat(evidence-export): offload bulk export to Trigger.dev background task The auditor bulk evidence export previously ran in the API process, peaking at ~20% memory per request. Multiple concurrent exports could OOM the container. Now the heavy work (DB queries, PDF generation, ZIP creation) runs in a Trigger.dev background task with its own memory. The API endpoint triggers the task and returns a runId for progress tracking. - Add export-organization-evidence Trigger.dev task (S3 upload + presigned URL) - Change POST /v1/evidence-export/all to trigger background task - Frontend uses useRealtimeRun for progress + auto-download on completion - API process memory stays flat regardless of export size Co-Authored-By: Claude Opus 4.6 (1M context) --- .../evidence-export.controller.spec.ts | 52 ++-- .../evidence-export.controller.ts | 57 ++-- .../export-organization-evidence.ts | 275 ++++++++++++++++++ .../components/ExportEvidenceButton.tsx | 216 +++++++++++--- .../tasks/components/TasksPageClient.tsx | 54 +++- apps/app/src/lib/evidence-download.ts | 24 +- 6 files changed, 537 insertions(+), 141 deletions(-) create mode 100644 apps/api/src/trigger/evidence-export/export-organization-evidence.ts diff --git a/apps/api/src/tasks/evidence-export/evidence-export.controller.spec.ts b/apps/api/src/tasks/evidence-export/evidence-export.controller.spec.ts index e9f29c13e1..62015060f9 100644 --- a/apps/api/src/tasks/evidence-export/evidence-export.controller.spec.ts +++ b/apps/api/src/tasks/evidence-export/evidence-export.controller.spec.ts @@ -1,5 +1,10 @@ // Mocks must be declared before any SUT import so guards' transitive deps // (Prisma, better-auth) don't instantiate in Jest. +const mockTrigger = jest.fn(); +jest.mock('@trigger.dev/sdk', () => ({ + tasks: { trigger: mockTrigger }, +})); + jest.mock('@db', () => ({ ...jest.requireActual('@prisma/client'), db: {}, @@ -223,18 +228,15 @@ describe('EvidenceExportController', () => { describe('AuditorEvidenceExportController', () => { let controller: AuditorEvidenceExportController; - let service: jest.Mocked< - Pick - >; beforeEach(async () => { - service = { - streamOrganizationEvidenceZip: jest.fn(), - }; + mockTrigger.mockReset().mockResolvedValue({ + id: 'run_123', + publicAccessToken: 'tok_abc', + }); const moduleRef = await Test.createTestingModule({ controllers: [AuditorEvidenceExportController], - providers: [{ provide: EvidenceExportService, useValue: service }], }) .overrideGuard(HybridAuthGuard) .useValue({ canActivate: () => true }) @@ -245,34 +247,16 @@ describe('AuditorEvidenceExportController', () => { controller = moduleRef.get(AuditorEvidenceExportController); }); - it('pipes the org-wide archive to response with correct headers', async () => { - const archive = makeFakeArchive(); - service.streamOrganizationEvidenceZip.mockResolvedValue({ - archive: archive as unknown as import('archiver').Archiver, - filename: 'acme_all-evidence_2026-04-22.zip', - }); - const req = makeFakeRequest(); - const res = makeFakeResponse(); + it('triggers a background task and returns runId + token', async () => { + const result = await controller.exportAllEvidence('org_1', 'true'); - await controller.exportAllEvidence( - 'org_1', - 'true', - req as unknown as import('express').Request, - res as unknown as import('express').Response, - ); - - expect(service.streamOrganizationEvidenceZip).toHaveBeenCalledWith( - 'org_1', - { includeRawJson: true }, - ); - expect(res.setHeader).toHaveBeenCalledWith( - 'Content-Type', - 'application/zip', + expect(mockTrigger).toHaveBeenCalledWith( + 'export-organization-evidence', + { organizationId: 'org_1', includeJson: true }, ); - expect(res.setHeader).toHaveBeenCalledWith( - 'Content-Disposition', - `attachment; filename="acme_all-evidence_2026-04-22.zip"`, - ); - expect(archive.pipe).toHaveBeenCalledWith(res); + expect(result).toEqual({ + runId: 'run_123', + publicAccessToken: 'tok_abc', + }); }); }); diff --git a/apps/api/src/tasks/evidence-export/evidence-export.controller.ts b/apps/api/src/tasks/evidence-export/evidence-export.controller.ts index 3a5cd924d6..5eb3601dd3 100644 --- a/apps/api/src/tasks/evidence-export/evidence-export.controller.ts +++ b/apps/api/src/tasks/evidence-export/evidence-export.controller.ts @@ -1,6 +1,7 @@ import { Controller, Get, + Post, Param, Query, Req, @@ -16,6 +17,7 @@ import { ApiSecurity, ApiTags, } from '@nestjs/swagger'; +import { tasks } from '@trigger.dev/sdk'; import type { Request, Response } from 'express'; import type { Archiver } from 'archiver'; import { AuditRead } from '../../audit/skip-audit-log.decorator'; @@ -200,7 +202,8 @@ export class EvidenceExportController { } /** - * Auditor-only controller for bulk evidence export + * Auditor-only controller for bulk evidence export. + * The heavy work runs in a Trigger.dev background task to avoid OOM in the API. */ @ApiTags('Evidence Export (Auditor)') @Controller({ path: 'evidence-export', version: '1' }) @@ -209,18 +212,13 @@ export class EvidenceExportController { export class AuditorEvidenceExportController { private readonly logger = new Logger(AuditorEvidenceExportController.name); - constructor(private readonly evidenceExportService: EvidenceExportService) {} - - /** - * Export all evidence for the organization (auditor only) - */ - @Get('all') + @Post('all') @RequirePermission('evidence', 'read') @AuditRead() @ApiOperation({ - summary: 'Export all organization evidence as ZIP (Auditor only)', + summary: 'Trigger bulk evidence export (Auditor only)', description: - 'Generate and download a ZIP file containing all automation evidence across all tasks. Only accessible by auditors.', + 'Starts a background job that generates a ZIP of all evidence. Returns a run ID for progress tracking.', }) @ApiQuery({ name: 'includeJson', @@ -229,46 +227,27 @@ export class AuditorEvidenceExportController { required: false, }) @ApiResponse({ - status: 200, - description: 'ZIP file generated successfully', - content: { - 'application/zip': {}, - }, - }) - @ApiResponse({ - status: 403, - description: 'Access denied - Auditor role required', + status: 201, + description: 'Export job started', }) async exportAllEvidence( @OrganizationId() organizationId: string, @Query('includeJson') includeJson: string, - @Req() req: Request, - @Res() res: Response, ) { - this.logger.log('Auditor exporting all evidence', { + this.logger.log('Auditor triggering bulk evidence export', { organizationId, includeJson: includeJson === 'true', }); - const { archive, filename } = - await this.evidenceExportService.streamOrganizationEvidenceZip( - organizationId, - { includeRawJson: includeJson === 'true' }, - ); - - res.setHeader('Content-Type', 'application/zip'); - res.setHeader( - 'Content-Disposition', - `attachment; filename="${filename}"`, - ); - - pipeArchiveToResponse({ - archive, - req, - res, - logger: this.logger, - tag: `org ${organizationId}`, + const handle = await tasks.trigger('export-organization-evidence', { + organizationId, + includeJson: includeJson === 'true', }); + + return { + runId: handle.id, + publicAccessToken: handle.publicAccessToken, + }; } } diff --git a/apps/api/src/trigger/evidence-export/export-organization-evidence.ts b/apps/api/src/trigger/evidence-export/export-organization-evidence.ts new file mode 100644 index 0000000000..cefef6b1ca --- /dev/null +++ b/apps/api/src/trigger/evidence-export/export-organization-evidence.ts @@ -0,0 +1,275 @@ +import { metadata, schemaTask } from '@trigger.dev/sdk'; +import { z } from 'zod'; +import { + S3Client, + PutObjectCommand, + GetObjectCommand, +} from '@aws-sdk/client-s3'; +import { getSignedUrl } from '@aws-sdk/s3-request-presigner'; +import archiver from 'archiver'; +import { db } from '@db'; +import { format } from 'date-fns'; +import { + getAutomationHeaders, + streamAutomationRuns, + findTasksWithEvidence, +} from '@/tasks/evidence-export/evidence-data-loader'; +import { + generateAutomationPDFFromStream, + generateTaskSummaryPDF, + sanitizeFilename, +} from '@/tasks/evidence-export/evidence-pdf-generator'; +import { buildAutomationJsonStream } from '@/tasks/evidence-export/evidence-json-builder'; +import { + getTaskAttachments, + appendAttachmentToArchive, + createFilenameTracker, +} from '@/tasks/evidence-export/evidence-attachment-streamer'; +import { configure as configureStringify } from 'safe-stable-stringify'; + +const safeStringify = configureStringify({ + bigint: true, + circularValue: '[Circular]', + deterministic: false, +}); + +const PRESIGNED_URL_EXPIRY = 3600; + +function createS3Client(): S3Client { + const region = process.env.APP_AWS_REGION || 'us-east-1'; + const accessKeyId = process.env.APP_AWS_ACCESS_KEY_ID; + const secretAccessKey = process.env.APP_AWS_SECRET_ACCESS_KEY; + + if (!accessKeyId || !secretAccessKey) { + throw new Error( + 'AWS S3 credentials missing. Set APP_AWS_ACCESS_KEY_ID and APP_AWS_SECRET_ACCESS_KEY.', + ); + } + + return new S3Client({ + region, + credentials: { accessKeyId, secretAccessKey }, + ...(process.env.APP_AWS_ENDPOINT + ? { + endpoint: process.env.APP_AWS_ENDPOINT, + forcePathStyle: true, + } + : {}), + }); +} + +function getBucketName(): string { + const bucket = process.env.APP_AWS_BUCKET_NAME; + if (!bucket) throw new Error('APP_AWS_BUCKET_NAME is not set.'); + return bucket; +} + +export const exportOrganizationEvidenceTask = schemaTask({ + id: 'export-organization-evidence', + maxDuration: 60 * 30, + retry: { maxAttempts: 0 }, + schema: z.object({ + organizationId: z.string(), + includeJson: z.boolean().default(false), + }), + run: async ({ organizationId, includeJson }) => { + metadata.set('status', 'starting'); + metadata.set('progress', 0); + + const organization = await db.organization.findUnique({ + where: { id: organizationId }, + select: { name: true }, + }); + if (!organization) throw new Error('Organization not found'); + + const taskIds = await findTasksWithEvidence(organizationId); + if (taskIds.length === 0) throw new Error('No tasks with evidence found'); + + metadata.set('tasksTotal', taskIds.length); + metadata.set('tasksCompleted', 0); + metadata.set('status', 'generating'); + + const orgFolder = sanitizeFilename(organization.name); + const exportDate = format(new Date(), 'yyyy-MM-dd'); + const runId = metadata.get('runId') ?? crypto.randomUUID().slice(0, 8); + const s3Key = `${organizationId}/exports/evidence-${exportDate}-${runId}.zip`; + + const s3Client = createS3Client(); + const bucket = getBucketName(); + + const archive = archiver('zip', { zlib: { level: 6 } }); + const zipBuffer = await buildZipBuffer({ + archive, + organizationId, + organizationName: organization.name, + orgFolder, + taskIds, + includeJson, + }); + + metadata.set('status', 'uploading'); + + await s3Client.send( + new PutObjectCommand({ + Bucket: bucket, + Key: s3Key, + Body: zipBuffer, + ContentType: 'application/zip', + }), + ); + + metadata.set('status', 'generating-link'); + metadata.set('progress', 95); + + const downloadUrl = await getSignedUrl( + s3Client, + new GetObjectCommand({ Bucket: bucket, Key: s3Key }), + { expiresIn: PRESIGNED_URL_EXPIRY }, + ); + + metadata.set('status', 'completed'); + metadata.set('progress', 100); + metadata.set('downloadUrl', downloadUrl); + + return { downloadUrl, s3Key }; + }, +}); + +async function buildZipBuffer(params: { + archive: archiver.Archiver; + organizationId: string; + organizationName: string; + orgFolder: string; + taskIds: string[]; + includeJson: boolean; +}): Promise { + const chunks: Buffer[] = []; + params.archive.on('data', (chunk: Buffer) => chunks.push(chunk)); + const finished = new Promise((resolve, reject) => { + params.archive.on('end', resolve); + params.archive.on('error', reject); + }); + await populateArchive(params); + await finished; + return Buffer.concat(chunks); +} + +async function populateArchive({ + archive, + organizationId, + organizationName, + orgFolder, + taskIds, + includeJson, +}: { + archive: archiver.Archiver; + organizationId: string; + organizationName: string; + orgFolder: string; + taskIds: string[]; + includeJson: boolean; +}): Promise { + const manifestEntries: Array<{ + id: string; + title: string; + automations: number; + attachments: number; + }> = []; + let totalAttachments = 0; + + for (let i = 0; i < taskIds.length; i++) { + const taskId = taskIds[i]; + try { + const [headers, attachments] = await Promise.all([ + getAutomationHeaders({ organizationId, taskId }), + getTaskAttachments(organizationId, taskId), + ]); + + if (headers.automations.length === 0 && attachments.length === 0) { + continue; + } + + const taskIdSuffix = headers.taskId.slice(-8); + const taskFolder = `${orgFolder}/${sanitizeFilename(headers.taskTitle)}-${taskIdSuffix}`; + + const summaryPdf = generateTaskSummaryPDF(headers, { + attachmentsCount: attachments.length, + }); + archive.append(summaryPdf, { name: `${taskFolder}/00-summary.pdf` }); + + if (attachments.length > 0) { + const uniqueName = createFilenameTracker(); + for (const attachment of attachments) { + await appendAttachmentToArchive({ + archive, + attachment, + folderPath: `${taskFolder}/01-attachments`, + uniqueName, + }); + } + } + + for (const automationHeader of headers.automations) { + const typePrefix = + automationHeader.type === 'app_automation' ? 'app' : 'custom'; + const automationName = sanitizeFilename(automationHeader.name); + const idSuffix = automationHeader.id.slice(-8); + const filePrefix = `${taskFolder}/${typePrefix}-${automationName}-${idSuffix}`; + + const pdfBuffer = await generateAutomationPDFFromStream( + automationHeader, + { organizationName, taskTitle: headers.taskTitle }, + streamAutomationRuns({ taskId, header: automationHeader }), + ); + archive.append(pdfBuffer, { name: `${filePrefix}.pdf` }); + + if (includeJson) { + const jsonStream = buildAutomationJsonStream({ + summary: headers, + header: automationHeader, + runBatches: streamAutomationRuns({ + taskId, + header: automationHeader, + }), + }); + archive.append(jsonStream, { name: `${filePrefix}.json` }); + } + } + + manifestEntries.push({ + id: headers.taskId, + title: headers.taskTitle, + automations: headers.automations.length, + attachments: attachments.length, + }); + totalAttachments += attachments.length; + } catch (error) { + console.warn( + `Failed to export task ${taskId}: ${error instanceof Error ? error.message : String(error)}`, + ); + } + + metadata.set('tasksCompleted', i + 1); + metadata.set( + 'progress', + Math.round(((i + 1) / taskIds.length) * 90), + ); + } + + manifestEntries.sort((a, b) => a.title.localeCompare(b.title)); + + const manifest = { + organization: organizationName, + organizationId, + exportedAt: new Date().toISOString(), + tasksCount: manifestEntries.length, + totalAttachments, + tasks: manifestEntries, + }; + archive.append( + Buffer.from(safeStringify(manifest, null, 2) ?? '{}', 'utf-8'), + { name: `${orgFolder}/manifest.json` }, + ); + + await archive.finalize(); +} diff --git a/apps/app/src/app/(app)/[orgId]/auditor/(overview)/components/ExportEvidenceButton.tsx b/apps/app/src/app/(app)/[orgId]/auditor/(overview)/components/ExportEvidenceButton.tsx index 8d52fb693d..604d434031 100644 --- a/apps/app/src/app/(app)/[orgId]/auditor/(overview)/components/ExportEvidenceButton.tsx +++ b/apps/app/src/app/(app)/[orgId]/auditor/(overview)/components/ExportEvidenceButton.tsx @@ -1,6 +1,6 @@ 'use client'; -import { downloadAllEvidenceZip } from '@/lib/evidence-download'; +import { triggerBulkEvidenceExport } from '@/lib/evidence-download'; import { Button, HStack, @@ -14,78 +14,133 @@ import { Text, } from '@trycompai/design-system'; import { ArrowDown } from '@trycompai/design-system/icons'; -import { useState } from 'react'; +import { useRealtimeRun } from '@trigger.dev/react-hooks'; +import { useCallback, useState } from 'react'; import { toast } from 'sonner'; interface ExportEvidenceButtonProps { organizationName: string; } -export function ExportEvidenceButton({ organizationName }: ExportEvidenceButtonProps) { - const [isDownloading, setIsDownloading] = useState(false); +type ExportState = + | { phase: 'idle' } + | { phase: 'triggering' } + | { phase: 'running'; runId: string; accessToken: string }; + +export function ExportEvidenceButton({ + organizationName, +}: ExportEvidenceButtonProps) { const [includeJson, setIncludeJson] = useState(false); const [isOpen, setIsOpen] = useState(false); + const [exportState, setExportState] = useState({ + phase: 'idle', + }); + + const isRunning = + exportState.phase === 'triggering' || exportState.phase === 'running'; - const handleDownload = async () => { - setIsDownloading(true); + const handleTrigger = async () => { + setExportState({ phase: 'triggering' }); try { - await downloadAllEvidenceZip({ organizationName, includeJson }); - toast.success('Evidence package downloaded successfully'); - setIsOpen(false); + const { runId, publicAccessToken } = await triggerBulkEvidenceExport({ + includeJson, + }); + setExportState({ phase: 'running', runId, accessToken: publicAccessToken }); } catch (err) { - toast.error('Failed to download evidence. Please try again.'); - console.error('Evidence download error:', err); - } finally { - setIsDownloading(false); + toast.error('Failed to start evidence export. Please try again.'); + console.error('Evidence export trigger error:', err); + setExportState({ phase: 'idle' }); } }; + const handleComplete = useCallback( + (run: { output?: { downloadUrl?: string } | null; metadata?: Record }) => { + const downloadUrl = + run.output?.downloadUrl ?? + (run.metadata?.downloadUrl as string | undefined); + + if (downloadUrl) { + const link = document.createElement('a'); + link.href = downloadUrl; + link.download = `${organizationName || 'evidence'}-export.zip`; + document.body.appendChild(link); + link.click(); + document.body.removeChild(link); + toast.success('Evidence package downloaded successfully'); + } else { + toast.error('Export completed but download link was not available.'); + } + + setExportState({ phase: 'idle' }); + setIsOpen(false); + }, + [organizationName], + ); + + const handleError = useCallback(() => { + toast.error('Evidence export failed. Please try again.'); + setExportState({ phase: 'idle' }); + }, []); + return ( <> - + { + if (!isRunning) setIsOpen(open); + }}> Export All Evidence - - Download every task's uploaded evidence as a single ZIP so - you can hand it to your auditor or keep an offline snapshot. - - - - - - Include raw JSON files - - - Adds machine-readable metadata alongside the evidence - files. + {exportState.phase === 'running' ? ( + + ) : ( + <> + + Download every task's uploaded evidence as a single ZIP + so you can hand it to your auditor or keep an offline + snapshot. - - - - - - - - + + + + + Include raw JSON files + + + Adds machine-readable metadata alongside the evidence + files. + + + + + + + + + + + )} @@ -93,3 +148,68 @@ export function ExportEvidenceButton({ organizationName }: ExportEvidenceButtonP ); } + +function ExportProgress({ + runId, + accessToken, + onComplete, + onError, +}: { + runId: string; + accessToken: string; + onComplete: (run: { output?: { downloadUrl?: string } | null; metadata?: Record }) => void; + onError: () => void; +}) { + const { run } = useRealtimeRun(runId, { + accessToken, + enabled: true, + onComplete: (run) => onComplete(run), + onError: () => onError(), + }); + + const meta = run?.metadata as + | { + status?: string; + progress?: number; + tasksCompleted?: number; + tasksTotal?: number; + } + | undefined; + + const progress = meta?.progress ?? 0; + const status = meta?.status ?? 'starting'; + const tasksCompleted = meta?.tasksCompleted ?? 0; + const tasksTotal = meta?.tasksTotal ?? 0; + + const statusLabel = + status === 'starting' + ? 'Starting export...' + : status === 'generating' + ? `Processing task ${tasksCompleted} of ${tasksTotal}...` + : status === 'generating-link' + ? 'Generating download link...' + : 'Preparing...'; + + return ( + + + {statusLabel} + +
+
+
+ + This may take a few minutes for large organizations. You can close this + dialog — the export will continue in the background. + + + ); +} diff --git a/apps/app/src/app/(app)/[orgId]/tasks/components/TasksPageClient.tsx b/apps/app/src/app/(app)/[orgId]/tasks/components/TasksPageClient.tsx index 4231a9e803..6550c92f88 100644 --- a/apps/app/src/app/(app)/[orgId]/tasks/components/TasksPageClient.tsx +++ b/apps/app/src/app/(app)/[orgId]/tasks/components/TasksPageClient.tsx @@ -1,7 +1,8 @@ 'use client'; import { UpdateOrganizationEvidenceApproval } from '@/components/forms/organization/update-organization-evidence-approval'; -import { downloadAllEvidenceZip } from '@/lib/evidence-download'; +import { triggerBulkEvidenceExport } from '@/lib/evidence-download'; +import { useRealtimeRun } from '@trigger.dev/react-hooks'; import type { Member, Task, User } from '@db'; import { Button, @@ -68,31 +69,64 @@ export function TasksPageClient({ const { tasks, createTask, mutate: mutateTasks } = useTasks({ initialData: initialTasks }); const { hasPermission } = usePermissions(); const [isCreateSheetOpen, setIsCreateSheetOpen] = useState(false); - const [isDownloadingAll, setIsDownloadingAll] = useState(false); const [includeRawJson, setIncludeRawJson] = useState(false); const [isPopoverOpen, setIsPopoverOpen] = useState(false); const [mainTab, setMainTab] = useState('evidence-list'); + const [exportRun, setExportRun] = useState<{ + runId: string; + accessToken: string; + } | null>(null); + const [isTriggering, setIsTriggering] = useState(false); + + const { run: realtimeRun } = useRealtimeRun(exportRun?.runId ?? '', { + accessToken: exportRun?.accessToken, + enabled: !!exportRun, + onComplete: (run) => { + const downloadUrl = + run.output?.downloadUrl ?? + (run.metadata?.downloadUrl as string | undefined); + if (downloadUrl) { + const link = document.createElement('a'); + link.href = downloadUrl; + link.download = `${organizationName || 'evidence'}-export.zip`; + document.body.appendChild(link); + link.click(); + document.body.removeChild(link); + toast.success('Evidence package downloaded successfully'); + } + setExportRun(null); + setIsPopoverOpen(false); + }, + onError: () => { + toast.error('Evidence export failed. Please try again.'); + setExportRun(null); + }, + }); + + const isDownloadingAll = isTriggering || !!exportRun; const handleDownloadAllEvidence = async () => { - setIsDownloadingAll(true); + setIsTriggering(true); try { - await downloadAllEvidenceZip({ - organizationName: organizationName ?? undefined, + const result = await triggerBulkEvidenceExport({ includeJson: includeRawJson, }); - toast.success('Evidence package downloaded successfully'); - setIsPopoverOpen(false); + setExportRun({ + runId: result.runId, + accessToken: result.publicAccessToken, + }); + toast.info('Evidence export started. You\'ll be notified when it\'s ready.'); } catch (err) { const noEvidence = err instanceof Error && err.message?.includes('No tasks with evidence found'); if (noEvidence) { toast.info('No tasks with evidence found to export.'); } else { - toast.error('Failed to download evidence. Please try again.'); + toast.error('Failed to start evidence export. Please try again.'); } - console.error('Evidence download error:', err); + console.error('Evidence export error:', err); } finally { - setIsDownloadingAll(false); + setIsTriggering(false); } }; diff --git a/apps/app/src/lib/evidence-download.ts b/apps/app/src/lib/evidence-download.ts index 0081126dff..b11f229d24 100644 --- a/apps/app/src/lib/evidence-download.ts +++ b/apps/app/src/lib/evidence-download.ts @@ -45,24 +45,28 @@ export async function downloadTaskEvidenceZip({ } /** - * Download all evidence for the organization (auditor only) + * Trigger bulk evidence export as a background job. + * Returns a run ID and access token for tracking progress via Trigger.dev realtime. */ -export async function downloadAllEvidenceZip({ - organizationName, +export async function triggerBulkEvidenceExport({ includeJson = false, }: { - organizationName?: string; includeJson?: boolean; -}): Promise { +}): Promise<{ runId: string; publicAccessToken: string }> { const baseUrl = env.NEXT_PUBLIC_API_URL || 'http://localhost:3333'; const endpoint = `/v1/evidence-export/all?includeJson=${includeJson}`; - await downloadFile(baseUrl + endpoint, { - fallbackBaseName: organizationName - ? `${organizationName}-all-evidence` - : 'all-evidence', - fallbackExtension: 'zip', + const response = await fetch(baseUrl + endpoint, { + method: 'POST', + credentials: 'include', }); + + if (!response.ok) { + const error = await response.text(); + throw new Error(error || `Failed to start export: ${response.statusText}`); + } + + return response.json(); } /** From 5fe153891c3c7cb3988c2d91517a9de3bb96699b Mon Sep 17 00:00:00 2001 From: Mariano Fuentes Date: Thu, 4 Jun 2026 12:10:59 -0400 Subject: [PATCH 3/4] fix(evidence-export): harden bulk-export trigger task against worker OOM + retry storms MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Follow-up to the Trigger.dev offload, addressing a 2026-06-04 prod downtime (3 concurrent in-process exports for one org starved the event loop → health checks failed → ECS restarted the tasks; ~11 min degraded). - Stream the ZIP straight to S3 via @aws-sdk/lib-storage Upload + a PassThrough instead of Buffer.concat'ing the whole archive. Peak worker memory is now bounded by one automation's PDF plus the uploader's part buffer (~40MB), never the full ZIP. Populate + upload run concurrently with explicit error forwarding and symmetric upload.abort() on either failure (no orphaned multipart parts). - Run the task on machine large-1x (8GB/4vCPU) — isolated from the API. - Serialize exports per org: task queue concurrencyLimit 1 + per-org concurrencyKey, plus an org+includeJson idempotencyKey (30m TTL) so a burst of retries collapses to a single run instead of N concurrent exports. - Record partial failures in the manifest (hasFailures + failedTasks[]) and in run metadata (tasksFailed) so a partial export is self-evident from the ZIP. - Use ctx.run.id for a unique, traceable S3 key. Tests: streamArchiveToS3 happy-path + both failure paths; controller asserts the per-org concurrency/idempotency trigger options. 43 passing. Co-Authored-By: Claude Opus 4.8 (1M context) --- apps/api/package.json | 1 + .../evidence-export.controller.spec.ts | 19 +++ .../evidence-export.controller.ts | 21 ++- .../export-organization-evidence.spec.ts | 123 ++++++++++++++ .../export-organization-evidence.ts | 157 +++++++++++------- bun.lock | 17 +- 6 files changed, 276 insertions(+), 62 deletions(-) create mode 100644 apps/api/src/trigger/evidence-export/export-organization-evidence.spec.ts diff --git a/apps/api/package.json b/apps/api/package.json index acf8fe92c4..77a03743c8 100644 --- a/apps/api/package.json +++ b/apps/api/package.json @@ -58,6 +58,7 @@ "@aws-sdk/client-sts": "^3.948.0", "@aws-sdk/client-transfer": "^3.948.0", "@aws-sdk/client-wafv2": "^3.948.0", + "@aws-sdk/lib-storage": "3.1013.0", "@aws-sdk/s3-request-presigner": "3.1013.0", "@browserbasehq/sdk": "2.6.0", "@browserbasehq/stagehand": "^3.2.1", diff --git a/apps/api/src/tasks/evidence-export/evidence-export.controller.spec.ts b/apps/api/src/tasks/evidence-export/evidence-export.controller.spec.ts index 62015060f9..18c3a9fdc4 100644 --- a/apps/api/src/tasks/evidence-export/evidence-export.controller.spec.ts +++ b/apps/api/src/tasks/evidence-export/evidence-export.controller.spec.ts @@ -253,10 +253,29 @@ describe('AuditorEvidenceExportController', () => { expect(mockTrigger).toHaveBeenCalledWith( 'export-organization-evidence', { organizationId: 'org_1', includeJson: true }, + { + concurrencyKey: 'org_1', + idempotencyKey: 'evidence-export:org_1:true', + idempotencyKeyTTL: '30m', + }, ); expect(result).toEqual({ runId: 'run_123', publicAccessToken: 'tok_abc', }); }); + + it('serializes per-org and dedupes on org + includeJson (includeJson=false)', async () => { + await controller.exportAllEvidence('org_2', undefined as unknown as string); + + expect(mockTrigger).toHaveBeenCalledWith( + 'export-organization-evidence', + { organizationId: 'org_2', includeJson: false }, + { + concurrencyKey: 'org_2', + idempotencyKey: 'evidence-export:org_2:false', + idempotencyKeyTTL: '30m', + }, + ); + }); }); diff --git a/apps/api/src/tasks/evidence-export/evidence-export.controller.ts b/apps/api/src/tasks/evidence-export/evidence-export.controller.ts index 5eb3601dd3..206c7cafd7 100644 --- a/apps/api/src/tasks/evidence-export/evidence-export.controller.ts +++ b/apps/api/src/tasks/evidence-export/evidence-export.controller.ts @@ -234,15 +234,26 @@ export class AuditorEvidenceExportController { @OrganizationId() organizationId: string, @Query('includeJson') includeJson: string, ) { + const includeJsonBool = includeJson === 'true'; this.logger.log('Auditor triggering bulk evidence export', { organizationId, - includeJson: includeJson === 'true', + includeJson: includeJsonBool, }); - const handle = await tasks.trigger('export-organization-evidence', { - organizationId, - includeJson: includeJson === 'true', - }); + const handle = await tasks.trigger( + 'export-organization-evidence', + { organizationId, includeJson: includeJsonBool }, + { + // Serialize exports per org (the task queue's concurrencyLimit is 1, and + // concurrencyKey gives each org its own lane) so a burst of clicks can + // never run multiple heavy exports for the same org at once. + concurrencyKey: organizationId, + // Collapse rapid duplicate triggers (double-click / retry) for the same + // org + options into a single run for the TTL window. + idempotencyKey: `evidence-export:${organizationId}:${includeJsonBool}`, + idempotencyKeyTTL: '30m', + }, + ); return { runId: handle.id, diff --git a/apps/api/src/trigger/evidence-export/export-organization-evidence.spec.ts b/apps/api/src/trigger/evidence-export/export-organization-evidence.spec.ts new file mode 100644 index 0000000000..87da540128 --- /dev/null +++ b/apps/api/src/trigger/evidence-export/export-organization-evidence.spec.ts @@ -0,0 +1,123 @@ +import { EventEmitter } from 'node:events'; + +// Mock module boundaries so importing the task does not connect to Postgres, +// the Trigger SDK, S3, or pull in jsPDF. We only exercise streamArchiveToS3 — +// the concurrent populate+upload orchestration and its error propagation. +jest.mock('@db', () => ({ db: { organization: { findUnique: jest.fn() } } })); + +jest.mock('@trigger.dev/sdk', () => ({ + metadata: { set: jest.fn(), get: jest.fn() }, + schemaTask: (config: unknown) => config, +})); + +const mockUploadDone = jest.fn(); +const mockUploadAbort = jest.fn().mockResolvedValue(undefined); +jest.mock('@aws-sdk/lib-storage', () => ({ + Upload: jest.fn().mockImplementation((opts: { params?: { Body?: unknown } }) => { + // A real Upload consumes the Body stream and surfaces its errors via done(). + // Mimic that so a destroyed PassThrough doesn't become an unhandled error. + const body = opts?.params?.Body as + | { on?: (e: string, cb: () => void) => void; resume?: () => void } + | undefined; + if (body && typeof body.on === 'function') { + body.on('error', () => {}); + if (typeof body.resume === 'function') body.resume(); + } + return { done: mockUploadDone, abort: mockUploadAbort }; + }), +})); + +jest.mock('@aws-sdk/client-s3', () => ({ + S3Client: jest.fn(), + GetObjectCommand: jest.fn(), +})); +jest.mock('@aws-sdk/s3-request-presigner', () => ({ getSignedUrl: jest.fn() })); + +// Fake archiver: an EventEmitter with the methods streamArchiveToS3 touches. +function makeFakeArchive() { + const emitter = new EventEmitter(); + return Object.assign(emitter, { + pipe: jest.fn(), + finalize: jest.fn().mockResolvedValue(undefined), + abort: jest.fn(), + append: jest.fn(), + }); +} +let fakeArchive: ReturnType; +jest.mock('archiver', () => jest.fn(() => fakeArchive)); + +// The task module pulls these in at import time; stub them (unused by the SUT). +jest.mock('@/tasks/evidence-export/evidence-data-loader', () => ({ + getAutomationHeaders: jest.fn(), + streamAutomationRuns: jest.fn(), + findTasksWithEvidence: jest.fn(), +})); +jest.mock('@/tasks/evidence-export/evidence-pdf-generator', () => ({ + generateAutomationPDFFromStream: jest.fn(), + generateTaskSummaryPDF: jest.fn(), + sanitizeFilename: (s: string) => s, +})); +jest.mock('@/tasks/evidence-export/evidence-json-builder', () => ({ + buildAutomationJsonStream: jest.fn(), +})); +jest.mock('@/tasks/evidence-export/evidence-attachment-streamer', () => ({ + getTaskAttachments: jest.fn(), + appendAttachmentToArchive: jest.fn(), + createFilenameTracker: jest.fn(), +})); + +import { streamArchiveToS3 } from './export-organization-evidence'; + +describe('streamArchiveToS3', () => { + const s3Client = {} as never; + const baseParams = { s3Client, bucket: 'b', key: 'k' }; + + beforeEach(() => { + fakeArchive = makeFakeArchive(); + mockUploadDone.mockReset().mockResolvedValue(undefined); + mockUploadAbort.mockClear(); + }); + + it('finalizes the archive once and resolves on the happy path', async () => { + const populate = jest.fn().mockResolvedValue(undefined); + + await streamArchiveToS3({ ...baseParams, populate }); + + expect(populate).toHaveBeenCalledWith(fakeArchive); + expect(fakeArchive.finalize).toHaveBeenCalledTimes(1); + expect(fakeArchive.pipe).toHaveBeenCalledTimes(1); + expect(mockUploadDone).toHaveBeenCalledTimes(1); + expect(mockUploadAbort).not.toHaveBeenCalled(); + }); + + it('aborts the archive + upload and rethrows when populate fails', async () => { + const boom = new Error('populate failed'); + const populate = jest.fn().mockRejectedValue(boom); + // Upload would hang until the stream ends; resolve it so allSettled settles. + mockUploadDone.mockResolvedValue(undefined); + + await expect(streamArchiveToS3({ ...baseParams, populate })).rejects.toBe( + boom, + ); + + expect(fakeArchive.abort).toHaveBeenCalledTimes(1); + expect(fakeArchive.finalize).not.toHaveBeenCalled(); + expect(mockUploadAbort).toHaveBeenCalledTimes(1); + }); + + it('aborts the multipart upload and rethrows when the S3 upload fails', async () => { + const uploadErr = new Error('s3 upload failed'); + const populate = jest.fn().mockResolvedValue(undefined); + mockUploadDone.mockRejectedValue(uploadErr); + + await expect(streamArchiveToS3({ ...baseParams, populate })).rejects.toBe( + uploadErr, + ); + + expect(fakeArchive.finalize).toHaveBeenCalledTimes(1); + // populate succeeded, so we do not abort the archive itself... + expect(fakeArchive.abort).not.toHaveBeenCalled(); + // ...but we DO cancel the multipart upload to avoid orphaned S3 parts. + expect(mockUploadAbort).toHaveBeenCalledTimes(1); + }); +}); diff --git a/apps/api/src/trigger/evidence-export/export-organization-evidence.ts b/apps/api/src/trigger/evidence-export/export-organization-evidence.ts index cefef6b1ca..649dd019b3 100644 --- a/apps/api/src/trigger/evidence-export/export-organization-evidence.ts +++ b/apps/api/src/trigger/evidence-export/export-organization-evidence.ts @@ -1,10 +1,8 @@ import { metadata, schemaTask } from '@trigger.dev/sdk'; import { z } from 'zod'; -import { - S3Client, - PutObjectCommand, - GetObjectCommand, -} from '@aws-sdk/client-s3'; +import { PassThrough } from 'node:stream'; +import { S3Client, GetObjectCommand } from '@aws-sdk/client-s3'; +import { Upload } from '@aws-sdk/lib-storage'; import { getSignedUrl } from '@aws-sdk/s3-request-presigner'; import archiver from 'archiver'; import { db } from '@db'; @@ -34,6 +32,10 @@ const safeStringify = configureStringify({ }); const PRESIGNED_URL_EXPIRY = 3600; +// 10 MB parts; the multipart uploader buffers at most queueSize * partSize (~40 MB), +// so worker memory stays flat regardless of total ZIP size. +const UPLOAD_PART_SIZE = 10 * 1024 * 1024; +const UPLOAD_QUEUE_SIZE = 4; function createS3Client(): S3Client { const region = process.env.APP_AWS_REGION || 'us-east-1'; @@ -50,10 +52,7 @@ function createS3Client(): S3Client { region, credentials: { accessKeyId, secretAccessKey }, ...(process.env.APP_AWS_ENDPOINT - ? { - endpoint: process.env.APP_AWS_ENDPOINT, - forcePathStyle: true, - } + ? { endpoint: process.env.APP_AWS_ENDPOINT, forcePathStyle: true } : {}), }); } @@ -66,13 +65,19 @@ function getBucketName(): string { export const exportOrganizationEvidenceTask = schemaTask({ id: 'export-organization-evidence', + // Runs on an isolated worker; 8 GB / 4 vCPU gives ample headroom for jsPDF + + // zlib across the largest orgs now that the ZIP streams to S3 (never buffered). + machine: { preset: 'large-1x' }, + // concurrencyLimit 1 + a per-org concurrencyKey (passed at trigger time) means + // at most one export runs per org at a time; different orgs still run in parallel. + queue: { name: 'evidence-export', concurrencyLimit: 1 }, maxDuration: 60 * 30, retry: { maxAttempts: 0 }, schema: z.object({ organizationId: z.string(), includeJson: z.boolean().default(false), }), - run: async ({ organizationId, includeJson }) => { + run: async ({ organizationId, includeJson }, { ctx }) => { metadata.set('status', 'starting'); metadata.set('progress', 0); @@ -91,33 +96,26 @@ export const exportOrganizationEvidenceTask = schemaTask({ const orgFolder = sanitizeFilename(organization.name); const exportDate = format(new Date(), 'yyyy-MM-dd'); - const runId = metadata.get('runId') ?? crypto.randomUUID().slice(0, 8); - const s3Key = `${organizationId}/exports/evidence-${exportDate}-${runId}.zip`; + const s3Key = `${organizationId}/exports/evidence-${exportDate}-${ctx.run.id}.zip`; const s3Client = createS3Client(); const bucket = getBucketName(); - const archive = archiver('zip', { zlib: { level: 6 } }); - const zipBuffer = await buildZipBuffer({ - archive, - organizationId, - organizationName: organization.name, - orgFolder, - taskIds, - includeJson, + await streamArchiveToS3({ + s3Client, + bucket, + key: s3Key, + populate: (archive) => + populateArchive({ + archive, + organizationId, + organizationName: organization.name, + orgFolder, + taskIds, + includeJson, + }), }); - metadata.set('status', 'uploading'); - - await s3Client.send( - new PutObjectCommand({ - Bucket: bucket, - Key: s3Key, - Body: zipBuffer, - ContentType: 'application/zip', - }), - ); - metadata.set('status', 'generating-link'); metadata.set('progress', 95); @@ -135,23 +133,68 @@ export const exportOrganizationEvidenceTask = schemaTask({ }, }); -async function buildZipBuffer(params: { - archive: archiver.Archiver; - organizationId: string; - organizationName: string; - orgFolder: string; - taskIds: string[]; - includeJson: boolean; -}): Promise { - const chunks: Buffer[] = []; - params.archive.on('data', (chunk: Buffer) => chunks.push(chunk)); - const finished = new Promise((resolve, reject) => { - params.archive.on('end', resolve); - params.archive.on('error', reject); +/** + * Pipe a freshly-built ZIP archive straight to S3 via multipart upload. The + * archive is populated and uploaded concurrently, so peak memory is bounded by + * one automation's PDF plus the uploader's part buffer — never the whole ZIP. + */ +export async function streamArchiveToS3(params: { + s3Client: S3Client; + bucket: string; + key: string; + populate: (archive: archiver.Archiver) => Promise; +}): Promise { + const { s3Client, bucket, key, populate } = params; + + const archive = archiver('zip', { zlib: { level: 6 } }); + const passThrough = new PassThrough(); + + archive.on('warning', (err) => { + console.warn(`Archive warning (${key}): ${err.message}`); + }); + // pipe() does not forward source errors to the destination — do it explicitly + // so a failed archive ends the upload stream and upload.done() rejects. + archive.on('error', (err) => passThrough.destroy(err)); + archive.pipe(passThrough); + + const upload = new Upload({ + client: s3Client, + params: { + Bucket: bucket, + Key: key, + Body: passThrough, + ContentType: 'application/zip', + }, + queueSize: UPLOAD_QUEUE_SIZE, + partSize: UPLOAD_PART_SIZE, }); - await populateArchive(params); - await finished; - return Buffer.concat(chunks); + + const populatePromise = (async () => { + try { + await populate(archive); + await archive.finalize(); + } catch (err) { + archive.abort(); + passThrough.destroy(err instanceof Error ? err : new Error(String(err))); + throw err; + } + })(); + + // allSettled so a populate failure cannot leave upload.done() pending forever. + const [populateResult, uploadResult] = await Promise.allSettled([ + populatePromise, + upload.done(), + ]); + + if (populateResult.status === 'rejected') { + await upload.abort().catch(() => {}); + throw populateResult.reason; + } + if (uploadResult.status === 'rejected') { + // Cancel the multipart upload so no orphaned parts linger on S3. + await upload.abort().catch(() => {}); + throw uploadResult.reason; + } } async function populateArchive({ @@ -175,6 +218,7 @@ async function populateArchive({ automations: number; attachments: number; }> = []; + const failedTasks: Array<{ taskId: string; reason: string }> = []; let totalAttachments = 0; for (let i = 0; i < taskIds.length; i++) { @@ -244,32 +288,33 @@ async function populateArchive({ }); totalAttachments += attachments.length; } catch (error) { - console.warn( - `Failed to export task ${taskId}: ${error instanceof Error ? error.message : String(error)}`, - ); + const reason = error instanceof Error ? error.message : String(error); + console.warn(`Failed to export task ${taskId}: ${reason}`); + failedTasks.push({ taskId, reason }); } metadata.set('tasksCompleted', i + 1); - metadata.set( - 'progress', - Math.round(((i + 1) / taskIds.length) * 90), - ); + metadata.set('tasksFailed', failedTasks.length); + metadata.set('progress', Math.round(((i + 1) / taskIds.length) * 90)); } manifestEntries.sort((a, b) => a.title.localeCompare(b.title)); + // Surface partial failures inside the ZIP itself so an auditor reading only the + // archive can tell the export is incomplete (not just via the Trigger run UI). const manifest = { organization: organizationName, organizationId, exportedAt: new Date().toISOString(), tasksCount: manifestEntries.length, totalAttachments, + hasFailures: failedTasks.length > 0, + failedTasks, tasks: manifestEntries, }; archive.append( Buffer.from(safeStringify(manifest, null, 2) ?? '{}', 'utf-8'), { name: `${orgFolder}/manifest.json` }, ); - - await archive.finalize(); + // Note: archive.finalize() is owned by streamArchiveToS3 (the caller). } diff --git a/bun.lock b/bun.lock index e5cfd7955e..d502aa06a9 100644 --- a/bun.lock +++ b/bun.lock @@ -125,6 +125,7 @@ "@aws-sdk/client-sts": "^3.948.0", "@aws-sdk/client-transfer": "^3.948.0", "@aws-sdk/client-wafv2": "^3.948.0", + "@aws-sdk/lib-storage": "3.1013.0", "@aws-sdk/s3-request-presigner": "3.1013.0", "@browserbasehq/sdk": "2.6.0", "@browserbasehq/stagehand": "^3.2.1", @@ -1007,6 +1008,8 @@ "@aws-sdk/endpoint-cache": ["@aws-sdk/endpoint-cache@3.972.5", "", { "dependencies": { "mnemonist": "0.38.3", "tslib": "^2.6.2" } }, "sha512-itVdge0NozgtgmtbZ25FVwWU3vGlE7x7feE/aOEJNkQfEpbkrF8Rj1QmnK+2blFfYE1xWt/iU+6/jUp/pv1+MA=="], + "@aws-sdk/lib-storage": ["@aws-sdk/lib-storage@3.1013.0", "", { "dependencies": { "@smithy/abort-controller": "^4.2.12", "@smithy/middleware-endpoint": "^4.4.26", "@smithy/smithy-client": "^4.12.6", "buffer": "5.6.0", "events": "3.3.0", "stream-browserify": "3.0.0", "tslib": "^2.6.2" }, "peerDependencies": { "@aws-sdk/client-s3": "^3.1013.0" } }, "sha512-I2fT+ve+R5iRiULrr8WqlhBnO2yXSXAc3FF8OGWjYph3zXNdIzZJC6l/o18iYHxbPlHXxp8g3tCZ42XdGmAfLQ=="], + "@aws-sdk/middleware-bucket-endpoint": ["@aws-sdk/middleware-bucket-endpoint@3.972.10", "", { "dependencies": { "@aws-sdk/types": "^3.973.8", "@aws-sdk/util-arn-parser": "^3.972.3", "@smithy/node-config-provider": "^4.3.14", "@smithy/protocol-http": "^5.3.14", "@smithy/types": "^4.14.1", "@smithy/util-config-provider": "^4.2.2", "tslib": "^2.6.2" } }, "sha512-Vbc2frZH7wXlMNd+ZZSXUEs/l1Sv8Jj4zUnIfwrYF5lwaLdXHZ9xx4U3rjUcaye3HRhFVc+E5DbBxpRAbB16BA=="], "@aws-sdk/middleware-endpoint-discovery": ["@aws-sdk/middleware-endpoint-discovery@3.972.11", "", { "dependencies": { "@aws-sdk/endpoint-cache": "^3.972.5", "@aws-sdk/types": "^3.973.8", "@smithy/node-config-provider": "^4.3.14", "@smithy/protocol-http": "^5.3.14", "@smithy/types": "^4.14.1", "tslib": "^2.6.2" } }, "sha512-vXARCZVFQHdsd6qPPZyC/hh+5x2XsCYKqUQDCqnUlpGpChMpDojOOacQWdLJ+FFXKN8X3cmLOGrtgx/zysCKqQ=="], @@ -2373,6 +2376,8 @@ "@sinonjs/fake-timers": ["@sinonjs/fake-timers@15.3.2", "", { "dependencies": { "@sinonjs/commons": "^3.0.1" } }, "sha512-mrn35Jl2pCpns+mE3HaZa1yPN5EYCRgiMI+135COjr2hr8Cls9DXqIZ57vZe2cz7y2XVSq92tcs6kGQcT1J8Rw=="], + "@smithy/abort-controller": ["@smithy/abort-controller@4.2.16", "", { "dependencies": { "@smithy/types": "^4.14.3", "tslib": "^2.6.2" } }, "sha512-YH0c/t1n8dhzz/NyIEoqoHyn4tBvrtESNRUV+ar6to4eqA3jN2OGqhe9hAicmVCokd+EC3puYXsMAG6PDZVEgg=="], + "@smithy/chunked-blob-reader": ["@smithy/chunked-blob-reader@5.2.2", "", { "dependencies": { "tslib": "^2.6.2" } }, "sha512-St+kVicSyayWQca+I1rGitaOEH6uKgE8IUWoYnnEX26SWdWQcL6LvMSD19Lg+vYHKdT9B2Zuu7rd3i6Wnyb/iw=="], "@smithy/chunked-blob-reader-native": ["@smithy/chunked-blob-reader-native@4.2.3", "", { "dependencies": { "@smithy/util-base64": "^4.3.2", "tslib": "^2.6.2" } }, "sha512-jA5k5Udn7Y5717L86h4EIv06wIr3xn8GM1qHRi/Nf31annXcXHJjBKvgztnbn2TxH3xWrPBfgwHsOwZf0UmQWw=="], @@ -3425,7 +3430,7 @@ "bser": ["bser@2.1.1", "", { "dependencies": { "node-int64": "^0.4.0" } }, "sha512-gQxTNE/GAfIIrmHLUE3oJyp5FO6HRBfhjnw4/wMmA63ZGDJnWBmgY/lyQBpnDUkGmAhbSe39tx2d/iTOAfglwQ=="], - "buffer": ["buffer@6.0.3", "", { "dependencies": { "base64-js": "^1.3.1", "ieee754": "^1.2.1" } }, "sha512-FTiCpNxtwiZZHEZbcbTIcZjERVICn9yq/pDFkTl95/AxzD1naBctN7YO68riM/gLSDY7sdrMby8hofADYuuqOA=="], + "buffer": ["buffer@5.6.0", "", { "dependencies": { "base64-js": "^1.0.2", "ieee754": "^1.1.4" } }, "sha512-/gDYp/UtU0eA1ys8bOs9J6a+E/KWIY+DZ+Q2WESNUA0jFRsJOc0SNUO6xJ5SGA1xueg3NL65W6s+NY5l9cunuw=="], "buffer-crc32": ["buffer-crc32@1.0.0", "", {}, "sha512-Db1SbgBS/fg/392AblrMJk97KggmvYhr4pB5ZIMTWtaivCPMWLkmb7m21cJvpvgK+J3nsU2CmmixNBZx4vFj/w=="], @@ -6095,6 +6100,8 @@ "stoppable": ["stoppable@1.1.0", "", {}, "sha512-KXDYZ9dszj6bzvnEMRYvxgeTHU74QBFL54XKtP3nyMuJ81CFYtABZ3bAzL2EdFUaEwJOBOgENyFj3R7oTzDyyw=="], + "stream-browserify": ["stream-browserify@3.0.0", "", { "dependencies": { "inherits": "~2.0.4", "readable-stream": "^3.5.0" } }, "sha512-H73RAHsVBapbim0tU2JwwOiXUj+fikfiaoYAKHF3VJfA0pe2BCzkhAHBlLG6REzE+2WNZcxOXjK7lkso+9euLA=="], + "stream-combiner2": ["stream-combiner2@1.1.1", "", { "dependencies": { "duplexer2": "~0.1.0", "readable-stream": "^2.0.2" } }, "sha512-3PnJbYgS56AeWgtKF5jtJRT6uFJe56Z0Hc5Ngg/6sI6rIt8iiMBTa9cvdyFfpMQjaVHr8dusbNeFGIIonxOvKw=="], "streamdown": ["streamdown@2.5.0", "", { "dependencies": { "clsx": "^2.1.1", "hast-util-to-jsx-runtime": "^2.3.6", "html-url-attributes": "^3.0.1", "marked": "^17.0.1", "mermaid": "^11.12.2", "rehype-harden": "^1.1.8", "rehype-raw": "^7.0.0", "rehype-sanitize": "^6.0.0", "remark-gfm": "^4.0.1", "remark-parse": "^11.0.0", "remark-rehype": "^11.1.2", "remend": "1.3.0", "tailwind-merge": "^3.4.0", "unified": "^11.0.5", "unist-util-visit": "^5.0.0", "unist-util-visit-parents": "^6.0.0" }, "peerDependencies": { "react": "^18.0.0 || ^19.0.0", "react-dom": "^18.0.0 || ^19.0.0" } }, "sha512-/tTnURfIOxZK/pqJAxsfCvETG/XCJHoWnk3jq9xLcuz6CSpnjjuxSRBTTL4PKGhxiZQf0lqPxGhImdpwcZ2XwA=="], @@ -7183,6 +7190,8 @@ "@react-email/components/@react-email/render": ["@react-email/render@1.1.2", "", { "dependencies": { "html-to-text": "^9.0.5", "prettier": "^3.5.3", "react-promise-suspense": "^0.3.4" }, "peerDependencies": { "react": "^18.0 || ^19.0 || ^19.0.0-rc", "react-dom": "^18.0 || ^19.0 || ^19.0.0-rc" } }, "sha512-RnRehYN3v9gVlNMehHPHhyp2RQo7+pSkHDtXPvg3s0GbzM9SQMW4Qrf8GRNvtpLC4gsI+Wt0VatNRUFqjvevbw=="], + "@react-three/fiber/buffer": ["buffer@6.0.3", "", { "dependencies": { "base64-js": "^1.3.1", "ieee754": "^1.2.1" } }, "sha512-FTiCpNxtwiZZHEZbcbTIcZjERVICn9yq/pDFkTl95/AxzD1naBctN7YO68riM/gLSDY7sdrMby8hofADYuuqOA=="], + "@react-three/postprocessing/maath": ["maath@0.6.0", "", { "peerDependencies": { "@types/three": ">=0.144.0", "three": ">=0.144.0" } }, "sha512-dSb2xQuP7vDnaYqfoKzlApeRcR2xtN8/f7WV/TMAkBC8552TwTLtOO0JTcSygkYMjNDPoo6V01jTw/aPi4JrMw=="], "@rollup/plugin-commonjs/estree-walker": ["estree-walker@2.0.2", "", {}, "sha512-Rfkk/Mp/DL7JVje3u18FxFujQlTNR2q6QfMSMB7AvCBx91NGj/ba3kCfza0f6dVDbw7YlRf/nDrn7pQrCCyQ/w=="], @@ -7231,6 +7240,8 @@ "@sentry/vercel-edge/@opentelemetry/api": ["@opentelemetry/api@1.9.1", "", {}, "sha512-gLyJlPHPZYdAk1JENA9LeHejZe1Ti77/pTeFm/nMXmQH/HFZlcS/O2XJB+L8fkbrNSqhdtlvjBVjxwUYanNH5Q=="], + "@smithy/abort-controller/@smithy/types": ["@smithy/types@4.14.3", "", { "dependencies": { "tslib": "^2.6.2" } }, "sha512-YupL0ZWmFtJexUN2cHzkvvF/b9pKrtAIfT1o7/oY/Ppu8IYeZ+lDPM5vZdQJaSeA132dJCqojjGC9NhXeF71VQ=="], + "@smithy/middleware-compression/fflate": ["fflate@0.8.1", "", {}, "sha512-/exOvEuc+/iaUm105QIiOt4LpBdMTWsXxqR0HDF35vx3fmaKzw7354gTilCh5rkzEt8WYyG//ku3h3nRmd7CHQ=="], "@streamdown/code/shiki": ["shiki@3.23.0", "", { "dependencies": { "@shikijs/core": "3.23.0", "@shikijs/engine-javascript": "3.23.0", "@shikijs/engine-oniguruma": "3.23.0", "@shikijs/langs": "3.23.0", "@shikijs/themes": "3.23.0", "@shikijs/types": "3.23.0", "@shikijs/vscode-textmate": "^10.0.2", "@types/hast": "^3.0.4" } }, "sha512-55Dj73uq9ZXL5zyeRPzHQsK7Nbyt6Y10k5s7OjuFZGMhpp4r/rsLBH0o/0fstIzX1Lep9VxefWljK/SKCzygIA=="], @@ -8185,6 +8196,8 @@ "read-yaml-file/strip-bom": ["strip-bom@4.0.0", "", {}, "sha512-3xurFv5tEgii33Zi8Jtp55wEIILR9eh34FAW00PZf+JnSsTmV/ioewSgQl97JHvgjoRGwPShsWm+IdrxB35d0w=="], + "readable-stream/buffer": ["buffer@6.0.3", "", { "dependencies": { "base64-js": "^1.3.1", "ieee754": "^1.2.1" } }, "sha512-FTiCpNxtwiZZHEZbcbTIcZjERVICn9yq/pDFkTl95/AxzD1naBctN7YO68riM/gLSDY7sdrMby8hofADYuuqOA=="], + "readdir-glob/minimatch": ["minimatch@5.1.9", "", { "dependencies": { "brace-expansion": "^2.0.1" } }, "sha512-7o1wEA2RyMP7Iu7GNba9vc0RWWGACJOCZBJX2GJWip0ikV+wcOsgVuY9uE8CPiyQhkGFSlhuSkZPavN7u1c2Fw=="], "recast/source-map": ["source-map@0.6.1", "", {}, "sha512-UjgapumWlbMhkBgzT7Ykc5YXUT46F0iKu8SGXq0bcwP5dz/h0Plj6enJqjz1Zbq2l5WaqYnrVbwWOWMyF3F47g=="], @@ -8265,6 +8278,8 @@ "stats-gl/three": ["three@0.170.0", "", {}, "sha512-FQK+LEpYc0fBD+J8g6oSEyyNzjp+Q7Ks1C568WWaoMRLW+TkNNWmenWeGgJjV105Gd+p/2ql1ZcjYvNiPZBhuQ=="], + "stream-browserify/readable-stream": ["readable-stream@3.6.2", "", { "dependencies": { "inherits": "^2.0.3", "string_decoder": "^1.1.1", "util-deprecate": "^1.0.1" } }, "sha512-9u/sniCrY3D5WdsERHzHE4G2YCXqoG5FTHUiCC4SIbr6XcLZBY05ya9EKjYek9O5xOAwjGq+1JdGBAS7Q9ScoA=="], + "stream-combiner2/readable-stream": ["readable-stream@2.3.8", "", { "dependencies": { "core-util-is": "~1.0.0", "inherits": "~2.0.3", "isarray": "~1.0.0", "process-nextick-args": "~2.0.0", "safe-buffer": "~5.1.1", "string_decoder": "~1.1.1", "util-deprecate": "~1.0.1" } }, "sha512-8p0AUk4XODgIewSi0l8Epjs+EVnWiK7NoDIEGU0HhE7+ZyY8D1IMY7odu5lRrFXGg71L15KG8QrPmum45RTtdA=="], "streamdown/marked": ["marked@17.0.6", "", { "bin": { "marked": "bin/marked.js" } }, "sha512-gB0gkNafnonOw0obSTEGZTT86IuhILt2Wfx0mWH/1Au83kybTayroZ/V6nS25mN7u8ASy+5fMhgB3XPNrOZdmA=="], From e91c547053d6a82e8ac4b0a1f5dbb1593b0252f1 Mon Sep 17 00:00:00 2001 From: Tofik Hasanov Date: Fri, 5 Jun 2026 12:34:40 -0400 Subject: [PATCH 4/4] fix(evidence-export): repair bulk-export progress UI (dialog trap, silent failures, invalid hook option) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three issues in the new bulk-evidence-export progress UI: 1. ExportEvidenceButton: the running export sheet could not be dismissed — onOpenChange swallowed every close while a run was active, directly contradicting the in-sheet copy ("You can close this dialog — the export will continue in the background"). Lifted the useRealtimeRun subscription from the unmount-on-close ExportProgress child up to the parent so the run keeps streaming and still auto-downloads on completion after the sheet is closed, then made the sheet freely dismissable (onOpenChange={setIsOpen}). (reported by cubic) 2. TasksPageClient: onComplete handled the happy path but, when the finished run had no downloadUrl, it silently closed the export UI with no download and no feedback — a completed-but-invalid export looked like a no-op. Added the missing error toast, matching ExportEvidenceButton. (reported by cubic) 3. Both files passed an `onError` option to useRealtimeRun, which does not exist on UseRealtimeSingleRunOptions — a type error that was failing the app build. Removed it and folded failure handling into onComplete(run, err), treating any non-COMPLETED terminal state (or an err) as a failure, the same way PolicyHeaderActions does. Tests: new ExportEvidenceButton.test.tsx (dismiss-while-running, download + success, metadata fallback, missing-link error, run-failure error, onComplete err) and added export-flow cases to TasksPageClient.test.tsx; also fixed that file's stale `downloadAllEvidenceZip` mock (now triggerBulkEvidenceExport). Co-Authored-By: Claude Opus 4.8 (1M context) --- .../components/ExportEvidenceButton.test.tsx | 234 ++++++++++++++++++ .../components/ExportEvidenceButton.tsx | 96 +++---- .../tasks/components/TasksPageClient.test.tsx | 113 ++++++++- .../tasks/components/TasksPageClient.tsx | 20 +- 4 files changed, 409 insertions(+), 54 deletions(-) create mode 100644 apps/app/src/app/(app)/[orgId]/auditor/(overview)/components/ExportEvidenceButton.test.tsx diff --git a/apps/app/src/app/(app)/[orgId]/auditor/(overview)/components/ExportEvidenceButton.test.tsx b/apps/app/src/app/(app)/[orgId]/auditor/(overview)/components/ExportEvidenceButton.test.tsx new file mode 100644 index 0000000000..9a466d39df --- /dev/null +++ b/apps/app/src/app/(app)/[orgId]/auditor/(overview)/components/ExportEvidenceButton.test.tsx @@ -0,0 +1,234 @@ +import { act, fireEvent, render, screen, waitFor } from '@testing-library/react'; +import type { ReactNode } from 'react'; +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; + +// Capture the options passed to useRealtimeRun so tests can invoke onComplete +// directly (simulating a Trigger.dev run reaching a terminal state). +const mockRealtime: { + onComplete?: ( + run: { + status?: string; + output?: { downloadUrl?: string } | null; + metadata?: Record; + }, + err?: Error, + ) => void; +} = {}; + +vi.mock('@trigger.dev/react-hooks', () => ({ + useRealtimeRun: ( + _runId: string, + options: { onComplete?: typeof mockRealtime.onComplete }, + ) => { + mockRealtime.onComplete = options.onComplete; + return { run: undefined }; + }, +})); + +const mockTriggerBulkEvidenceExport = vi.fn(); +vi.mock('@/lib/evidence-download', () => ({ + triggerBulkEvidenceExport: (args: { includeJson?: boolean }) => + mockTriggerBulkEvidenceExport(args), +})); + +const mockToastSuccess = vi.fn(); +const mockToastError = vi.fn(); +vi.mock('sonner', () => ({ + toast: { + success: (msg: string) => mockToastSuccess(msg), + error: (msg: string) => mockToastError(msg), + }, +})); + +// Minimal design-system mocks so the component renders in jsdom. +vi.mock('@trycompai/design-system', () => ({ + Button: ({ + children, + onClick, + disabled, + }: { + children: ReactNode; + onClick?: () => void; + disabled?: boolean; + }) => ( + + ), + HStack: ({ children }: { children: ReactNode }) =>
{children}
, + Stack: ({ children }: { children: ReactNode }) =>
{children}
, + Text: ({ children }: { children: ReactNode }) => {children}, + Switch: ({ + checked, + onCheckedChange, + }: { + checked?: boolean; + onCheckedChange?: (checked: boolean) => void; + }) => ( + onCheckedChange?.(e.target.checked)} + /> + ), + Sheet: ({ + children, + open, + onOpenChange, + }: { + children: ReactNode; + open: boolean; + onOpenChange: (open: boolean) => void; + }) => + open ? ( +
+ {/* Stands in for the overlay / ESC / X close affordances, all of which + route through onOpenChange. */} +
+ ) : null, + SheetContent: ({ children }: { children: ReactNode }) =>
{children}
, + SheetHeader: ({ children }: { children: ReactNode }) =>
{children}
, + SheetTitle: ({ children }: { children: ReactNode }) =>

{children}

, + SheetBody: ({ children }: { children: ReactNode }) =>
{children}
, +})); + +vi.mock('@trycompai/design-system/icons', () => ({ + ArrowDown: () => , +})); + +import { ExportEvidenceButton } from './ExportEvidenceButton'; + +async function startExport() { + fireEvent.click(screen.getByRole('button', { name: 'Export All Evidence' })); + fireEvent.click(screen.getByRole('button', { name: 'Export' })); + // Wait for the trigger promise to resolve and the running UI to render. + await waitFor(() => + expect(screen.getByText('Starting export...')).toBeInTheDocument(), + ); +} + +describe('ExportEvidenceButton', () => { + beforeEach(() => { + vi.clearAllMocks(); + mockRealtime.onComplete = undefined; + mockTriggerBulkEvidenceExport.mockResolvedValue({ + runId: 'run_1', + publicAccessToken: 'tok_1', + }); + vi.spyOn(HTMLAnchorElement.prototype, 'click').mockImplementation(() => {}); + }); + + afterEach(() => { + vi.restoreAllMocks(); + }); + + it('opens the export sheet when the trigger button is clicked', () => { + render(); + + expect(screen.queryByTestId('sheet')).not.toBeInTheDocument(); + fireEvent.click( + screen.getByRole('button', { name: 'Export All Evidence' }), + ); + expect(screen.getByTestId('sheet')).toBeInTheDocument(); + expect(screen.getByText('Include raw JSON files')).toBeInTheDocument(); + }); + + it('can be dismissed while an export is running (does not trap the user)', async () => { + render(); + await startExport(); + + // The copy promises the export can be closed and continues in the + // background — the sheet must honor a close request while running. + fireEvent.click(screen.getByTestId('sheet-request-close')); + + await waitFor(() => + expect(screen.queryByTestId('sheet')).not.toBeInTheDocument(), + ); + }); + + it('auto-downloads and toasts success when the run completes with a download URL', async () => { + const clickSpy = vi.spyOn(HTMLAnchorElement.prototype, 'click'); + render(); + await startExport(); + + act(() => { + mockRealtime.onComplete?.({ + status: 'COMPLETED', + output: { downloadUrl: 'https://example.com/evidence.zip' }, + }); + }); + + expect(clickSpy).toHaveBeenCalledTimes(1); + expect(mockToastSuccess).toHaveBeenCalledWith( + 'Evidence package downloaded successfully', + ); + expect(mockToastError).not.toHaveBeenCalled(); + }); + + it('falls back to metadata.downloadUrl when output has none', async () => { + const clickSpy = vi.spyOn(HTMLAnchorElement.prototype, 'click'); + render(); + await startExport(); + + act(() => { + mockRealtime.onComplete?.({ + status: 'COMPLETED', + output: null, + metadata: { downloadUrl: 'https://example.com/from-metadata.zip' }, + }); + }); + + expect(clickSpy).toHaveBeenCalledTimes(1); + expect(mockToastSuccess).toHaveBeenCalledTimes(1); + }); + + it('toasts an error when the run completes without any download URL', async () => { + render(); + await startExport(); + + act(() => { + mockRealtime.onComplete?.({ status: 'COMPLETED', output: null, metadata: {} }); + }); + + expect(mockToastError).toHaveBeenCalledWith( + 'Export completed but download link was not available.', + ); + expect(mockToastSuccess).not.toHaveBeenCalled(); + }); + + it('toasts a failure error when the run ends in a non-COMPLETED state', async () => { + render(); + await startExport(); + + act(() => { + mockRealtime.onComplete?.({ status: 'FAILED' }); + }); + + expect(mockToastError).toHaveBeenCalledWith( + 'Evidence export failed. Please try again.', + ); + expect(mockToastSuccess).not.toHaveBeenCalled(); + }); + + it('toasts a failure error when onComplete reports an error', async () => { + render(); + await startExport(); + + act(() => { + mockRealtime.onComplete?.( + { status: 'COMPLETED', output: { downloadUrl: 'https://example.com/x.zip' } }, + new Error('subscription failed'), + ); + }); + + expect(mockToastError).toHaveBeenCalledWith( + 'Evidence export failed. Please try again.', + ); + expect(mockToastSuccess).not.toHaveBeenCalled(); + }); +}); diff --git a/apps/app/src/app/(app)/[orgId]/auditor/(overview)/components/ExportEvidenceButton.tsx b/apps/app/src/app/(app)/[orgId]/auditor/(overview)/components/ExportEvidenceButton.tsx index 604d434031..43c3e20a0e 100644 --- a/apps/app/src/app/(app)/[orgId]/auditor/(overview)/components/ExportEvidenceButton.tsx +++ b/apps/app/src/app/(app)/[orgId]/auditor/(overview)/components/ExportEvidenceButton.tsx @@ -36,9 +36,6 @@ export function ExportEvidenceButton({ phase: 'idle', }); - const isRunning = - exportState.phase === 'triggering' || exportState.phase === 'running'; - const handleTrigger = async () => { setExportState({ phase: 'triggering' }); try { @@ -54,7 +51,23 @@ export function ExportEvidenceButton({ }; const handleComplete = useCallback( - (run: { output?: { downloadUrl?: string } | null; metadata?: Record }) => { + ( + run: { + status?: string; + output?: { downloadUrl?: string } | null; + metadata?: Record; + }, + err?: Error, + ) => { + // useRealtimeRun fires onComplete on any terminal state (and surfaces + // subscription errors via `err`), so treat anything that isn't a clean + // COMPLETED run as a failure. + if (err || (run.status && run.status !== 'COMPLETED')) { + toast.error('Evidence export failed. Please try again.'); + setExportState({ phase: 'idle' }); + return; + } + const downloadUrl = run.output?.downloadUrl ?? (run.metadata?.downloadUrl as string | undefined); @@ -77,18 +90,34 @@ export function ExportEvidenceButton({ [organizationName], ); - const handleError = useCallback(() => { - toast.error('Evidence export failed. Please try again.'); - setExportState({ phase: 'idle' }); - }, []); + // Subscribe to the run from the parent (not from the progress UI inside the + // sheet) so the export keeps streaming — and still auto-downloads on + // completion — even after the user dismisses the sheet, which the copy below + // explicitly invites them to do. + const { run } = useRealtimeRun( + exportState.phase === 'running' ? exportState.runId : '', + { + accessToken: + exportState.phase === 'running' ? exportState.accessToken : undefined, + enabled: exportState.phase === 'running', + onComplete: handleComplete, + }, + ); + + const meta = run?.metadata as + | { + status?: string; + progress?: number; + tasksCompleted?: number; + tasksTotal?: number; + } + | undefined; return ( <> - { - if (!isRunning) setIsOpen(open); - }}> + Export All Evidence @@ -97,10 +126,10 @@ export function ExportEvidenceButton({ {exportState.phase === 'running' ? ( ) : ( <> @@ -150,37 +179,16 @@ export function ExportEvidenceButton({ } function ExportProgress({ - runId, - accessToken, - onComplete, - onError, + status, + progress, + tasksCompleted, + tasksTotal, }: { - runId: string; - accessToken: string; - onComplete: (run: { output?: { downloadUrl?: string } | null; metadata?: Record }) => void; - onError: () => void; + status: string; + progress: number; + tasksCompleted: number; + tasksTotal: number; }) { - const { run } = useRealtimeRun(runId, { - accessToken, - enabled: true, - onComplete: (run) => onComplete(run), - onError: () => onError(), - }); - - const meta = run?.metadata as - | { - status?: string; - progress?: number; - tasksCompleted?: number; - tasksTotal?: number; - } - | undefined; - - const progress = meta?.progress ?? 0; - const status = meta?.status ?? 'starting'; - const tasksCompleted = meta?.tasksCompleted ?? 0; - const tasksTotal = meta?.tasksTotal ?? 0; - const statusLabel = status === 'starting' ? 'Starting export...' diff --git a/apps/app/src/app/(app)/[orgId]/tasks/components/TasksPageClient.test.tsx b/apps/app/src/app/(app)/[orgId]/tasks/components/TasksPageClient.test.tsx index 6231d35e49..d0eca58a66 100644 --- a/apps/app/src/app/(app)/[orgId]/tasks/components/TasksPageClient.test.tsx +++ b/apps/app/src/app/(app)/[orgId]/tasks/components/TasksPageClient.test.tsx @@ -1,5 +1,5 @@ -import { render, screen } from '@testing-library/react'; -import { beforeEach, describe, expect, it, vi } from 'vitest'; +import { act, fireEvent, render, screen, waitFor } from '@testing-library/react'; +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { setMockPermissions, mockHasPermission, @@ -35,9 +35,43 @@ vi.mock('./TaskList', () => ({ TaskList: () =>
, })); -// Mock evidence download +// Mock evidence download (bulk export is triggered as a background job) +const mockTriggerBulkEvidenceExport = vi.fn(); vi.mock('@/lib/evidence-download', () => ({ - downloadAllEvidenceZip: vi.fn(), + triggerBulkEvidenceExport: (args: { includeJson?: boolean }) => + mockTriggerBulkEvidenceExport(args), +})); + +// Capture useRealtimeRun's onComplete so tests can simulate a finished run. +const mockRealtime: { + onComplete?: ( + run: { + status?: string; + output?: { downloadUrl?: string } | null; + metadata?: Record; + }, + err?: Error, + ) => void; +} = {}; +vi.mock('@trigger.dev/react-hooks', () => ({ + useRealtimeRun: ( + _runId: string, + options: { onComplete?: typeof mockRealtime.onComplete }, + ) => { + mockRealtime.onComplete = options.onComplete; + return { run: undefined }; + }, +})); + +const mockToastSuccess = vi.fn(); +const mockToastError = vi.fn(); +const mockToastInfo = vi.fn(); +vi.mock('sonner', () => ({ + toast: { + success: (msg: string) => mockToastSuccess(msg), + error: (msg: string) => mockToastError(msg), + info: (msg: string) => mockToastInfo(msg), + }, })); // Mock UpdateOrganizationEvidenceApproval @@ -164,3 +198,74 @@ describe('TasksPageClient permission gating', () => { expect(screen.getByTestId('task-list')).toBeInTheDocument(); }); }); + +describe('TasksPageClient evidence export', () => { + const exportProps = { ...defaultProps, hasEvidenceExportAccess: true }; + + beforeEach(() => { + vi.clearAllMocks(); + mockRealtime.onComplete = undefined; + setMockPermissions({}); + mockTriggerBulkEvidenceExport.mockResolvedValue({ + runId: 'run_1', + publicAccessToken: 'tok_1', + }); + vi.spyOn(HTMLAnchorElement.prototype, 'click').mockImplementation(() => {}); + }); + + afterEach(() => { + vi.restoreAllMocks(); + }); + + async function startExport() { + fireEvent.click(screen.getByRole('button', { name: 'Export' })); + await waitFor(() => expect(mockToastInfo).toHaveBeenCalled()); + } + + it('downloads and toasts success when the run completes with a URL', async () => { + const clickSpy = vi.spyOn(HTMLAnchorElement.prototype, 'click'); + render(); + await startExport(); + + act(() => { + mockRealtime.onComplete?.({ + status: 'COMPLETED', + output: { downloadUrl: 'https://example.com/e.zip' }, + }); + }); + + expect(clickSpy).toHaveBeenCalledTimes(1); + expect(mockToastSuccess).toHaveBeenCalledWith( + 'Evidence package downloaded successfully', + ); + expect(mockToastError).not.toHaveBeenCalled(); + }); + + it('toasts an error when the run completes without a download URL', async () => { + render(); + await startExport(); + + act(() => { + mockRealtime.onComplete?.({ status: 'COMPLETED', output: null, metadata: {} }); + }); + + expect(mockToastError).toHaveBeenCalledWith( + 'Export completed but download link was not available.', + ); + expect(mockToastSuccess).not.toHaveBeenCalled(); + }); + + it('toasts a failure error when the run does not complete cleanly', async () => { + render(); + await startExport(); + + act(() => { + mockRealtime.onComplete?.({ status: 'FAILED' }); + }); + + expect(mockToastError).toHaveBeenCalledWith( + 'Evidence export failed. Please try again.', + ); + expect(mockToastSuccess).not.toHaveBeenCalled(); + }); +}); diff --git a/apps/app/src/app/(app)/[orgId]/tasks/components/TasksPageClient.tsx b/apps/app/src/app/(app)/[orgId]/tasks/components/TasksPageClient.tsx index 6550c92f88..4fea4d83cf 100644 --- a/apps/app/src/app/(app)/[orgId]/tasks/components/TasksPageClient.tsx +++ b/apps/app/src/app/(app)/[orgId]/tasks/components/TasksPageClient.tsx @@ -78,10 +78,20 @@ export function TasksPageClient({ } | null>(null); const [isTriggering, setIsTriggering] = useState(false); - const { run: realtimeRun } = useRealtimeRun(exportRun?.runId ?? '', { + // Subscribe for the onComplete side effect (download / error toast); the + // returned run isn't rendered here, so we don't destructure it. + useRealtimeRun(exportRun?.runId ?? '', { accessToken: exportRun?.accessToken, enabled: !!exportRun, - onComplete: (run) => { + onComplete: (run, err) => { + // useRealtimeRun fires onComplete on any terminal state (and surfaces + // subscription errors via `err`), so treat anything that isn't a clean + // COMPLETED run as a failure. + if (err || run.status !== 'COMPLETED') { + toast.error('Evidence export failed. Please try again.'); + setExportRun(null); + return; + } const downloadUrl = run.output?.downloadUrl ?? (run.metadata?.downloadUrl as string | undefined); @@ -93,14 +103,12 @@ export function TasksPageClient({ link.click(); document.body.removeChild(link); toast.success('Evidence package downloaded successfully'); + } else { + toast.error('Export completed but download link was not available.'); } setExportRun(null); setIsPopoverOpen(false); }, - onError: () => { - toast.error('Evidence export failed. Please try again.'); - setExportRun(null); - }, }); const isDownloadingAll = isTriggering || !!exportRun;