refactor(ui): merge connection test into account selection step

ui/actions/organizations/organizations.adapter.ts

@@ -1,4 +1,4 @@
-import { Building2, FolderTree, ShieldCheck } from "lucide-react";
+import { Box, Folder, FolderTree } from "lucide-react";
 
 import {
   APPLY_STATUS,
@@ -10,7 +10,8 @@ import { TreeDataItem } from "@/types/tree";
 /**
  * Transforms flat API discovery arrays into hierarchical TreeDataItem[] for TreeView.
  *
- * Structure: Root → OUs → Accounts (leaf nodes)
+ * Structure: OUs → nested OUs/Accounts (leaf nodes)
+ * Root nodes are used only internally for parent linking and are not rendered.
  * Accounts with apply_status === "blocked" are marked disabled.
  */
 export function buildOrgTreeData(result: DiscoveryResult): TreeDataItem[] {
@@ -32,7 +33,7 @@ export function buildOrgTreeData(result: DiscoveryResult): TreeDataItem[] {
     nodeMap.set(ou.id, {
       id: ou.id,
       name: ou.name,
-      icon: Building2,
+      icon: Folder,
       children: [],
     });
   }
@@ -45,7 +46,7 @@ export function buildOrgTreeData(result: DiscoveryResult): TreeDataItem[] {
     nodeMap.set(account.id, {
       id: account.id,
       name: `${account.id} — ${account.name}`,
-      icon: ShieldCheck,
+      icon: Box,
       disabled: isBlocked,
     });
   }
@@ -70,10 +71,11 @@ export function buildOrgTreeData(result: DiscoveryResult): TreeDataItem[] {
     }
   }
 
-  // 7. Return root-level nodes as the tree
-  return result.roots
-    .map((root) => nodeMap.get(root.id))
-    .filter((node): node is TreeDataItem => node !== undefined);
+  // 7. Return root children as top-level nodes (roots are not rendered).
+  return result.roots.flatMap((root) => {
+    const rootNode = nodeMap.get(root.id);
+    return rootNode?.children ?? [];
+  });
 }
 
 /**

ui/components/providers/organizations/org-account-selection.tsx

@@ -1,56 +1,213 @@
 "use client";
 
-import { useEffect } from "react";
+import { AlertTriangle } from "lucide-react";
+import { useEffect, useRef, useState } from "react";
 
+import { applyDiscovery } from "@/actions/organizations/organizations";
 import {
   buildAccountLookup,
   buildOrgTreeData,
+  getOuIdsForSelectedAccounts,
+  getSelectableAccountIds,
 } from "@/actions/organizations/organizations.adapter";
+import { checkConnectionProvider } from "@/actions/providers/providers";
+import { getTask } from "@/actions/task/tasks";
+import { AWSProviderBadge } from "@/components/icons/providers-badge";
 import {
   WIZARD_FOOTER_ACTION_TYPE,
   WizardFooterConfig,
 } from "@/components/providers/wizard/steps/footer-controls";
-import { Badge } from "@/components/shadcn";
+import { Alert, AlertDescription } from "@/components/shadcn/alert";
 import { TreeView } from "@/components/shadcn/tree-view";
+import { checkTaskStatus } from "@/lib";
 import { useOrgSetupStore } from "@/store/organizations/store";
+import {
+  CONNECTION_TEST_STATUS,
+  ConnectionTestStatus,
+} from "@/types/organizations";
+import { TREE_ITEM_STATUS, TreeDataItem } from "@/types/tree";
 
 import { OrgAccountTreeItem, TREE_ITEM_MODE } from "./org-account-tree-item";
 
 interface OrgAccountSelectionProps {
   onBack: () => void;
   onNext: () => void;
+  onSkip: () => void;
   onFooterChange: (config: WizardFooterConfig) => void;
 }
 
+interface SelectionState {
+  hasSelectableDescendants: boolean;
+  allSelectableDescendantsSelected: boolean;
+}
+
+function collectFullySelectedNodeIds(
+  node: TreeDataItem,
+  selectedAccountIdSet: Set<string>,
+  selectableAccountIdSet: Set<string>,
+  selectedNodeIds: Set<string>,
+): SelectionState {
+  if (selectableAccountIdSet.has(node.id)) {
+    return {
+      hasSelectableDescendants: true,
+      allSelectableDescendantsSelected: selectedAccountIdSet.has(node.id),
+    };
+  }
+
+  const children = node.children ?? [];
+  let hasSelectableDescendants = false;
+  let allSelectableDescendantsSelected = true;
+
+  for (const child of children) {
+    const childSelectionState = collectFullySelectedNodeIds(
+      child,
+      selectedAccountIdSet,
+      selectableAccountIdSet,
+      selectedNodeIds,
+    );
+
+    if (!childSelectionState.hasSelectableDescendants) {
+      continue;
+    }
+
+    hasSelectableDescendants = true;
+    allSelectableDescendantsSelected =
+      allSelectableDescendantsSelected &&
+      childSelectionState.allSelectableDescendantsSelected;
+  }
+
+  if (hasSelectableDescendants && allSelectableDescendantsSelected) {
+    selectedNodeIds.add(node.id);
+  }
+
+  return {
+    hasSelectableDescendants,
+    allSelectableDescendantsSelected,
+  };
+}
+
+function buildTreeSelectedIds(
+  treeData: TreeDataItem[],
+  selectedAccountIds: string[],
+  selectableAccountIdSet: Set<string>,
+): string[] {
+  const selectedAccountIdSet = new Set(selectedAccountIds);
+  const selectedNodeIds = new Set<string>();
+
+  for (const rootNode of treeData) {
+    collectFullySelectedNodeIds(
+      rootNode,
+      selectedAccountIdSet,
+      selectableAccountIdSet,
+      selectedNodeIds,
+    );
+  }
+
+  return [...selectedAccountIds, ...Array.from(selectedNodeIds)];
+}
+
+function buildTreeWithConnectionState(
+  nodes: TreeDataItem[],
+  selectedAccountIdsSet: Set<string>,
+  accountToProviderMap: Map<string, string>,
+  connectionResults: Record<string, ConnectionTestStatus>,
+  connectionErrors: Record<string, string>,
+  showPendingState: boolean,
+): TreeDataItem[] {
+  return nodes.map((node) => {
+    const children = node.children
+      ? buildTreeWithConnectionState(
+          node.children,
+          selectedAccountIdsSet,
+          accountToProviderMap,
+          connectionResults,
+          connectionErrors,
+          showPendingState,
+        )
+      : undefined;
+
+    let isLoading = node.isLoading;
+    let status = node.status;
+    let errorMessage = node.errorMessage;
+
+    if (selectedAccountIdsSet.has(node.id)) {
+      const providerId = accountToProviderMap.get(node.id);
+      const connectionStatus = providerId
+        ? connectionResults[providerId]
+        : undefined;
+
+      if (connectionStatus === CONNECTION_TEST_STATUS.SUCCESS) {
+        isLoading = false;
+        status = TREE_ITEM_STATUS.SUCCESS;
+        errorMessage = undefined;
+      } else if (connectionStatus === CONNECTION_TEST_STATUS.ERROR) {
+        isLoading = false;
+        status = TREE_ITEM_STATUS.ERROR;
+        errorMessage =
+          (providerId && connectionErrors[providerId]) || "Connection failed.";
+      } else if (
+        showPendingState ||
+        connectionStatus === CONNECTION_TEST_STATUS.PENDING
+      ) {
+        isLoading = true;
+        status = undefined;
+        errorMessage = undefined;
+      }
+    }
+
+    return {
+      ...node,
+      children,
+      isLoading,
+      status,
+      errorMessage,
+    };
+  });
+}
+
+function extractErrorMessage(response: unknown, fallback: string): string {
+  if (!response || typeof response !== "object") {
+    return fallback;
+  }
+
+  const responseRecord = response as {
+    error?: string;
+    errors?: Array<{ detail?: string }>;
+  };
+  const detailedError = responseRecord.errors?.[0]?.detail;
+  return detailedError || responseRecord.error || fallback;
+}
+
 export function OrgAccountSelection({
   onBack,
   onNext,
+  onSkip,
   onFooterChange,
 }: OrgAccountSelectionProps) {
   const {
-    organizationName,
+    organizationId,
     organizationExternalId,
+    discoveryId,
     discoveryResult,
     selectedAccountIds,
     accountAliases,
+    createdProviderIds,
+    connectionResults,
+    connectionErrors,
     setSelectedAccountIds,
     setAccountAlias,
+    setCreatedProviderIds,
+    setConnectionError,
+    setConnectionResult,
   } = useOrgSetupStore();
-
-  const selectedCount = selectedAccountIds.length;
-
-  useEffect(() => {
-    onFooterChange({
-      showBack: true,
-      backLabel: "Back",
-      onBack,
-      showAction: true,
-      actionLabel: "Next",
-      actionDisabled: selectedCount === 0,
-      actionType: WIZARD_FOOTER_ACTION_TYPE.BUTTON,
-      onAction: onNext,
-    });
-  }, [onBack, onFooterChange, onNext, selectedCount]);
+  const [isTestingView, setIsTestingView] = useState(false);
+  const [isApplying, setIsApplying] = useState(false);
+  const [isTesting, setIsTesting] = useState(false);
+  const [applyError, setApplyError] = useState<string | null>(null);
+  const [accountToProviderMap, setAccountToProviderMap] = useState<
+    Map<string, string>
+  >(new Map());
+  const hasAppliedRef = useRef(false);
 
   if (!discoveryResult) {
     return (
@@ -62,35 +219,320 @@ export function OrgAccountSelection({
 
   const treeData = buildOrgTreeData(discoveryResult);
   const accountLookup = buildAccountLookup(discoveryResult);
+  const selectableAccountIds = getSelectableAccountIds(discoveryResult);
+  const selectableAccountIdSet = new Set(selectableAccountIds);
+  const sanitizedSelectedAccountIds = selectedAccountIds.filter((id) =>
+    selectableAccountIdSet.has(id),
+  );
+  const selectedIdsForTree = buildTreeSelectedIds(
+    treeData,
+    sanitizedSelectedAccountIds,
+    selectableAccountIdSet,
+  );
+  const selectedAccountIdSet = new Set(sanitizedSelectedAccountIds);
+  const selectedCount = sanitizedSelectedAccountIds.length;
+  const totalAccounts = selectableAccountIds.length;
+  const hasConnectionErrors = Object.values(connectionResults).some(
+    (status) => status === CONNECTION_TEST_STATUS.ERROR,
+  );
+  const showHeaderHelperText = !isTestingView || isApplying || isTesting;
+  const treeDataWithConnectionState = isTestingView
+    ? buildTreeWithConnectionState(
+        treeData,
+        selectedAccountIdSet,
+        accountToProviderMap,
+        connectionResults,
+        connectionErrors,
+        isApplying || isTesting,
+      )
+    : treeData;
 
-  const totalAccounts = discoveryResult.accounts.length;
+  useEffect(() => {
+    if (sanitizedSelectedAccountIds.length === selectedAccountIds.length) {
+      return;
+    }
+
+    setSelectedAccountIds(sanitizedSelectedAccountIds);
+  }, [sanitizedSelectedAccountIds, selectedAccountIds, setSelectedAccountIds]);
+
+  const testAllConnections = async (providerIds: string[]) => {
+    setIsTesting(true);
+
+    for (const id of providerIds) {
+      setConnectionResult(id, CONNECTION_TEST_STATUS.PENDING);
+      setConnectionError(id, null);
+    }
+
+    const testPromises = providerIds.map(async (providerId) => {
+      try {
+        const formData = new FormData();
+        formData.set("providerId", providerId);
+
+        const checkResult = await checkConnectionProvider(formData);
+        if (checkResult?.error || checkResult?.errors?.length) {
+          setConnectionResult(providerId, CONNECTION_TEST_STATUS.ERROR);
+          setConnectionError(
+            providerId,
+            extractErrorMessage(checkResult, "Connection test failed."),
+          );
+          return;
+        }
+
+        const taskId = checkResult?.data?.id;
+        if (!taskId) {
+          setConnectionResult(providerId, CONNECTION_TEST_STATUS.SUCCESS);
+          setConnectionError(providerId, null);
+          return;
+        }
+
+        const taskResult = await checkTaskStatus(taskId);
+        if (!taskResult.completed) {
+          setConnectionResult(providerId, CONNECTION_TEST_STATUS.ERROR);
+          setConnectionError(
+            providerId,
+            taskResult.error || "Connection test task failed.",
+          );
+          return;
+        }
+
+        const task = await getTask(taskId);
+        if (task?.error) {
+          setConnectionResult(providerId, CONNECTION_TEST_STATUS.ERROR);
+          setConnectionError(providerId, task.error);
+          return;
+        }
+
+        const resultPayload = task?.data?.attributes?.result as
+          | { connected?: boolean; error?: string }
+          | undefined;
+        const isConnected = resultPayload?.connected ?? taskResult.completed;
+
+        setConnectionResult(
+          providerId,
+          isConnected
+            ? CONNECTION_TEST_STATUS.SUCCESS
+            : CONNECTION_TEST_STATUS.ERROR,
+        );
+        setConnectionError(
+          providerId,
+          isConnected
+            ? null
+            : resultPayload?.error || "Connection failed for this account.",
+        );
+      } catch {
+        setConnectionResult(providerId, CONNECTION_TEST_STATUS.ERROR);
+        setConnectionError(
+          providerId,
+          "Unexpected error during connection test.",
+        );
+      }
+    });
+
+    await Promise.allSettled(testPromises);
+    setIsTesting(false);
+
+    const latestResults = useOrgSetupStore.getState().connectionResults;
+    const allPassed =
+      providerIds.length > 0 &&
+      providerIds.every(
+        (providerId) =>
+          latestResults[providerId] === CONNECTION_TEST_STATUS.SUCCESS,
+      );
+
+    if (allPassed) {
+      onNext();
+    }
+  };
+
+  const handleApplyAndTest = async () => {
+    if (!organizationId || !discoveryId) {
+      return;
+    }
+
+    setApplyError(null);
+    setIsApplying(true);
+
+    const accounts = sanitizedSelectedAccountIds.map((id) => ({
+      id,
+      ...(accountAliases[id] ? { alias: accountAliases[id] } : {}),
+    }));
+    const ouIds = getOuIdsForSelectedAccounts(
+      discoveryResult,
+      sanitizedSelectedAccountIds,
+    );
+    const organizationalUnits = ouIds.map((id) => ({ id }));
+
+    const result = await applyDiscovery(
+      organizationId,
+      discoveryId,
+      accounts,
+      organizationalUnits,
+    );
+
+    if (result?.error) {
+      setApplyError(result.error);
+      setIsApplying(false);
+      hasAppliedRef.current = false;
+      return;
+    }
+
+    const providerIds: string[] =
+      result.data?.relationships?.providers?.data?.map(
+        (provider: { id: string }) => provider.id,
+      ) ?? [];
+
+    setCreatedProviderIds(providerIds);
+
+    const mapping = new Map<string, string>();
+    sanitizedSelectedAccountIds.forEach((accountId, index) => {
+      if (providerIds[index]) {
+        mapping.set(accountId, providerIds[index]);
+      }
+    });
+    setAccountToProviderMap(mapping);
+    setIsApplying(false);
+
+    await testAllConnections(providerIds);
+  };
+
+  const handleStartTesting = () => {
+    setIsTestingView(true);
+
+    if (applyError) {
+      setApplyError(null);
+      hasAppliedRef.current = false;
+    }
+
+    if (!hasAppliedRef.current) {
+      hasAppliedRef.current = true;
+      void handleApplyAndTest();
+      return;
+    }
+
+    const failedProviderIds = createdProviderIds.filter(
+      (providerId) =>
+        connectionResults[providerId] === CONNECTION_TEST_STATUS.ERROR,
+    );
+    const providerIdsToTest =
+      failedProviderIds.length > 0 ? failedProviderIds : createdProviderIds;
+    void testAllConnections(providerIdsToTest);
+  };
+
+  useEffect(() => {
+    if (!isTestingView) {
+      onFooterChange({
+        showBack: true,
+        backLabel: "Back",
+        onBack,
+        showSecondaryAction: false,
+        secondaryActionLabel: "",
+        secondaryActionVariant: "outline",
+        secondaryActionType: WIZARD_FOOTER_ACTION_TYPE.BUTTON,
+        showAction: true,
+        actionLabel: "Test Connections",
+        actionDisabled: selectedCount === 0,
+        actionType: WIZARD_FOOTER_ACTION_TYPE.BUTTON,
+        onAction: handleStartTesting,
+      });
+      return;
+    }
+
+    const canRetry = hasConnectionErrors || Boolean(applyError);
+
+    onFooterChange({
+      showBack: true,
+      backLabel: "Back",
+      backDisabled: isApplying || isTesting,
+      onBack: () => setIsTestingView(false),
+      showSecondaryAction: true,
+      secondaryActionLabel: "Skip Connection Validation",
+      secondaryActionDisabled: false,
+      secondaryActionVariant: "link",
+      secondaryActionType: WIZARD_FOOTER_ACTION_TYPE.BUTTON,
+      onSecondaryAction: onSkip,
+      showAction: isApplying || isTesting || canRetry,
+      actionLabel: "Test Connections",
+      actionDisabled: isApplying || isTesting || !canRetry,
+      actionType: WIZARD_FOOTER_ACTION_TYPE.BUTTON,
+      onAction: canRetry ? handleStartTesting : undefined,
+    });
+  }, [
+    applyError,
+    connectionErrors,
+    hasConnectionErrors,
+    isApplying,
+    isTesting,
+    isTestingView,
+    onBack,
+    onFooterChange,
+    onSkip,
+    selectedCount,
+    setConnectionError,
+  ]);
+
+  const handleTreeSelectionChange = (ids: string[]) => {
+    setSelectedAccountIds(ids.filter((id) => selectableAccountIdSet.has(id)));
+  };
 
   return (
-    <div className="flex flex-col gap-5">
+    <div className="flex min-h-0 flex-1 flex-col gap-5">
       {/* Header */}
