fix: organization setup redirect if it didn't create correctly

apps/app/src/app/[locale]/(app)/(dashboard)/[orgId]/frameworks/data/getAllFrameworkInstancesWithControls.ts

@@ -1,31 +1,34 @@
 "use server";
 
+import { cache } from "react";
 import { db } from "@comp/db";
 import type { FrameworkInstanceWithControls } from "../types";
 
-export async function getAllFrameworkInstancesWithControls({
-	organizationId,
-}: {
-	organizationId: string;
-}): Promise<FrameworkInstanceWithControls[]> {
-	const frameworksWithControls = await db.frameworkInstance.findMany({
-		where: {
-			organizationId,
-		},
-		include: {
-			controls: {
-				include: {
-					artifacts: {
-						include: {
-							policy: true,
-							evidence: true,
+export const getAllFrameworkInstancesWithControls = cache(
+	async ({
+		organizationId,
+	}: {
+		organizationId: string;
+	}): Promise<FrameworkInstanceWithControls[]> => {
+		const frameworksWithControls = await db.frameworkInstance.findMany({
+			where: {
+				organizationId,
+			},
+			include: {
+				controls: {
+					include: {
+						artifacts: {
+							include: {
+								policy: true,
+								evidence: true,
+							},
 						},
+						requirementsMapped: true,
 					},
-					requirementsMapped: true,
 				},
 			},
-		},
-	});
+		});
 
-	return frameworksWithControls;
-}
+		return frameworksWithControls;
+	},
+);

apps/app/src/app/[locale]/(app)/(dashboard)/[orgId]/frameworks/data/getSingleFrameworkInstanceWithControls.ts

@@ -1,34 +1,37 @@
 "use server";
 
+import { cache } from "react";
 import { db } from "@comp/db";
 import type { FrameworkInstanceWithControls } from "../types";
 
-export async function getSingleFrameworkInstanceWithControls({
-	organizationId,
-	frameworkInstanceId,
-}: {
-	organizationId: string;
-	frameworkInstanceId: string;
-}): Promise<FrameworkInstanceWithControls | null> {
-	const frameworkInstance = await db.frameworkInstance.findUnique({
-		where: {
-			organizationId,
-			id: frameworkInstanceId,
-		},
-		include: {
-			controls: {
-				include: {
-					artifacts: {
-						include: {
-							policy: true,
-							evidence: true,
+export const getSingleFrameworkInstanceWithControls = cache(
+	async ({
+		organizationId,
+		frameworkInstanceId,
+	}: {
+		organizationId: string;
+		frameworkInstanceId: string;
+	}): Promise<FrameworkInstanceWithControls | null> => {
+		const frameworkInstance = await db.frameworkInstance.findUnique({
+			where: {
+				organizationId,
+				id: frameworkInstanceId,
+			},
+			include: {
+				controls: {
+					include: {
+						artifacts: {
+							include: {
+								policy: true,
+								evidence: true,
+							},
 						},
+						requirementsMapped: true,
 					},
-					requirementsMapped: true,
 				},
 			},
-		},
-	});
+		});
 
-	return frameworkInstance;
-}
+		return frameworkInstance;
+	},
+);

apps/app/src/app/[locale]/(app)/(dashboard)/[orgId]/frameworks/page.tsx

@@ -2,37 +2,37 @@ import { getI18n } from "@/locales/server";
 import { auth } from "@/utils/auth";
 import { headers } from "next/headers";
 import { redirect } from "next/navigation";
-import PageWithBreadcrumb from "../../../../../../components/pages/PageWithBreadcrumb";
+import PageWithBreadcrumb from "@/components/pages/PageWithBreadcrumb";
 import { FrameworksOverview } from "./components/FrameworksOverview";
 import { getAllFrameworkInstancesWithControls } from "./data/getAllFrameworkInstancesWithControls";
 
 export async function generateMetadata() {
-	const t = await getI18n();
+  const t = await getI18n();
 
-	return {
-		title: t("sidebar.frameworks"),
-	};
+  return {
+    title: t("sidebar.frameworks"),
+  };
 }
 
 export default async function DashboardPage() {
-	const session = await auth.api.getSession({
-		headers: await headers(),
-	});
+  const session = await auth.api.getSession({
+    headers: await headers(),
+  });
 
-	const organizationId = session?.session.activeOrganizationId;
+  const organizationId = session?.session.activeOrganizationId;
 
-	if (!organizationId) {
-		console.log("Redirect on page.tsx");
-		redirect("/");
-	}
+  if (!organizationId) {
+    console.log("Redirect on page.tsx");
+    redirect("/");
+  }
 
-	const frameworksWithControls = await getAllFrameworkInstancesWithControls({
-		organizationId,
-	});
+  const frameworksWithControls = await getAllFrameworkInstancesWithControls({
+    organizationId,
+  });
 
-	return (
-		<PageWithBreadcrumb breadcrumbs={[{ label: "Frameworks", current: true }]}>
-			<FrameworksOverview frameworksWithControls={frameworksWithControls} />
-		</PageWithBreadcrumb>
-	);
+  return (
+    <PageWithBreadcrumb breadcrumbs={[{ label: "Frameworks", current: true }]}>
+      <FrameworksOverview frameworksWithControls={frameworksWithControls} />
+    </PageWithBreadcrumb>
+  );
 }

apps/app/src/app/[locale]/(app)/(dashboard)/[orgId]/layout.tsx

@@ -4,54 +4,72 @@ import { AssistantSheet } from "@/components/sheets/assistant-sheet";
 import { Sidebar } from "@/components/sidebar";
 import { SidebarProvider } from "@/context/sidebar-context";
 import { auth } from "@/utils/auth";
+import { db } from "@comp/db";
 import dynamic from "next/dynamic";
 import { cookies, headers } from "next/headers";
 import { redirect } from "next/navigation";
+import { cache } from "react";
 
 const HotKeys = dynamic(
-	() => import("@/components/hot-keys").then((mod) => mod.HotKeys),
-	{
-		ssr: true,
-	},
+  () => import("@/components/hot-keys").then((mod) => mod.HotKeys),
+  {
+    ssr: true,
+  },
 );
 
 export default async function Layout({
-	children,
-	params,
+  children,
+  params,
 }: {
-	children: React.ReactNode;
-	params: Promise<{ orgId: string }>;
+  children: React.ReactNode;
+  params: Promise<{ orgId: string }>;
 }) {
-	const session = await auth.api.getSession({
-		headers: await headers(),
-	});
-
-	const orgId = (await params).orgId;
-	const cookieStore = await cookies();
-	const isCollapsed = cookieStore.get("sidebar-collapsed")?.value === "true";
-
-	if (!session) {
-		redirect("/auth");
-	}
-
-	if (!orgId) {
-		redirect("/");
-	}
-
-	if (!session?.session.activeOrganizationId) {
-		redirect("/");
-	}
-
-	return (
-		<SidebarProvider initialIsCollapsed={isCollapsed}>
-			<AnimatedLayout sidebar={<Sidebar />} isCollapsed={isCollapsed}>
-				<div className="mx-4 md:ml-[95px] md:mr-10 pb-8">
-					<Header />
-					<main>{children}</main>
-				</div>
-				<AssistantSheet />
-			</AnimatedLayout>
-			<HotKeys />
-		</SidebarProvider>
-	);
+  const session = await auth.api.getSession({
+    headers: await headers(),
+  });
+
+  const orgId = (await params).orgId;
+  await getOrganization(orgId);
+  const cookieStore = await cookies();
+  const isCollapsed = cookieStore.get("sidebar-collapsed")?.value === "true";
+
+  if (!session) {
+    redirect("/auth");
+  }
+
+  if (!session?.session.activeOrganizationId) {
+    redirect("/");
+  }
+
+  return (
+    <SidebarProvider initialIsCollapsed={isCollapsed}>
+      <AnimatedLayout sidebar={<Sidebar />} isCollapsed={isCollapsed}>
+        <div className="mx-4 md:ml-[95px] md:mr-10 pb-8">
+          <Header />
+          <main>{children}</main>
+        </div>
+        <AssistantSheet />
+      </AnimatedLayout>
+      <HotKeys />
+    </SidebarProvider>
+  );
 }
+
+const getOrganization = cache(async (orgId: string) => {
+  try {
+    const organization = await db.organization.findUnique({
+      where: {
+        id: orgId,
+      },
+    });
+
+    if (!organization) {
+      redirect("/setup");
+    }
+
+    return organization;
+  } catch (error) {
+    console.error(error);
+    redirect("/setup");
+  }
+});

apps/app/src/app/[locale]/(app)/setup/page.tsx

@@ -7,50 +7,46 @@ import { redirect } from "next/navigation";
 import { AcceptInvite } from "./components/accept-invite";
 
 export const metadata: Metadata = {
-	title: "Organization Setup | Comp AI",
+  title: "Organization Setup | Comp AI",
 };
 
 export default async function Page() {
-	const headersList = await headers();
-
-	const session = await auth.api.getSession({
-		headers: headersList,
-	});
-
-	if (session?.session?.activeOrganizationId) {
-		redirect(`/${session.session.activeOrganizationId}/frameworks`);
-	}
-
-	const organization = await db.organization.findFirst({
-		where: {
-			members: {
-				some: {
-					userId: session?.user.id,
-				},
-			},
-		},
-	});
-
-	if (!session?.session?.activeOrganizationId && organization?.id) {
-		await auth.api.setActiveOrganization({
-			headers: headersList,
-			body: {
-				organizationId: organization.id,
-			},
-		});
-		redirect(`/${organization.id}/frameworks`);
-	}
-
-	const hasInvite = await db.invitation.findFirst({
-		where: {
-			email: session?.user.email,
-			status: "pending",
-		},
-	});
-
-	if (hasInvite) {
-		return <AcceptInvite inviteCode={hasInvite.id} />;
-	}
-
-	return <OnboardingClient />;
+  const headersList = await headers();
+
+  const session = await auth.api.getSession({
+    headers: headersList,
+  });
+
+  const organization = await db.organization.findFirst({
+    where: {
+      members: {
+        some: {
+          userId: session?.user.id,
+        },
+      },
+    },
+  });
+
+  if (!session?.session?.activeOrganizationId && organization?.id) {
+    await auth.api.setActiveOrganization({
+      headers: headersList,
+      body: {
+        organizationId: organization.id,
+      },
+    });
+    redirect(`/${organization.id}/frameworks`);
+  }
+
+  const hasInvite = await db.invitation.findFirst({
+    where: {
+      email: session?.user.email,
+      status: "pending",
+    },
+  });
+
+  if (hasInvite) {
+    return <AcceptInvite inviteCode={hasInvite.id} />;
+  }
+
+  return <OnboardingClient />;
 }