Merge pull request #2567 from trycompai/fix/remediation-preview-trigger-dev

fix(cloud-tests): move remediation preview to Trigger.dev

Author: tofikwest

apps/app/src/app/(app)/[orgId]/cloud-tests/actions/single-fix.ts

@@ -4,6 +4,50 @@ import { auth as triggerAuth, tasks } from '@trigger.dev/sdk';
 import { auth } from '@/utils/auth';
 import { headers } from 'next/headers';
 
+interface PreviewInput {
+  connectionId: string;
+  checkResultId: string;
+  remediationKey: string;
+  cachedPermissions?: string[];
+}
+
+export async function startPreview(
+  input: PreviewInput,
+): Promise<{ data?: { runId: string; accessToken: string }; error?: string }> {
+  try {
+    const session = await auth.api.getSession({
+      headers: await headers(),
+    });
+
+    if (!session?.user?.id) {
+      return { error: 'Unauthorized' };
+    }
+
+    const organizationId = session.session?.activeOrganizationId;
+    if (!organizationId) {
+      return { error: 'No active organization' };
+    }
+
+    const handle = await tasks.trigger('remediate-preview', {
+      connectionId: input.connectionId,
+      organizationId,
+      checkResultId: input.checkResultId,
+      remediationKey: input.remediationKey,
+      userId: session.user.id,
+      cachedPermissions: input.cachedPermissions,
+    });
+
+    const accessToken = await triggerAuth.createPublicToken({
+      scopes: { read: { runs: [handle.id] } },
+    });
+
+    return { data: { runId: handle.id, accessToken } };
+  } catch (err) {
+    console.error('Failed to start preview:', err);
+    return { error: err instanceof Error ? err.message : 'Failed to load preview' };
+  }
+}
+
 interface SingleFixInput {
   connectionId: string;
   checkResultId: string;

apps/app/src/app/(app)/[orgId]/cloud-tests/components/RemediationDialog.tsx

@@ -1,6 +1,5 @@
 'use client';
 
-import { useApi } from '@/hooks/use-api';
 import { Badge } from '@trycompai/ui/badge';
 import { Button } from '@trycompai/ui/button';
 import {
@@ -14,10 +13,16 @@ import { useRealtimeRun } from '@trigger.dev/react-hooks';
 import { AlertTriangle, ListOrdered, Loader2, RotateCcw } from 'lucide-react';
 import { useCallback, useEffect, useRef, useState } from 'react';
 import { toast } from 'sonner';
-import { startSingleFix } from '../actions/single-fix';
+import { startPreview, startSingleFix } from '../actions/single-fix';
 import { AcknowledgmentPanel } from './AcknowledgmentPanel';
 import { PermissionErrorPanel } from './PermissionErrorPanel';
 
+interface PreviewProgress {
+  phase: 'analyzing' | 'complete' | 'failed';
+  error?: string;
+  preview?: PreviewData;
+}
+
 interface SingleFixProgress {
   phase: 'executing' | 'success' | 'failed' | 'needs_permissions';
   error?: string;
@@ -267,7 +272,6 @@ export function RemediationDialog({
   description,
   onComplete,
 }: RemediationDialogProps) {
-  const api = useApi();
   const [preview, setPreview] = useState<PreviewData | null>(null);
   const [isLoadingPreview, setIsLoadingPreview] = useState(false);
   const [isExecuting, setIsExecuting] = useState(false);
@@ -277,18 +281,53 @@ export function RemediationDialog({
   const [permissionError, setPermissionError] = useState<{ missingActions: string[]; fixScript?: string } | null>(null);
   const [acknowledgment, setAcknowledgment] = useState<string | null>(null);
 
-  // Trigger.dev state for async execution
-  const [runId, setRunId] = useState<string | null>(null);
-  const [triggerAccessToken, setTriggerAccessToken] = useState<string | null>(null);
+  // Trigger.dev state for preview (async)
+  const [previewRunId, setPreviewRunId] = useState<string | null>(null);
+  const [previewAccessToken, setPreviewAccessToken] = useState<string | null>(null);
 
-  const { run } = useRealtimeRun(runId ?? '', {
-    accessToken: triggerAccessToken ?? undefined,
-    enabled: Boolean(runId && triggerAccessToken),
+  // Trigger.dev state for execute (async)
+  const [executeRunId, setExecuteRunId] = useState<string | null>(null);
+  const [executeAccessToken, setExecuteAccessToken] = useState<string | null>(null);
+
+  const { run: previewRun } = useRealtimeRun(previewRunId ?? '', {
+    accessToken: previewAccessToken ?? undefined,
+    enabled: Boolean(previewRunId && previewAccessToken),
   });
 
-  // Watch task progress and update dialog state
+  const { run: executeRun } = useRealtimeRun(executeRunId ?? '', {
+    accessToken: executeAccessToken ?? undefined,
+    enabled: Boolean(executeRunId && executeAccessToken),
+  });
+
+  // Ref to store permissions across rechecks (avoids stale closure in useCallback)
+  const permissionsRef = useRef<string[] | undefined>(undefined);
+
+  // Watch preview task progress
   useEffect(() => {
-    const progress = (run?.metadata as { progress?: SingleFixProgress } | undefined)?.progress;
+    const progress = (previewRun?.metadata as { progress?: PreviewProgress } | undefined)?.progress;
+    if (!progress || progress.phase === 'analyzing') return;
+
+    if (progress.phase === 'complete' && progress.preview) {
+      const previewData = progress.preview as unknown as PreviewData;
+      setPreview(previewData);
+      if (previewData.allRequiredPermissions) {
+        permissionsRef.current = previewData.allRequiredPermissions;
+      }
+      setIsLoadingPreview(false);
+      setPreviewRunId(null);
+      setPreviewAccessToken(null);
+    } else if (progress.phase === 'failed') {
+      setError(progress.error || 'Failed to load preview');
+      setIsLoadingPreview(false);
+      setPreviewRunId(null);
+      setPreviewAccessToken(null);
+    }
+  // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [previewRun?.metadata]);
+
+  // Watch execute task progress
+  useEffect(() => {
+    const progress = (executeRun?.metadata as { progress?: SingleFixProgress } | undefined)?.progress;
     if (!progress || progress.phase === 'executing') return;
 
     if (progress.phase === 'success') {
@@ -301,73 +340,61 @@ export function RemediationDialog({
       setTimeout(() => {
         onOpenChange(false);
         setSucceeded(false);
-        setRunId(null);
-        setTriggerAccessToken(null);
+        setExecuteRunId(null);
+        setExecuteAccessToken(null);
       }, 4000);
     } else if (progress.phase === 'failed') {
       setIsExecuting(false);
       setError(progress.error || 'Remediation failed');
-      setRunId(null);
-      setTriggerAccessToken(null);
+      setExecuteRunId(null);
+      setExecuteAccessToken(null);
     } else if (progress.phase === 'needs_permissions') {
       setIsExecuting(false);
       setError(progress.error || 'Missing permissions');
       if (progress.permissionError) {
         setPermissionError(progress.permissionError);
       }
-      setRunId(null);
-      setTriggerAccessToken(null);
+      setExecuteRunId(null);
+      setExecuteAccessToken(null);
     }
   // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, [run?.metadata]);
-
-  // Ref to store permissions across rechecks (avoids stale closure in useCallback)
-  const permissionsRef = useRef<string[] | undefined>(undefined);
+  }, [executeRun?.metadata]);
 
   const loadPreview = useCallback(async (recheck = false) => {
     setIsLoadingPreview(true);
     setError(null);
     try {
-      const response = await api.post(
-        '/v1/cloud-security/remediation/preview',
-        {
-          connectionId,
-          checkResultId,
-          remediationKey,
-          // On recheck, send the cached permissions so backend doesn't re-run AI
-          ...(recheck && permissionsRef.current && {
-            cachedPermissions: permissionsRef.current,
-          }),
-        },
-      );
-      if (response.error) {
-        setError(
-          typeof response.error === 'string'
-            ? response.error
-            : 'Failed to load preview',
-        );
+      const result = await startPreview({
+        connectionId,
+        checkResultId,
+        remediationKey,
+        ...(recheck && permissionsRef.current && {
+          cachedPermissions: permissionsRef.current,
+        }),
+      });
+      if (result.error || !result.data) {
+        setError(result.error || 'Failed to load preview');
+        setIsLoadingPreview(false);
         return;
       }
-      const previewData = response.data as PreviewData;
-      setPreview(previewData);
-      // Store permissions in ref so Recheck can access them without stale closure
-      if (previewData.allRequiredPermissions) {
-        permissionsRef.current = previewData.allRequiredPermissions;
-      }
+      // Task started — preview effect handles the rest
+      setPreviewRunId(result.data.runId);
+      setPreviewAccessToken(result.data.accessToken);
     } catch {
       setError('Failed to load preview');
-    } finally {
       setIsLoadingPreview(false);
     }
-  }, [api, connectionId, checkResultId, remediationKey]);
+  }, [connectionId, checkResultId, remediationKey]);
 
   useEffect(() => {
     if (!open) return;
     setError(null);
     setPermissionError(null);
     setAcknowledgment(null);
-    setRunId(null);
-    setTriggerAccessToken(null);
+    setPreviewRunId(null);
+    setPreviewAccessToken(null);
+    setExecuteRunId(null);
+    setExecuteAccessToken(null);
     setSucceeded(false);
 
     // Guided-only: skip API call, use local data
@@ -406,9 +433,9 @@ export function RemediationDialog({
         setIsExecuting(false);
         return;
       }
-      // Task started — useRealtimeRun effect handles the rest
-      setRunId(result.data.runId);
-      setTriggerAccessToken(result.data.accessToken);
+      // Task started — execute effect handles the rest
+      setExecuteRunId(result.data.runId);
+      setExecuteAccessToken(result.data.accessToken);
     } catch {
       setError('Failed to start fix');
       setIsExecuting(false);

apps/app/src/trigger/tasks/cloud-security/remediate-preview.ts

@@ -0,0 +1,82 @@
+import { logger, metadata, task } from '@trigger.dev/sdk';
+
+interface PreviewProgress {
+  phase: 'analyzing' | 'complete' | 'failed';
+  error?: string;
+  preview?: Record<string, unknown>;
+}
+
+const getApiBaseUrl = () =>
+  process.env.NEXT_PUBLIC_API_URL || process.env.API_BASE_URL || 'http://localhost:3333';
+
+function makeHeaders(organizationId: string, userId?: string): Record<string, string> {
+  return {
+    'Content-Type': 'application/json',
+    'x-service-token': process.env.SERVICE_TOKEN_TRIGGER!,
+    'x-organization-id': organizationId,
+    ...(userId && { 'x-user-id': userId }),
+  };
+}
+
+function sync(progress: PreviewProgress) {
+  metadata.set('progress', JSON.parse(JSON.stringify(progress)));
+}
+
+export const remediatePreview = task({
+  id: 'remediate-preview',
+  maxDuration: 60 * 3, // 3 minutes
+  retry: { maxAttempts: 1 },
+  run: async (payload: {
+    connectionId: string;
+    organizationId: string;
+    checkResultId: string;
+    remediationKey: string;
+    userId: string;
+    cachedPermissions?: string[];
+  }) => {
+    const { connectionId, organizationId, checkResultId, remediationKey, userId, cachedPermissions } = payload;
+
+    logger.info(`Preview: ${remediationKey} on ${checkResultId} (user: ${userId})`);
+
+    const progress: PreviewProgress = { phase: 'analyzing' };
+    sync(progress);
+
+    try {
+      const url = `${getApiBaseUrl()}/v1/cloud-security/remediation/preview`;
+      const resp = await fetch(url, {
+        method: 'POST',
+        headers: makeHeaders(organizationId, userId),
+        body: JSON.stringify({
+          connectionId,
+          checkResultId,
+          remediationKey,
+          ...(cachedPermissions && { cachedPermissions }),
+        }),
+      });
+
+      const json = await resp.json();
+
+      if (!resp.ok) {
+        const errorMsg = (json as { message?: string }).message ?? `HTTP ${resp.status}`;
+        progress.phase = 'failed';
+        progress.error = errorMsg;
+        sync(progress);
+        logger.error(`Preview failed: ${errorMsg}`);
+        return { success: false, error: errorMsg };
+      }
+
+      progress.phase = 'complete';
+      progress.preview = json as Record<string, unknown>;
+      sync(progress);
+      logger.info(`Preview complete for ${remediationKey}`);
+      return { success: true, preview: json };
+    } catch (err) {
+      const errorMsg = err instanceof Error ? err.message : String(err);
+      progress.phase = 'failed';
+      progress.error = errorMsg;
+      sync(progress);
+      logger.error(`Preview exception: ${errorMsg}`);
+      return { success: false, error: errorMsg };
+    }
+  },
+});