feat(cloud-tests): enhance legacy integration filtering and add support for multiple connections

tofikwest · tofikwest · commit a91f30f8f0bc · 2026-02-03T11:35:35.000-05:00
diff --git a/apps/app/src/app/(app)/[orgId]/cloud-tests/actions/connect-cloud.ts b/apps/app/src/app/(app)/[orgId]/cloud-tests/actions/connect-cloud.ts
@@ -5,7 +5,7 @@ import { getIntegrationHandler } from '@comp/integrations';
 import { db } from '@db';
 import { Prisma } from '@prisma/client';
 import { revalidatePath } from 'next/cache';
-import { cookies, headers } from 'next/headers';
+import { headers } from 'next/headers';
 import { z } from 'zod';
 import { authActionClient } from '../../../../../actions/safe-action';
 import { runTests } from './run-tests';
@@ -113,12 +113,9 @@ export const connectCloudAction = authActionClient
       });
 
       // Trigger immediate scan for only this new connection
+      // runTests now waits for completion before returning
       const runResult = await runTests(newIntegration.id);
 
-      if (runResult.success && runResult.publicAccessToken) {
-        (await cookies()).set('publicAccessToken', runResult.publicAccessToken);
-      }
-
       // Revalidate the path
       const headersList = await headers();
       let path = headersList.get('x-pathname') || headersList.get('referer') || '';
@@ -130,7 +127,6 @@ export const connectCloudAction = authActionClient
         trigger: runResult.success
           ? {
               taskId: runResult.taskId ?? undefined,
-              publicAccessToken: runResult.publicAccessToken ?? undefined,
             }
           : undefined,
         runErrors: runResult.success ? undefined : (runResult.errors ?? undefined),
diff --git a/apps/app/src/app/(app)/[orgId]/cloud-tests/actions/run-tests.ts b/apps/app/src/app/(app)/[orgId]/cloud-tests/actions/run-tests.ts
@@ -2,12 +2,15 @@
 
 import { runIntegrationTests } from '@/trigger/tasks/integration/run-integration-tests';
 import { auth } from '@/utils/auth';
-import { tasks } from '@trigger.dev/sdk';
+import { runs, tasks } from '@trigger.dev/sdk';
 import { revalidatePath } from 'next/cache';
 import { headers } from 'next/headers';
 
+const MAX_POLL_ATTEMPTS = 60; // Max 2 minutes (60 * 2 seconds)
+const POLL_INTERVAL_MS = 2000;
+
 /**
- * Run integration tests.
+ * Run integration tests and wait for completion.
  * @param integrationId - Optional. If provided, only run tests for this specific connection.
  *                        If not provided, run tests for all connections in the organization.
  */
@@ -32,29 +35,72 @@ export const runTests = async (integrationId?: string) => {
   }
 
   try {
+    // Trigger the task
     const handle = await tasks.trigger<typeof runIntegrationTests>('run-integration-tests', {
       organizationId: orgId,
       ...(integrationId ? { integrationId } : {}),
     });
 
-    const headersList = await headers();
-    let path = headersList.get('x-pathname') || headersList.get('referer') || '';
-    path = path.replace(/\/[a-z]{2}\//, '/');
+    // Poll for completion
+    let attempts = 0;
+    while (attempts < MAX_POLL_ATTEMPTS) {
+      const run = await runs.retrieve(handle.id);
+
+      // Check if the run is in a terminal state
+      if (run.isCompleted) {
+        const headersList = await headers();
+        let path = headersList.get('x-pathname') || headersList.get('referer') || '';
+        path = path.replace(/\/[a-z]{2}\//, '/');
+        revalidatePath(path);
+
+        if (run.isSuccess) {
+          const output = run.output as {
+            success?: boolean;
+            errors?: string[];
+            failedIntegrations?: Array<{ name: string; error: string }>;
+          } | null;
+
+          if (output?.success === false) {
+            return {
+              success: false,
+              errors: output.errors || ['Scan completed with errors'],
+              taskId: run.id,
+            };
+          }
 
-    revalidatePath(path);
+          return {
+            success: true,
+            errors: null,
+            taskId: run.id,
+          };
+        }
 
+        if (run.isFailed || run.isCancelled) {
+          return {
+            success: false,
+            errors: ['Task failed or was canceled'],
+            taskId: run.id,
+          };
+        }
+      }
+
+      // Wait before polling again
+      await new Promise((resolve) => setTimeout(resolve, POLL_INTERVAL_MS));
+      attempts++;
+    }
+
+    // Timeout - task is taking too long
     return {
-      success: true,
-      errors: null,
+      success: false,
+      errors: ['Scan is taking longer than expected. Check the status in Trigger.dev dashboard.'],
       taskId: handle.id,
-      publicAccessToken: handle.publicAccessToken,
     };
   } catch (error) {
-    console.error('Error triggering integration tests:', error);
+    console.error('Error running integration tests:', error);
 
     return {
       success: false,
-      errors: [error instanceof Error ? error.message : 'Failed to trigger integration tests'],
+      errors: [error instanceof Error ? error.message : 'Failed to run integration tests'],
     };
   }
 };
diff --git a/apps/app/src/app/(app)/[orgId]/cloud-tests/components/TestsLayout.tsx b/apps/app/src/app/(app)/[orgId]/cloud-tests/components/TestsLayout.tsx
@@ -83,7 +83,9 @@ export function TestsLayout({ initialFindings, initialProviders, orgId }: TestsL
     },
     {
       fallbackData: initialProviders,
+      refreshInterval: 10000, // Refresh providers every 10 seconds
       revalidateOnFocus: true,
+      revalidateOnMount: true, // Always revalidate on mount to get fresh data
     },
   );
 