-      <div className="flex flex-col gap-2">
-        <div className="flex items-center gap-2">
-          <Badge variant="outline">AWS</Badge>
-          <span className="text-sm font-medium">{organizationName}</span>
-          {organizationExternalId && (
-            <Badge variant="secondary">{organizationExternalId}</Badge>
-          )}
+      <div className="flex flex-col gap-3">
+        <div className="flex items-center gap-4">
+          <AWSProviderBadge size={32} />
+          <h3 className="text-base font-semibold">My Organization</h3>
         </div>
-        <p className="text-muted-foreground text-sm">
-          Select the accounts you want to connect. {selectedCount} of{" "}
-          {totalAccounts} accounts selected.
-        </p>
+
+        <div className="ml-12 flex items-center gap-3">
+          <span className="text-text-neutral-tertiary text-xs">UID:</span>
+          <div className="bg-bg-neutral-tertiary border-border-input-primary inline-flex h-10 items-center rounded-full border px-4">
+            <span className="text-xs font-medium">
+              {organizationExternalId || "N/A"}
+            </span>
+          </div>
+        </div>
+
+        {showHeaderHelperText && (
+          <p className="text-muted-foreground text-sm">
+            {isTestingView
+              ? "Testing account connections..."
+              : "Confirm all accounts under this Organization you want to add to Prowler."}{" "}
+            {!isTestingView &&
+              `${selectedCount} of ${totalAccounts} accounts selected.`}
+          </p>
+        )}
       </div>
 
