feat(ui): add organization and wizard types and stores (#10154)

ui/actions/organizations/organizations.adapter.ts

@@ -0,0 +1,141 @@
+import { Box, Folder, FolderTree } from "lucide-react";
+
+import {
+  APPLY_STATUS,
+  DiscoveredAccount,
+  DiscoveryResult,
+} from "@/types/organizations";
+import { TreeDataItem } from "@/types/tree";
+
+/**
+ * Transforms flat API discovery arrays into hierarchical TreeDataItem[] for TreeView.
+ *
+ * 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[] {
+  const nodeMap = new Map<string, TreeDataItem>();
+
+  for (const root of result.roots) {
+    nodeMap.set(root.id, {
+      id: root.id,
+      name: root.name,
+      icon: FolderTree,
+      children: [],
+    });
+  }
+
+  for (const ou of result.organizational_units) {
+    nodeMap.set(ou.id, {
+      id: ou.id,
+      name: ou.name,
+      icon: Folder,
+      children: [],
+    });
+  }
+
+  for (const account of result.accounts) {
+    const isBlocked =
+      account.registration?.apply_status === APPLY_STATUS.BLOCKED;
+
+    nodeMap.set(account.id, {
+      id: account.id,
+      name: `${account.id} — ${account.name}`,
+      icon: Box,
+      disabled: isBlocked,
+    });
+  }
+
+  for (const ou of result.organizational_units) {
+    const parent = nodeMap.get(ou.parent_id);
+    if (parent?.children) {
+      const ouNode = nodeMap.get(ou.id);
+      if (ouNode) {
+        parent.children.push(ouNode);
+      }
+    }
+  }
+
+  for (const account of result.accounts) {
+    const parent = nodeMap.get(account.parent_id);
+    if (!parent) {
+      continue;
+    }
+
+    if (!parent.children) {
+      parent.children = [];
+    }
+
+    const accountNode = nodeMap.get(account.id);
+    if (accountNode) {
+      parent.children.push(accountNode);
+    }
+  }
+
+  return result.roots.flatMap((root) => {
+    const rootNode = nodeMap.get(root.id);
+    return rootNode?.children ?? [];
+  });
+}
+
+/**
+ * Returns IDs of accounts that can be selected.
+ * Accounts are selectable when registration is READY or not yet present.
+ * Accounts with explicit non-ready states are excluded.
+ */
+export function getSelectableAccountIds(result: DiscoveryResult): string[] {
+  return result.accounts
+    .filter((account) => {
+      const applyStatus = account.registration?.apply_status;
+      if (!applyStatus) {
+        return true;
+      }
+      return applyStatus === APPLY_STATUS.READY;
+    })
+    .map((account) => account.id);
+}
+
+/**
+ * Creates a lookup map from account ID to DiscoveredAccount.
+ */
+export function buildAccountLookup(
+  result: DiscoveryResult,
+): Map<string, DiscoveredAccount> {
+  const map = new Map<string, DiscoveredAccount>();
+  for (const account of result.accounts) {
+    map.set(account.id, account);
+  }
+  return map;
+}
+
+/**
+ * Given selected account IDs, returns OU IDs that are ancestors of selected accounts.
+ */
+export function getOuIdsForSelectedAccounts(
+  result: DiscoveryResult,
+  selectedAccountIds: string[],
+): string[] {
+  const selectedSet = new Set(selectedAccountIds);
+  const ouIds = new Set<string>();
+  const allOuIds = new Set(result.organizational_units.map((ou) => ou.id));
+  const ouParentMap = new Map<string, string>();
+
+  for (const ou of result.organizational_units) {
+    ouParentMap.set(ou.id, ou.parent_id);
+  }
+
+  for (const account of result.accounts) {
+    if (!selectedSet.has(account.id)) {
+      continue;
+    }
+
+    let currentParentId = account.parent_id;
+    while (currentParentId && allOuIds.has(currentParentId)) {
+      ouIds.add(currentParentId);
+      currentParentId = ouParentMap.get(currentParentId) ?? "";
+    }
+  }
+
+  return Array.from(ouIds);
+}

ui/store/index.ts

@@ -1 +1,3 @@
+export * from "./organizations/store";
+export * from "./provider-wizard/store";
 export * from "./ui/store";

ui/store/organizations/store.test.ts

@@ -0,0 +1,33 @@
+import { beforeEach, describe, expect, it } from "vitest";
+
+import { useOrgSetupStore } from "./store";
+
+describe("useOrgSetupStore", () => {
+  beforeEach(() => {
+    sessionStorage.clear();
+    localStorage.clear();
+    useOrgSetupStore.getState().reset();
+  });
+
+  it("persists organization wizard state in sessionStorage", () => {
+    // Given
+    useOrgSetupStore
+      .getState()
+      .setOrganization("org-1", "My Org", "o-abc123def4");
+    useOrgSetupStore.getState().setDiscovery("discovery-1", {
+      roots: [],
+      organizational_units: [],
+      accounts: [],
+    });
+    useOrgSetupStore
+      .getState()
+      .setSelectedAccountIds(["111111111111", "222222222222"]);
+
+    // When
+    const persistedValue = sessionStorage.getItem("org-setup-store");
+
+    // Then
+    expect(persistedValue).toBeTruthy();
+    expect(localStorage.getItem("org-setup-store")).toBeNull();
+  });
+});

ui/store/organizations/store.ts

@@ -0,0 +1,203 @@
+import { create } from "zustand";
+import { createJSONStorage, persist } from "zustand/middleware";
+
+import {
+  buildAccountLookup,
+  buildOrgTreeData,
+  getSelectableAccountIds,
+} from "@/actions/organizations/organizations.adapter";
+import {
+  ConnectionTestStatus,
+  DiscoveredAccount,
+  DiscoveryResult,
+} from "@/types/organizations";
+import { TreeDataItem } from "@/types/tree";
+
+interface DerivedDiscoveryState {
+  treeData: TreeDataItem[];
+  accountLookup: Map<string, DiscoveredAccount>;
+  selectableAccountIds: string[];
+  selectableAccountIdSet: Set<string>;
+}
+
+function buildDerivedDiscoveryState(
+  discoveryResult: DiscoveryResult | null,
+): DerivedDiscoveryState {
+  if (!discoveryResult) {
+    return {
+      treeData: [],
+      accountLookup: new Map<string, DiscoveredAccount>(),
+      selectableAccountIds: [],
+      selectableAccountIdSet: new Set<string>(),
+    };
+  }
+
+  const selectableAccountIds = getSelectableAccountIds(discoveryResult);
+  return {
+    treeData: buildOrgTreeData(discoveryResult),
+    accountLookup: buildAccountLookup(discoveryResult),
+    selectableAccountIds,
+    selectableAccountIdSet: new Set(selectableAccountIds),
+  };
+}
+
+interface OrgSetupState {
+  // Identity
+  organizationId: string | null;
+  organizationName: string | null;
+  organizationExternalId: string | null;
+  discoveryId: string | null;
+
+  // Discovery
+  discoveryResult: DiscoveryResult | null;
+  treeData: TreeDataItem[];
+  accountLookup: Map<string, DiscoveredAccount>;
+  selectableAccountIds: string[];
+  selectableAccountIdSet: Set<string>;
+
+  // Selection + aliases
+  selectedAccountIds: string[];
+  accountAliases: Record<string, string>;
+
+  // Apply result
+  createdProviderIds: string[];
+
+  // Connection test results
+  connectionResults: Record<string, ConnectionTestStatus>;
+  connectionErrors: Record<string, string>;
+
+  // Actions
+  setOrganization: (id: string, name: string, externalId: string) => void;
+  setDiscovery: (id: string, result: DiscoveryResult) => void;
+  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,
+  ) => void;
+  reset: () => void;
+}
+
+const initialState = {
+  organizationId: null,
+  organizationName: null,
+  organizationExternalId: null,
+  discoveryId: null,
+  discoveryResult: null,
+  treeData: [],
+  accountLookup: new Map<string, DiscoveredAccount>(),
+  selectableAccountIds: [],
+  selectableAccountIdSet: new Set<string>(),
+  selectedAccountIds: [],
+  accountAliases: {},
+  createdProviderIds: [],
+  connectionResults: {},
+  connectionErrors: {},
+};
+
+export const useOrgSetupStore = create<OrgSetupState>()(
+  persist(
+    (set) => ({
+      ...initialState,
+
+      setOrganization: (id, name, externalId) =>
+        set({
+          organizationId: id,
+          organizationName: name,
+          organizationExternalId: externalId,
+        }),
+
+      setDiscovery: (id, result) =>
+        set((state) => {
+          const derivedState = buildDerivedDiscoveryState(result);
+          return {
+            discoveryId: id,
+            discoveryResult: result,
+            ...derivedState,
+            selectedAccountIds: state.selectedAccountIds.filter((accountId) =>
+              derivedState.selectableAccountIdSet.has(accountId),
+            ),
+          };
+        }),
+
+      setSelectedAccountIds: (ids) =>
+        set((state) => ({
+          selectedAccountIds: ids.filter((accountId) =>
+            state.selectableAccountIdSet.has(accountId),
+          ),
+        })),
+
+      setAccountAlias: (accountId, alias) =>
+        set((state) => ({
+          accountAliases: { ...state.accountAliases, [accountId]: alias },
+        })),
+
+      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: {
+            ...state.connectionResults,
+            [providerId]: status,
+          },
+        })),
+
+      reset: () => set(initialState),
+    }),
+    {
+      name: "org-setup-store",
+      storage: createJSONStorage(() => sessionStorage),
+      merge: (persistedState, currentState) => {
+        const mergedState = {
+          ...currentState,
+          ...(persistedState as Partial<OrgSetupState>),
+        };
+        const derivedState = buildDerivedDiscoveryState(
+          mergedState.discoveryResult,
+        );
+
+        return {
+          ...mergedState,
+          ...derivedState,
+          selectedAccountIds: mergedState.selectedAccountIds.filter(
+            (accountId) => derivedState.selectableAccountIdSet.has(accountId),
+          ),
+        };
+      },
+      partialize: (state) => ({
+        organizationId: state.organizationId,
+        organizationName: state.organizationName,
+        organizationExternalId: state.organizationExternalId,
+        discoveryId: state.discoveryId,
+        discoveryResult: state.discoveryResult,
+        selectedAccountIds: state.selectedAccountIds,
+        accountAliases: state.accountAliases,
+      }),
+    },
+  ),
+);

