fix: flawed organization guards

apps/app/src/app/(main)/organizations/_components/organization-guard.tsx

@@ -1,7 +1,6 @@
 import { apiServer } from '@bulkit/app/api/api.server'
 import { OrganizationProvider } from '@bulkit/app/app/(main)/organizations/_components/selected-organization-provider'
-import { ORGANIZATION_COOKIE_NAME } from '@bulkit/app/app/(main)/organizations/organizations.constants'
-import { cookies } from 'next/headers'
+import { fetchServerOrganization } from '@bulkit/app/app/(main)/organizations/_utils/fetch-server-organization'
 import { redirect } from 'next/navigation'
 import type { PropsWithChildren } from 'react'
 
@@ -12,22 +11,11 @@ export async function OrganizationGuard(props: PropsWithChildren) {
     redirect('/login')
   }
 
-  const selectedOrganizationId = (await cookies()).get(ORGANIZATION_COOKIE_NAME)?.value
-  const [orgsResp, selectedOrganizationResp] = await Promise.all([
-    apiServer.organizations.index.get({
-      query: {
-        limit: 1,
-        cursor: 0,
-      },
-    }),
-    selectedOrganizationId ? apiServer.organizations({ id: selectedOrganizationId }).get() : null,
-  ])
+  const organization = await fetchServerOrganization()
 
-  if (!selectedOrganizationResp?.data && !orgsResp.data?.data[0]) {
+  if (!organization) {
     redirect('/onboarding/organization')
   }
 
-  const organization = selectedOrganizationResp?.data ?? orgsResp.data?.data[0]!
-
   return <OrganizationProvider organization={organization}>{props.children}</OrganizationProvider>
 }

apps/app/src/app/(main)/organizations/_components/plan-guard.tsx

@@ -1,17 +1,20 @@
 import { apiServer } from '@bulkit/app/api/api.server'
-import { ORGANIZATION_COOKIE_NAME } from '@bulkit/app/app/(main)/organizations/organizations.constants'
-import { cookies } from 'next/headers'
+import { fetchServerOrganization } from '@bulkit/app/app/(main)/organizations/_utils/fetch-server-organization'
+import { buildOrganizationHeaders } from '@bulkit/app/app/(main)/organizations/_utils/organizations.utils'
 import { redirect } from 'next/navigation'
 import type { PropsWithChildren } from 'react'
 
 export async function PlanGuard(props: PropsWithChildren) {
-  const selectedOrgId = (await cookies()).get(ORGANIZATION_COOKIE_NAME)?.value
+  const organization = await fetchServerOrganization()
 
-  if (!selectedOrgId) {
+  if (!organization) {
     redirect('/onboarding/organization')
   }
 
-  const planResp = await apiServer.plans.active.get()
+  const planResp = await apiServer.plans.active.get({
+    headers: buildOrganizationHeaders(organization.id),
+  })
+
   if (!planResp.data) {
     redirect('/onboarding/plan')
   }

apps/app/src/app/(main)/organizations/_utils/fetch-server-organization.tsx

@@ -0,0 +1,18 @@
+import { apiServer } from '@bulkit/app/api/api.server'
+import { ORGANIZATION_COOKIE_NAME } from '@bulkit/app/app/(main)/organizations/organizations.constants'
+import { cookies } from 'next/headers'
+
+export async function fetchServerOrganization() {
+  const selectedOrganizationId = (await cookies()).get(ORGANIZATION_COOKIE_NAME)?.value
+  const [orgsResp, selectedOrganizationResp] = await Promise.all([
+    apiServer.organizations.index.get({
+      query: {
+        limit: 1,
+        cursor: 0,
+      },
+    }),
+    selectedOrganizationId ? apiServer.organizations({ id: selectedOrganizationId }).get() : null,
+  ])
+
+  return selectedOrganizationResp?.data ?? orgsResp.data?.data[0]!
+}

apps/app/src/app/onboarding/plan/_components/plan-selection-form.tsx

@@ -1,8 +1,8 @@
 'use client'
 
 import { apiClient } from '@bulkit/app/api/api.client'
+import { buildOrganizationHeaders } from '@bulkit/app/app/(main)/organizations/_utils/organizations.utils'
 import { PlanSelection } from '@bulkit/app/app/onboarding/organization/_components/plan-selection'
-import { env } from '@bulkit/app/env'
 import type { AvailablePlan } from '@bulkit/shared/modules/plans/plans.schemas'
 import { Button } from '@bulkit/ui/components/ui/button'
 import { Form } from '@bulkit/ui/components/ui/form'
@@ -43,10 +43,15 @@ export function PlanSelectionForm(props: PlanSelectionFormProps) {
         throw new Error('Selected plan not found')
       }
 
-      const resp = await apiClient.plans.checkout.post({
-        planId: selectedPlan.id,
-        variantId: data.variantId,
-      })
+      const resp = await apiClient.plans.checkout.post(
+        {
+          planId: selectedPlan.id,
+          variantId: data.variantId,
+        },
+        {
+          headers: buildOrganizationHeaders(props.organizationId),
+        }
+      )
 
       if (resp.error) {
         throw new Error(resp.error.value.message)

apps/app/src/app/onboarding/plan/page.tsx

@@ -1,7 +1,7 @@
 import { apiServer } from '@bulkit/app/api/api.server'
-import { ORGANIZATION_COOKIE_NAME } from '@bulkit/app/app/(main)/organizations/organizations.constants'
+import { fetchServerOrganization } from '@bulkit/app/app/(main)/organizations/_utils/fetch-server-organization'
+import { buildOrganizationHeaders } from '@bulkit/app/app/(main)/organizations/_utils/organizations.utils'
 import { PlanSelectionForm } from '@bulkit/app/app/onboarding/plan/_components/plan-selection-form'
-import { cookies } from 'next/headers'
 import { redirect } from 'next/navigation'
 
 export default async function PlanSelectionPage() {
@@ -11,24 +11,18 @@ export default async function PlanSelectionPage() {
     redirect('/login')
   }
 
-  const selectedOrgId = (await cookies()).get(ORGANIZATION_COOKIE_NAME)?.value
+  const organization = await fetchServerOrganization()
 
-  if (!selectedOrgId) {
+  if (!organization) {
     // just redirect to the home page,
     // the organization guard will handle all the organization setter logic
-    redirect('/')
-  }
-
-  // Get organization data
-  const orgResp = await apiServer.organizations({ id: selectedOrgId! }).get()
-  if (!orgResp.data) {
-    // just redirect to the home page,
-    // the organization guard will handle all the organization setter logic
-    redirect('/')
+    redirect('/onboarding/organization')
   }
 
   // Check if organization already has an active plan
-  const activePlanResp = await apiServer.plans.active.get()
+  const activePlanResp = await apiServer.plans.active.get({
+    headers: buildOrganizationHeaders(organization.id),
+  })
   if (activePlanResp.data) {
     // Organization already has a plan, redirect to home
     redirect('/')
@@ -49,7 +43,7 @@ export default async function PlanSelectionPage() {
           </p>
         </div>
 
-        <PlanSelectionForm organizationId={orgResp.data.id} initialPlans={plansResponse.data} />
+        <PlanSelectionForm organizationId={organization.id} initialPlans={plansResponse.data} />
       </div>
     </main>
   )