+      {isTestingView && applyError && (
+        <Alert variant="error">
+          <AlertTriangle />
+          <AlertDescription className="text-text-error-primary">
+            {applyError}
+          </AlertDescription>
+        </Alert>
+      )}
+
+      {isTestingView && hasConnectionErrors && !isTesting && (
+        <Alert variant="error">
+          <AlertTriangle />
+          <AlertDescription className="text-text-error-primary">
+            There was a problem connecting to some accounts. Hover each account
+            to check the error.
+          </AlertDescription>
+        </Alert>
+      )}
+
       {/* Tree */}
-      <div className="max-h-80 overflow-y-auto rounded-md border p-2">
+      <div className="border-border-neutral-secondary min-h-0 flex-1 overflow-y-auto rounded-md border p-2">
         <TreeView
-          data={treeData}
+          data={treeDataWithConnectionState}
           showCheckboxes
           enableSelectChildren
           expandAll
-          selectedIds={selectedAccountIds}
-          onSelectionChange={setSelectedAccountIds}
+          selectedIds={selectedIdsForTree}
+          onSelectionChange={
+            isTestingView ? () => {} : handleTreeSelectionChange
+          }
           renderItem={(params) => (
             <OrgAccountTreeItem
               params={params}

ui/components/providers/organizations/org-account-tree-item.tsx

@@ -1,6 +1,6 @@
 "use client";
 
-import { AlertCircle, Check, Loader2 } from "lucide-react";
+import { AlertCircle } from "lucide-react";
 
 import { Input } from "@/components/shadcn/input/input";
 import {
@@ -11,15 +11,12 @@ import {
 import { cn } from "@/lib/utils";
 import {
   APPLY_STATUS,
-  CONNECTION_TEST_STATUS,
-  ConnectionTestStatus,
   DiscoveredAccount,
 } from "@/types/organizations";
 import { TreeRenderItemParams } from "@/types/tree";
 
 const TREE_ITEM_MODE = {
   SELECTION: "selection",
-  TESTING: "testing",
 } as const;
 
 type TreeItemMode = (typeof TREE_ITEM_MODE)[keyof typeof TREE_ITEM_MODE];
@@ -30,8 +27,6 @@ interface OrgAccountTreeItemProps {
   accountLookup: Map<string, DiscoveredAccount>;
   aliases: Record<string, string>;
   onAliasChange?: (accountId: string, alias: string) => void;
-  connectionResults?: Record<string, ConnectionTestStatus>;
-  accountToProviderMap?: Map<string, string>;
 }
 
 export function OrgAccountTreeItem({
@@ -40,14 +35,46 @@ export function OrgAccountTreeItem({
   accountLookup,
   aliases,
   onAliasChange,
-  connectionResults,
-  accountToProviderMap,
 }: OrgAccountTreeItemProps) {
   const { item, isLeaf } = params;
   const account = accountLookup.get(item.id);
+  const isOuNode = item.id.startsWith("ou-");
+  const ItemIcon = item.icon;
+  const idColumnClass = "w-44 shrink-0";
+  const aliasInputClass = "h-9 w-full max-w-64 text-sm";
 
-  // Non-leaf nodes (roots/OUs) just render their name
-  if (!isLeaf || !account) {
+  // OU nodes: show OU id + alias/name (input in selection mode).
+  if (!account && isOuNode) {
+    const ouDisplayName = aliases[item.id] ?? item.name;
+    const isSelectionMode = mode === TREE_ITEM_MODE.SELECTION && onAliasChange;
+
+    return (
+      <div className="flex flex-1 items-center gap-3">
+        <div className={`${idColumnClass} flex items-center gap-2`}>
+          {ItemIcon && <ItemIcon className="text-muted-foreground size-4 shrink-0" />}
+          <span className="text-sm">{item.id}</span>
+        </div>
+        <div className="min-w-0 flex-1">
+          {isSelectionMode ? (
+            <Input
+              className={aliasInputClass}
+              placeholder="Name (optional)"
+              value={ouDisplayName}
+              onChange={(e) => onAliasChange(item.id, e.target.value)}
+              onClick={(e) => e.stopPropagation()}
+            />
+          ) : (
+            <span className="text-muted-foreground line-clamp-1 text-xs">
+              {ouDisplayName}
+            </span>
+          )}
+        </div>
+      </div>
+    );
+  }
+
+  // Any remaining non-account node (unexpected fallback).
+  if (!account || !isLeaf) {
     return <span className="text-sm font-medium">{item.name}</span>;
   }
 
@@ -56,36 +83,30 @@ export function OrgAccountTreeItem({
 
   return (
     <div className="flex flex-1 items-center gap-3">
-      {/* Status icon for testing mode */}
-      {mode === TREE_ITEM_MODE.TESTING && (
-        <TestStatusIcon
-          accountId={account.id}
-          connectionResults={connectionResults}
-          accountToProviderMap={accountToProviderMap}
-        />
-      )}
-
       {/* Account ID */}
-      <span
-        className={cn("shrink-0 text-sm", isBlocked && "text-muted-foreground")}
-      >
-        {account.id}
-      </span>
+      <div className={cn(idColumnClass, "flex items-center gap-2")}>
+        {ItemIcon && <ItemIcon className="text-muted-foreground size-4 shrink-0" />}
+        <span className={cn("text-sm", isBlocked && "text-muted-foreground")}>
+          {account.id}
+        </span>
+      </div>
 
       {/* Name / alias input */}
-      {mode === TREE_ITEM_MODE.SELECTION && !isBlocked && onAliasChange ? (
-        <Input
-          className="h-7 max-w-48 text-xs"
-          placeholder="Name (optional)"
-          value={aliases[account.id] ?? account.name}
-          onChange={(e) => onAliasChange(account.id, e.target.value)}
-          onClick={(e) => e.stopPropagation()}
-        />
-      ) : (
-        <span className="text-muted-foreground text-xs">
-          {aliases[account.id] || account.name}
-        </span>
-      )}
+      <div className="min-w-0 flex-1">
+        {mode === TREE_ITEM_MODE.SELECTION && !isBlocked && onAliasChange ? (
+          <Input
+            className={aliasInputClass}
+            placeholder="Name (optional)"
+            value={aliases[account.id] ?? account.name}
+            onChange={(e) => onAliasChange(account.id, e.target.value)}
+            onClick={(e) => e.stopPropagation()}
+          />
+        ) : (
+          <span className="text-muted-foreground line-clamp-1 text-xs">
+            {aliases[account.id] || account.name}
+          </span>
+        )}
+      </div>
 
       {/* Blocked reason tooltip */}
       {isBlocked && blockedReasons.length > 0 && (
@@ -102,34 +123,4 @@ export function OrgAccountTreeItem({
   );
 }
 
-interface TestStatusIconProps {
-  accountId: string;
-  connectionResults?: Record<string, ConnectionTestStatus>;
-  accountToProviderMap?: Map<string, string>;
-}
-
-function TestStatusIcon({
-  accountId,
-  connectionResults,
-  accountToProviderMap,
-}: TestStatusIconProps) {
-  const providerId = accountToProviderMap?.get(accountId);
-  if (!providerId || !connectionResults) return null;
-
-  const status = connectionResults[providerId];
-
-  if (status === CONNECTION_TEST_STATUS.SUCCESS) {
-    return <Check className="size-4 shrink-0 text-green-500" />;
-  }
-
-  if (status === CONNECTION_TEST_STATUS.ERROR) {
-    return <AlertCircle className="text-destructive size-4 shrink-0" />;
-  }
-
-  // pending or undefined = loading
-  return (
-    <Loader2 className="text-muted-foreground size-4 shrink-0 animate-spin" />
-  );
-}
-
 export { TREE_ITEM_MODE, type TreeItemMode };

ui/components/providers/organizations/org-connection-test.tsx

@@ -1,282 +0,0 @@
-"use client";
-
-import { AlertTriangle, Loader2 } from "lucide-react";
-import { useEffect, useRef, useState } from "react";
-
-import { applyDiscovery } from "@/actions/organizations/organizations";
-import {
-  buildAccountLookup,
-  buildOrgTreeData,
-  getOuIdsForSelectedAccounts,
-} from "@/actions/organizations/organizations.adapter";
-import { checkConnectionProvider } from "@/actions/providers/providers";
-import {
-  WIZARD_FOOTER_ACTION_TYPE,
-  WizardFooterConfig,
-} from "@/components/providers/wizard/steps/footer-controls";
-import { Badge } from "@/components/shadcn";
-import { TreeView } from "@/components/shadcn/tree-view";
-import { checkTaskStatus } from "@/lib";
-import { useOrgSetupStore } from "@/store/organizations/store";
-import { CONNECTION_TEST_STATUS } from "@/types/organizations";
-
-import { OrgAccountTreeItem, TREE_ITEM_MODE } from "./org-account-tree-item";
-
-interface OrgConnectionTestProps {
-  onBack: () => void;
-  onNext: () => void;
-  onSkip: () => void;
-  onFooterChange: (config: WizardFooterConfig) => void;
-}
-
-export function OrgConnectionTest({
-  onBack,
-  onNext,
-  onSkip,
-  onFooterChange,
-}: OrgConnectionTestProps) {
-  const {
-    organizationId,
-    organizationName,
-    organizationExternalId,
-    discoveryId,
-    discoveryResult,
-    selectedAccountIds,
-    accountAliases,
-    createdProviderIds,
-    connectionResults,
-    setCreatedProviderIds,
-    setConnectionResult,
-  } = useOrgSetupStore();
-
-  const [isApplying, setIsApplying] = useState(false);
-  const [applyError, setApplyError] = useState<string | null>(null);
-  const [isTesting, setIsTesting] = useState(false);
-  const [accountToProviderMap, setAccountToProviderMap] = useState<
-    Map<string, string>
-  >(new Map());
-
-  const hasApplied = useRef(false);
-  const retryFailedActionRef = useRef<() => void>(() => {});
-
-  // On mount: apply discovery, then test connections
-  useEffect(() => {
-    if (hasApplied.current) return;
-    hasApplied.current = true;
-    handleApplyAndTest();
-    // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, []);
-
-  const handleApplyAndTest = async () => {
-    if (!organizationId || !discoveryId || !discoveryResult) return;
-
-    setIsApplying(true);
-    setApplyError(null);
-
-    // Build apply payload
-    const accounts = selectedAccountIds.map((id) => ({
-      id,
-      ...(accountAliases[id] ? { alias: accountAliases[id] } : {}),
-    }));
-
-    const ouIds = getOuIdsForSelectedAccounts(
-      discoveryResult,
-      selectedAccountIds,
-    );
-    const organizationalUnits = ouIds.map((id) => ({ id }));
-
-    const result = await applyDiscovery(
-      organizationId,
-      discoveryId,
-      accounts,
-      organizationalUnits,
-    );
-
-    if (result?.error) {
-      setApplyError(result.error);
-      setIsApplying(false);
-      return;
-    }
-
-    // Extract created provider IDs from relationships
-    const providerIds: string[] =
-      result.data?.relationships?.providers?.data?.map(
-        (p: { id: string }) => p.id,
-      ) ?? [];
-
-    setCreatedProviderIds(providerIds);
-
-    // Build account -> provider mapping
-    // The providers are returned in the same order as the accounts submitted
-    const mapping = new Map<string, string>();
-    selectedAccountIds.forEach((accountId, index) => {
-      if (providerIds[index]) {
-        mapping.set(accountId, providerIds[index]);
-      }
-    });
-    setAccountToProviderMap(mapping);
-
-    setIsApplying(false);
-
-    // Now test all connections
-    await testAllConnections(providerIds);
-  };
-
-  const testAllConnections = async (providerIds: string[]) => {
-    setIsTesting(true);
-
-    // Initialize all as pending
-    for (const id of providerIds) {
-      setConnectionResult(id, CONNECTION_TEST_STATUS.PENDING);
-    }
-
-    // Test all concurrently
-    const testPromises = providerIds.map(async (providerId) => {
-      try {
-        const formData = new FormData();
-        formData.set("providerId", providerId);
-
-        const checkResult = await checkConnectionProvider(formData);
-
-        if (checkResult?.error) {
-          setConnectionResult(providerId, CONNECTION_TEST_STATUS.ERROR);
-          return;
-        }
-
-        // Poll for task completion
-        const taskId = checkResult?.data?.id;
-        if (taskId) {
-          const taskResult = await checkTaskStatus(taskId);
-          setConnectionResult(
-            providerId,
-            taskResult.completed
-              ? CONNECTION_TEST_STATUS.SUCCESS
-              : CONNECTION_TEST_STATUS.ERROR,
-          );
-        } else {
-          // No task returned — consider it success (connection already verified)
-          setConnectionResult(providerId, CONNECTION_TEST_STATUS.SUCCESS);
-        }
-      } catch {
-        setConnectionResult(providerId, CONNECTION_TEST_STATUS.ERROR);
-      }
-    });
-
-    await Promise.allSettled(testPromises);
-    setIsTesting(false);
-
-    // Check if all passed — auto-advance
-    const allResults = useOrgSetupStore.getState().connectionResults;
-    const allPassed = providerIds.every(
-      (id) => allResults[id] === CONNECTION_TEST_STATUS.SUCCESS,
-    );
-    if (allPassed && providerIds.length > 0) {
-      onNext();
-    }
-  };
-
-  retryFailedActionRef.current = () => {
-    const failedProviderIds = createdProviderIds.filter(
-      (id) => connectionResults[id] === CONNECTION_TEST_STATUS.ERROR,
-    );
-    testAllConnections(failedProviderIds);
-  };
-
-  const hasErrors = Object.values(connectionResults).some(
-    (s) => s === CONNECTION_TEST_STATUS.ERROR,
-  );
-
-  useEffect(() => {
-    onFooterChange({
-      showBack: true,
-      backLabel: "Back",
-      backDisabled: isApplying || isTesting,
-      onBack,
-      showAction: hasErrors || isTesting,
-      actionLabel: isTesting ? "Testing..." : "Test connections",
-      actionDisabled: isApplying || isTesting || !hasErrors,
-      actionType: WIZARD_FOOTER_ACTION_TYPE.BUTTON,
-      onAction: hasErrors ? () => retryFailedActionRef.current() : undefined,
-    });
-  }, [hasErrors, isApplying, isTesting, onBack, onFooterChange]);
-
-  if (!discoveryResult) return null;
-
-  const treeData = buildOrgTreeData(discoveryResult);
-  const accountLookup = buildAccountLookup(discoveryResult);
-
-  return (
-    <div className="flex flex-col gap-5">
-      {/* Header */}
-      <div className="flex items-center gap-2">
-        <Badge variant="outline">AWS</Badge>
-        <span className="text-sm font-medium">{organizationName}</span>
-        {organizationExternalId && (
-          <Badge variant="secondary">{organizationExternalId}</Badge>
-        )}
-      </div>
-
-      {/* Apply error */}
-      {applyError && (
-        <div className="border-destructive/50 bg-destructive/10 text-destructive rounded-md border px-4 py-3 text-sm">
-          {applyError}
-        </div>
-      )}
-
-      {/* Applying state */}
-      {isApplying && (
-        <div className="flex items-center gap-3 py-8">
-          <Loader2 className="text-primary size-5 animate-spin" />
-          <span className="text-sm">
-            Applying discovery results and creating providers...
-          </span>
-        </div>
-      )}
-
-      {/* Connection test error banner */}
-      {hasErrors && !isTesting && (
-        <div className="border-destructive/50 bg-destructive/10 text-destructive flex items-start gap-3 rounded-md border px-4 py-3 text-sm">
-          <AlertTriangle className="mt-0.5 size-4 shrink-0" />
-          <span>
-            There was a problem connecting to some accounts. Ensure the Prowler
-            StackSet has successfully deployed then retry testing connections.
-          </span>
-        </div>
-      )}
-
-      {/* Tree with test results */}
-      {!isApplying && !applyError && (
-        <div className="max-h-80 overflow-y-auto rounded-md border p-2">
-          <TreeView
-            data={treeData}
-            expandAll
-            selectedIds={selectedAccountIds}
-            renderItem={(params) => (
-              <OrgAccountTreeItem
-                params={params}
-                mode={TREE_ITEM_MODE.TESTING}
-                accountLookup={accountLookup}
-                aliases={accountAliases}
-                connectionResults={connectionResults}
-                accountToProviderMap={accountToProviderMap}
-              />
-            )}
-          />
-        </div>
-      )}
-
-      {/* Optional skip */}
-      <div className="flex justify-end">
-        <div className="flex items-center gap-3">
-          <button
-            type="button"
-            onClick={onSkip}
-            className="text-muted-foreground hover:text-foreground text-sm underline"
-          >
-            Skip Connection Validation
-          </button>
-        </div>
-      </div>
-    </div>
-  );
-}

ui/components/providers/organizations/org-setup-form.tsx

@@ -87,8 +87,12 @@ export function OrgSetupForm({
   const [setupPhase, setSetupPhase] = useState<OrgSetupPhase>(
     ORG_SETUP_PHASE.DETAILS,
   );
-  const { setOrganization, setDiscovery, setSelectedAccountIds } =
-    useOrgSetupStore();
+  const {
+    setOrganization,
+    setDiscovery,
+    setSelectedAccountIds,
+    clearValidationState,
+  } = useOrgSetupStore();
   const formId = "org-wizard-setup-form";
 
   const {
@@ -189,6 +193,7 @@ export function OrgSetupForm({
   const onSubmit = async (data: OrgSetupFormData) => {
     try {
       setApiError(null);
+      clearValidationState();
       const resolvedOrganizationName =
         data.organizationName?.trim() || data.awsOrgId;
 

ui/components/providers/wizard/provider-wizard-modal.tsx

@@ -5,7 +5,6 @@ import { useRouter } from "next/navigation";
 import { useEffect, useState } from "react";
 
 import { OrgAccountSelection } from "@/components/providers/organizations/org-account-selection";
-import { OrgConnectionTest } from "@/components/providers/organizations/org-connection-test";
 import { OrgDiscoveryLoader } from "@/components/providers/organizations/org-discovery-loader";
 import { OrgLaunchScan } from "@/components/providers/organizations/org-launch-scan";
 import { OrgSetupForm } from "@/components/providers/organizations/org-setup-form";
@@ -51,7 +50,6 @@ type WizardVariant = (typeof WIZARD_VARIANT)[keyof typeof WIZARD_VARIANT];
 const VALIDATE_PHASE = {
   DISCOVERY: "discovery",
   SELECTION: "selection",
-  TESTING: "testing",
 } as const;
 
 type ValidatePhase = (typeof VALIDATE_PHASE)[keyof typeof VALIDATE_PHASE];
@@ -59,6 +57,10 @@ type ValidatePhase = (typeof VALIDATE_PHASE)[keyof typeof VALIDATE_PHASE];
 const EMPTY_FOOTER_CONFIG: WizardFooterConfig = {
   showBack: false,
   backLabel: "Back",
+  showSecondaryAction: false,
+  secondaryActionLabel: "",
+  secondaryActionVariant: "outline",
+  secondaryActionType: WIZARD_FOOTER_ACTION_TYPE.BUTTON,
   showAction: false,
   actionLabel: "Next",
   actionType: WIZARD_FOOTER_ACTION_TYPE.BUTTON,
@@ -236,6 +238,10 @@ export function ProviderWizardModal({
           showBack: true,
           backLabel: "Back",
           onBack: () => setCurrentStep(PROVIDER_WIZARD_STEP.TEST),
+          showSecondaryAction: false,
+          secondaryActionLabel: "",
+          secondaryActionVariant: "outline",
+          secondaryActionType: WIZARD_FOOTER_ACTION_TYPE.BUTTON,
           showAction: true,
           actionLabel: "Go to scans",
           actionType: WIZARD_FOOTER_ACTION_TYPE.BUTTON,
@@ -262,6 +268,10 @@ export function ProviderWizardModal({
           setOrgCurrentStep(ORG_WIZARD_STEP.SETUP);
           setOrgSetupPhase(ORG_SETUP_PHASE.DETAILS);
         },
+        showSecondaryAction: false,
+        secondaryActionLabel: "",
+        secondaryActionVariant: "outline",
+        secondaryActionType: WIZARD_FOOTER_ACTION_TYPE.BUTTON,
         showAction: false,
         actionLabel: "Next",
         actionType: WIZARD_FOOTER_ACTION_TYPE.BUTTON,
@@ -377,20 +387,6 @@ export function ProviderWizardModal({
                     setOrgCurrentStep(ORG_WIZARD_STEP.SETUP);
                     setOrgSetupPhase(ORG_SETUP_PHASE.DETAILS);
                   }}
-                  onNext={() => {
-                    setValidatePhase(VALIDATE_PHASE.TESTING);
-                  }}
-                  onFooterChange={setFooterConfig}
-                />
-              )}
-
-            {!isProviderFlow &&
-              orgCurrentStep === ORG_WIZARD_STEP.VALIDATE &&
-              validatePhase === VALIDATE_PHASE.TESTING && (
-                <OrgConnectionTest
-                  onBack={() => {
-                    setValidatePhase(VALIDATE_PHASE.SELECTION);
-                  }}
                   onNext={() => {
                     setOrgCurrentStep(ORG_WIZARD_STEP.LAUNCH);
                   }}
@@ -406,7 +402,7 @@ export function ProviderWizardModal({
                 onClose={handleClose}
                 onBack={() => {
                   setOrgCurrentStep(ORG_WIZARD_STEP.VALIDATE);
-                  setValidatePhase(VALIDATE_PHASE.TESTING);
+                  setValidatePhase(VALIDATE_PHASE.SELECTION);
                 }}
                 onFooterChange={setFooterConfig}
               />
@@ -426,7 +422,9 @@ export function ProviderWizardModal({
         </div>
       </div>
 
-      {(resolvedFooterConfig.showBack || resolvedFooterConfig.showAction) && (
+      {(resolvedFooterConfig.showBack ||
+        resolvedFooterConfig.showSecondaryAction ||
+        resolvedFooterConfig.showAction) && (
         <div className="mt-8 pt-6">
           <div className="flex items-center justify-between">
             <div>
@@ -442,7 +440,39 @@ export function ProviderWizardModal({
                 </Button>
               )}
             </div>
-            <div>
+            <div className="flex items-center gap-6">
+              {resolvedFooterConfig.showSecondaryAction && (
+                <Button
+                  size={
+                    resolvedFooterConfig.secondaryActionVariant === "link"
+                      ? "link-sm"
+                      : "xl"
+                  }
+                  className={
+                    resolvedFooterConfig.secondaryActionVariant === "link"
+                      ? "h-auto p-0"
+                      : undefined
+                  }
+                  variant={resolvedFooterConfig.secondaryActionVariant}
+                  type={
+                    resolvedFooterConfig.secondaryActionType ===
+                    WIZARD_FOOTER_ACTION_TYPE.SUBMIT
+                      ? "submit"
+                      : "button"
+                  }
+                  form={resolvedFooterConfig.secondaryActionFormId}
+                  disabled={resolvedFooterConfig.secondaryActionDisabled}
+                  onClick={
+                    resolvedFooterConfig.secondaryActionType ===
+                    WIZARD_FOOTER_ACTION_TYPE.BUTTON
+                      ? resolvedFooterConfig.onSecondaryAction
+                      : undefined
+                  }
+                >
+                  {resolvedFooterConfig.secondaryActionLabel}
+                </Button>
+              )}
+
               {resolvedFooterConfig.showAction && (
                 <Button
                   size="xl"

ui/components/providers/wizard/wizard-stepper.tsx

@@ -1,7 +1,6 @@
 "use client";
 
 import {
-  AlertCircle,
   CircleCheckBig,
   FolderGit2,
   KeyRound,
@@ -52,7 +51,7 @@ const STEPS: StepConfig[] = [
 
 export function WizardStepper({
   currentStep,
-  hasConnectionErrors = false,
+  hasConnectionErrors: _hasConnectionErrors = false,
   stepOffset = 0,
 }: WizardStepperProps) {
   const activeVisualStep = Math.max(
@@ -66,17 +65,11 @@ export function WizardStepper({
         const isComplete = index < activeVisualStep;
         const isActive = index === activeVisualStep;
         const isInactive = index > activeVisualStep;
-        const isErrorStep = index === 2 && hasConnectionErrors && isActive;
 
         return (
           <div key={step.label} className="flex items-start gap-3">
             <div className="flex flex-col items-center">
-              <StepCircle
-                isComplete={isComplete}
-                isActive={isActive}
-                isError={isErrorStep}
-                icon={step.icon}
-              />
+              <StepCircle isComplete={isComplete} isActive={isActive} icon={step.icon} />
               {index < STEPS.length - 1 && (
                 <StepConnector isComplete={isComplete} />
               )}
@@ -89,7 +82,6 @@ export function WizardStepper({
                   isActive && "text-[#f4f4f5]",
                   isComplete && "text-[#f4f4f5]",
                   isInactive && "text-[#f4f4f5]/60",
-                  isErrorStep && "text-destructive",
                 )}
               >
                 {step.label}
@@ -108,24 +100,14 @@ export function WizardStepper({
 interface StepCircleProps {
   isComplete: boolean;
   isActive: boolean;
-  isError: boolean;
   icon: IconComponent;
 }
 
 function StepCircle({
   isComplete,
   isActive,
-  isError,
   icon: Icon,
 }: StepCircleProps) {
-  if (isError) {
-    return (
-      <div className="border-destructive flex size-[44px] shrink-0 items-center justify-center rounded-full border bg-[#121110]">
-        <AlertCircle className="text-destructive size-6" />
-      </div>
-    );
-  }
-
   if (isComplete) {
     return (
       <div className="bg-button-primary-press flex size-[44px] shrink-0 items-center justify-center rounded-full">

ui/store/organizations/store.ts

@@ -22,6 +22,7 @@ interface OrgSetupState {
 
   // Connection test results
   connectionResults: Record<string, ConnectionTestStatus>;
+  connectionErrors: Record<string, string>;
 
   // Actions
   setOrganization: (id: string, name: string, externalId: string) => void;
@@ -29,6 +30,8 @@ interface OrgSetupState {
   setSelectedAccountIds: (ids: string[]) => void;
   setAccountAlias: (accountId: string, alias: string) => void;
   setCreatedProviderIds: (ids: string[]) => void;
+  clearValidationState: () => void;
+  setConnectionError: (providerId: string, error: string | null) => void;
   setConnectionResult: (
     providerId: string,
     status: ConnectionTestStatus,
@@ -46,6 +49,7 @@ const initialState = {
   accountAliases: {},
   createdProviderIds: [],
   connectionResults: {},
+  connectionErrors: {},
 };
 
 export const useOrgSetupStore = create<OrgSetupState>()(
@@ -72,6 +76,28 @@ export const useOrgSetupStore = create<OrgSetupState>()(
 
       setCreatedProviderIds: (ids) => set({ createdProviderIds: ids }),
 
+      clearValidationState: () =>
+        set({
+          createdProviderIds: [],
+          connectionResults: {},
+          connectionErrors: {},
+        }),
+
+      setConnectionError: (providerId, error) =>
+        set((state) => {
+          if (!error) {
+            const { [providerId]: _, ...rest } = state.connectionErrors;
+            return { connectionErrors: rest };
+          }
+
+          return {
+            connectionErrors: {
+              ...state.connectionErrors,
+              [providerId]: error,
+            },
+          };
+        }),
+
       setConnectionResult: (providerId, status) =>
         set((state) => ({
           connectionResults: {