@@ -134,16 +136,22 @@ export function TestsLayout({ initialFindings, initialProviders, orgId }: TestsL
     }
 
     setIsScanning(true);
+    const startTime = Date.now();
     toast.message(`Starting ${targetProvider.name} security scan...`);
 
     try {
       if (targetProvider.isLegacy) {
         // Run legacy check for this specific connection
+        // runTests now waits for completion (uses triggerAndPoll)
         const { runTests } = await import('../actions/run-tests');
-        // Pass the unique connection ID to only scan this specific connection
         const result = await runTests(targetProvider.id);
+
         if (!result.success) {
           console.error('Legacy scan error:', result.errors);
+          toast.error(
+            `Scan failed: ${result.errors?.join(', ') || 'Unknown error'}`,
+          );
+          return null;
         }
       } else {
         // Use dedicated cloud security endpoint
@@ -155,9 +163,11 @@ export function TestsLayout({ initialFindings, initialProviders, orgId }: TestsL
         }
       }
 
-      toast.success('Scan completed! Results updated.');
-      await mutateProviders(); // Refresh to get updated lastRunAt
-      await mutateFindings();
+      // Refresh data to get updated results
+      await Promise.all([mutateProviders(), mutateFindings()]);
+
+      const elapsed = Math.round((Date.now() - startTime) / 1000);
+      toast.success(`Scan completed in ${elapsed}s! Results updated.`);
       return 'completed';
     } catch (error) {
       console.error('Scan error:', error);
diff --git a/apps/app/src/app/(app)/[orgId]/cloud-tests/page.tsx b/apps/app/src/app/(app)/[orgId]/cloud-tests/page.tsx
@@ -107,6 +107,7 @@ export default async function CloudTestsPage({ params }: { params: Promise<{ org
     requiredVariables: string[];
     accountId?: string;
     regions?: string[];
+    supportsMultipleConnections?: boolean;
   };
 
   const newProviders: Provider[] = newConnections.map((conn) => {
@@ -117,6 +118,7 @@ export default async function CloudTestsPage({ params }: { params: Promise<{ org
     const regions = Array.isArray(metadata.regions)
       ? metadata.regions.filter((region): region is string => typeof region === 'string')
       : undefined;
+    const manifest = getManifest(conn.provider.slug);
 
     return {
       id: conn.id,
@@ -133,6 +135,7 @@ export default async function CloudTestsPage({ params }: { params: Promise<{ org
       requiredVariables: getRequiredVariables(conn.provider.slug),
       accountId,
       regions,
+      supportsMultipleConnections: manifest?.supportsMultipleConnections ?? false,
     };
   });
 
@@ -144,6 +147,7 @@ export default async function CloudTestsPage({ params }: { params: Promise<{ org
     const regions = Array.isArray(settings.regions)
       ? settings.regions.filter((region): region is string => typeof region === 'string')
       : undefined;
+    const manifest = getManifest(integration.integrationId);
 
     return {
       id: integration.id,
@@ -160,6 +164,7 @@ export default async function CloudTestsPage({ params }: { params: Promise<{ org
       requiredVariables: getRequiredVariables(integration.integrationId),
       accountId,
       regions,
+      supportsMultipleConnections: manifest?.supportsMultipleConnections ?? false,
     };
   });
 
@@ -231,9 +236,17 @@ export default async function CloudTestsPage({ params }: { params: Promise<{ org
 
   // ====================================================================
   // Fetch findings from OLD platform (IntegrationResult)
+  // Only show results from the most recent scan for each integration
   // ====================================================================
   const legacyIntegrationIds = activeLegacyIntegrations.map((i) => i.id);
 
+  // Create a map of integration ID to lastRunAt for filtering
+  const integrationLastRunMap = new Map(
+    activeLegacyIntegrations
+      .filter((i) => i.lastRunAt)
+      .map((i) => [i.id, i.lastRunAt!]),
+  );
+
   const legacyResults =
     legacyIntegrationIds.length > 0
       ? await db.integrationResult.findMany({
@@ -254,17 +267,33 @@ export default async function CloudTestsPage({ params }: { params: Promise<{ org
               select: {
                 integrationId: true,
                 id: true,
+                lastRunAt: true,
               },
             },
           },
           orderBy: {
             completedAt: 'desc',
           },
-          take: 500,
         })
       : [];
 
-  const legacyFindings = legacyResults.map((result) => ({
+  // Filter to only include results from the most recent scan
+  // Results are considered from the "latest scan" if they were completed
+  // within 5 minutes of the integration's lastRunAt
+  const SCAN_WINDOW_MS = 5 * 60 * 1000; // 5 minutes
+
+  const filteredLegacyResults = legacyResults.filter((result) => {
+    const lastRunAt = integrationLastRunMap.get(result.integration.id);
+    if (!lastRunAt || !result.completedAt) return false;
+
+    const lastRunTime = lastRunAt.getTime();
+    const completedTime = result.completedAt.getTime();
+
+    // Include if completed within the scan window of the last run
+    return Math.abs(completedTime - lastRunTime) <= SCAN_WINDOW_MS;
+  });
+
+  const legacyFindings = filteredLegacyResults.map((result) => ({
     id: result.id,
     title: result.title,
     description: result.description,
diff --git a/apps/app/src/app/api/cloud-tests/findings/route.ts b/apps/app/src/app/api/cloud-tests/findings/route.ts
@@ -136,9 +136,17 @@ export async function GET(request: NextRequest) {
 
     // ====================================================================
     // Fetch findings from OLD platform (IntegrationResult)
+    // Only show results from the most recent scan for each integration
     // ====================================================================
     const legacyIntegrationIds = activeLegacyIntegrations.map((i) => i.id);
 
+    // Create a map of integration ID to lastRunAt for filtering
+    const integrationLastRunMap = new Map(
+      activeLegacyIntegrations
+        .filter((i) => i.lastRunAt)
+        .map((i) => [i.id, i.lastRunAt!]),
+    );
+
     const legacyResults =
       legacyIntegrationIds.length > 0
         ? await db.integrationResult.findMany({
@@ -159,17 +167,33 @@ export async function GET(request: NextRequest) {
                 select: {
                   integrationId: true,
                   id: true,
+                  lastRunAt: true,
                 },
               },
             },
             orderBy: {
               completedAt: 'desc',
             },
-            take: 500,
           })
         : [];
 
-    const legacyFindings = legacyResults.map((result) => ({
+    // Filter to only include results from the most recent scan
+    // Results are considered from the "latest scan" if they were completed
+    // within 5 minutes of the integration's lastRunAt
+    const SCAN_WINDOW_MS = 5 * 60 * 1000; // 5 minutes
+
+    const filteredLegacyResults = legacyResults.filter((result) => {
+      const lastRunAt = integrationLastRunMap.get(result.integration.id);
+      if (!lastRunAt || !result.completedAt) return false;
+
+      const lastRunTime = lastRunAt.getTime();
+      const completedTime = result.completedAt.getTime();
+
+      // Include if completed within the scan window of the last run
+      return Math.abs(completedTime - lastRunTime) <= SCAN_WINDOW_MS;
+    });
+
+    const legacyFindings = filteredLegacyResults.map((result) => ({
       id: result.id,
       title: result.title,
       description: result.description,