ui/store/provider-wizard/store.test.ts

@@ -0,0 +1,62 @@
+import { beforeEach, describe, expect, it } from "vitest";
+
+import { PROVIDER_WIZARD_MODE } from "@/types/provider-wizard";
+
+import { useProviderWizardStore } from "./store";
+
+describe("useProviderWizardStore", () => {
+  beforeEach(() => {
+    sessionStorage.clear();
+    localStorage.clear();
+    useProviderWizardStore.getState().reset();
+  });
+
+  it("stores provider identity and mode, then resets to defaults", () => {
+    useProviderWizardStore.getState().setProvider({
+      id: "provider-1",
+      type: "aws",
+      uid: "123456789012",
+      alias: "prod-account",
+    });
+    useProviderWizardStore.getState().setVia("role");
+    useProviderWizardStore.getState().setSecretId("secret-1");
+    useProviderWizardStore.getState().setMode(PROVIDER_WIZARD_MODE.UPDATE);
+
+    const afterSet = useProviderWizardStore.getState();
+    expect(afterSet.providerId).toBe("provider-1");
+    expect(afterSet.providerType).toBe("aws");
+    expect(afterSet.providerUid).toBe("123456789012");
+    expect(afterSet.providerAlias).toBe("prod-account");
+    expect(afterSet.via).toBe("role");
+    expect(afterSet.secretId).toBe("secret-1");
+    expect(afterSet.mode).toBe(PROVIDER_WIZARD_MODE.UPDATE);
+
+    useProviderWizardStore.getState().reset();
+    const afterReset = useProviderWizardStore.getState();
+
+    expect(afterReset.providerId).toBeNull();
+    expect(afterReset.providerType).toBeNull();
+    expect(afterReset.providerUid).toBeNull();
+    expect(afterReset.providerAlias).toBeNull();
+    expect(afterReset.via).toBeNull();
+    expect(afterReset.secretId).toBeNull();
+    expect(afterReset.mode).toBe(PROVIDER_WIZARD_MODE.ADD);
+  });
+
+  it("persists provider wizard state in sessionStorage", () => {
+    // Given
+    useProviderWizardStore.getState().setProvider({
+      id: "provider-1",
+      type: "aws",
+      uid: "123456789012",
+      alias: "prod-account",
+    });
+
+    // When
+    const persistedValue = sessionStorage.getItem("provider-wizard-store");
+
+    // Then
+    expect(persistedValue).toBeTruthy();
+    expect(localStorage.getItem("provider-wizard-store")).toBeNull();
+  });
+});

ui/store/provider-wizard/store.ts

@@ -0,0 +1,57 @@
+import { create } from "zustand";
+import { createJSONStorage, persist } from "zustand/middleware";
+
+import {
+  PROVIDER_WIZARD_MODE,
+  ProviderWizardIdentity,
+  ProviderWizardMode,
+} from "@/types/provider-wizard";
+import { ProviderType } from "@/types/providers";
+
+interface ProviderWizardState {
+  providerId: string | null;
+  providerType: ProviderType | null;
+  providerUid: string | null;
+  providerAlias: string | null;
+  via: string | null;
+  secretId: string | null;
+  mode: ProviderWizardMode;
+  setProvider: (provider: ProviderWizardIdentity) => void;
+  setVia: (via: string | null) => void;
+  setSecretId: (secretId: string | null) => void;
+  setMode: (mode: ProviderWizardMode) => void;
+  reset: () => void;
+}
+
+const initialState = {
+  providerId: null,
+  providerType: null,
+  providerUid: null,
+  providerAlias: null,
+  via: null,
+  secretId: null,
+  mode: PROVIDER_WIZARD_MODE.ADD,
+};
+
+export const useProviderWizardStore = create<ProviderWizardState>()(
+  persist(
+    (set) => ({
+      ...initialState,
+      setProvider: (provider) =>
+        set({
+          providerId: provider.id,
+          providerType: provider.type,
+          providerUid: provider.uid,
+          providerAlias: provider.alias,
+        }),
+      setVia: (via) => set({ via }),
+      setSecretId: (secretId) => set({ secretId }),
+      setMode: (mode) => set({ mode }),
+      reset: () => set(initialState),
+    }),
+    {
+      name: "provider-wizard-store",
+      storage: createJSONStorage(() => sessionStorage),
+    },
+  ),
+);

ui/types/index.ts

@@ -2,7 +2,9 @@ export * from "./authFormSchema";
 export * from "./components";
 export * from "./filters";
 export * from "./formSchemas";
+export * from "./organizations";
 export * from "./processors";
+export * from "./provider-wizard";
 export * from "./providers";
 export * from "./resources";
 export * from "./scans";

ui/types/organizations.ts

@@ -0,0 +1,205 @@
+// ─── Const Enums ──────────────────────────────────────────────────────────────
+
+export const DISCOVERY_STATUS = {
+  PENDING: "pending",
+  RUNNING: "running",
+  SUCCEEDED: "succeeded",
+  FAILED: "failed",
+} as const;
+
+export type DiscoveryStatus =
+  (typeof DISCOVERY_STATUS)[keyof typeof DISCOVERY_STATUS];
+
+export const APPLY_STATUS = {
+  READY: "ready",
+  BLOCKED: "blocked",
+} as const;
+
+export type ApplyStatus = (typeof APPLY_STATUS)[keyof typeof APPLY_STATUS];
+
+export const ORG_RELATION = {
+  ALREADY_LINKED: "already_linked",
+  LINK_REQUIRED: "link_required",
+  LINKED_TO_OTHER: "linked_to_other_organization",
+} as const;
+
+export type OrgRelation = (typeof ORG_RELATION)[keyof typeof ORG_RELATION];
+
+export const OU_RELATION = {
+  NOT_APPLICABLE: "not_applicable",
+  ALREADY_LINKED: "already_linked",
+  LINK_REQUIRED: "link_required",
+  LINKED_TO_OTHER_OU: "linked_to_other_ou",
+  UNCHANGED: "unchanged",
+} as const;
+
+export type OuRelation = (typeof OU_RELATION)[keyof typeof OU_RELATION];
+
+export const SECRET_STATE = {
+  ALREADY_EXISTS: "already_exists",
+  WILL_CREATE: "will_create",
+  MANUAL_REQUIRED: "manual_required",
+} as const;
+
+export type SecretState = (typeof SECRET_STATE)[keyof typeof SECRET_STATE];
+
+export const ORG_WIZARD_STEP = {
+  SETUP: 0,
+  VALIDATE: 1,
+  LAUNCH: 2,
+} as const;
+
+export type OrgWizardStep =
+  (typeof ORG_WIZARD_STEP)[keyof typeof ORG_WIZARD_STEP];
+
+export const ORG_SETUP_PHASE = {
+  DETAILS: "details",
+  ACCESS: "access",
+} as const;
+
+export type OrgSetupPhase =
+  (typeof ORG_SETUP_PHASE)[keyof typeof ORG_SETUP_PHASE];
+
+export const DISCOVERED_ACCOUNT_STATUS = {
+  ACTIVE: "ACTIVE",
+  SUSPENDED: "SUSPENDED",
+  PENDING_CLOSURE: "PENDING_CLOSURE",
+  CLOSED: "CLOSED",
+} as const;
+
+export type DiscoveredAccountStatus =
+  (typeof DISCOVERED_ACCOUNT_STATUS)[keyof typeof DISCOVERED_ACCOUNT_STATUS];
+
+export const DISCOVERED_ACCOUNT_JOINED_METHOD = {
+  INVITED: "INVITED",
+  CREATED: "CREATED",
+} as const;
+
+export type DiscoveredAccountJoinedMethod =
+  (typeof DISCOVERED_ACCOUNT_JOINED_METHOD)[keyof typeof DISCOVERED_ACCOUNT_JOINED_METHOD];
+
+export interface OrganizationPolicyType {
+  Type: string;
+  Status: string;
+}
+
+export const ORGANIZATION_TYPE = {
+  AWS: "aws",
+  AZURE: "azure",
+  GCP: "gcp",
+} as const;
+
+export type OrganizationType =
+  (typeof ORGANIZATION_TYPE)[keyof typeof ORGANIZATION_TYPE];
+
+// ─── Discovery Result Interfaces ──────────────────────────────────────────────
+
+export interface AccountRegistration {
+  provider_exists: boolean;
+  provider_id: string | null;
+  organization_relation: OrgRelation;
+  organizational_unit_relation: OuRelation;
+  provider_secret_state: SecretState;
+  apply_status: ApplyStatus;
+  blocked_reasons: string[];
+}
+
+export interface DiscoveredAccount {
+  id: string;
+  name: string;
+  arn: string;
+  email: string;
+  status: DiscoveredAccountStatus;
+  joined_method: DiscoveredAccountJoinedMethod;
+  joined_timestamp: string;
+  parent_id: string;
+  registration?: AccountRegistration;
+}
+
+export interface DiscoveredOu {
+  id: string;
+  name: string;
+  arn: string;
+  parent_id: string;
+}
+
+export interface DiscoveredRoot {
+  id: string;
+  arn: string;
+  name: string;
+  policy_types: OrganizationPolicyType[];
+}
+
+export interface DiscoveryResult {
+  roots: DiscoveredRoot[];
+  organizational_units: DiscoveredOu[];
+  accounts: DiscoveredAccount[];
+}
+
+// ─── JSON:API Resource Interfaces ─────────────────────────────────────────────
+
+export interface OrganizationAttributes {
+  name: string;
+  org_type: OrganizationType;
+  external_id: string;
+  metadata: Record<string, unknown>;
+  root_external_id: string | null;
+  inserted_at?: string;
+  updated_at?: string;
+}
+
+export interface OrganizationResource {
+  id: string;
+  type: "organizations";
+  attributes: OrganizationAttributes;
+}
+
+export interface DiscoveryAttributes {
+  status: DiscoveryStatus;
+  result: DiscoveryResult | Record<string, never>;
+  error: string | null;
+  inserted_at: string;
+  updated_at: string;
+}
+
+export interface DiscoveryResource {
+  id: string;
+  type: "organization-discoveries";
+  attributes: DiscoveryAttributes;
+}
+
+export interface ApplyResultAttributes {
+  providers_created_count: number;
+  providers_linked_count: number;
+  providers_applied_count: number;
+  organizational_units_created_count: number;
+}
+
+export interface ApplyResultRelationships {
+  providers: {
+    data: Array<{ type: "providers"; id: string }>;
+    meta: { count: number };
+  };
+  organizational_units: {
+    data: Array<{ type: "organizational-units"; id: string }>;
+    meta: { count: number };
+  };
+}
+
+export interface ApplyResultResource {
+  id: string;
+  type: "organization-discovery-apply-results";
+  attributes: ApplyResultAttributes;
+  relationships: ApplyResultRelationships;
+}
+
+// ─── Connection Test Status ───────────────────────────────────────────────────
+
+export const CONNECTION_TEST_STATUS = {
+  PENDING: "pending",
+  SUCCESS: "success",
+  ERROR: "error",
+} as const;
+
+export type ConnectionTestStatus =
+  (typeof CONNECTION_TEST_STATUS)[keyof typeof CONNECTION_TEST_STATUS];

ui/types/provider-wizard.ts

@@ -0,0 +1,26 @@
+import { ProviderType } from "./providers";
+
+export const PROVIDER_WIZARD_STEP = {
+  CONNECT: 0,
+  CREDENTIALS: 1,
+  TEST: 2,
+  LAUNCH: 3,
+} as const;
+
+export type ProviderWizardStep =
+  (typeof PROVIDER_WIZARD_STEP)[keyof typeof PROVIDER_WIZARD_STEP];
+
+export const PROVIDER_WIZARD_MODE = {
+  ADD: "add",
+  UPDATE: "update",
+} as const;
+
+export type ProviderWizardMode =
+  (typeof PROVIDER_WIZARD_MODE)[keyof typeof PROVIDER_WIZARD_MODE];
+
+export interface ProviderWizardIdentity {
+  id: string;
+  type: ProviderType;
+  uid: string | null;
+  alias: string | null;
+}