feat: enhance invitation handling and user redirection

- Implemented server-side redirection for users accepting invitations, improving user experience by directing them to the appropriate organization page.
- Added logic to check for pending invitations in the layout and root page, ensuring users are redirected to the invite page if applicable.
- Updated the invite page to handle email mismatches, prompting users to switch accounts if their signed-in email does not match the invited email.
- Refactored the AcceptInvite component to utilize redirection after accepting an invitation, streamlining the flow for users.

Author: Marfuen

apps/app/src/actions/organization/accept-invitation.ts

@@ -2,6 +2,7 @@
 
 import { db } from '@db';
 import { revalidatePath, revalidateTag } from 'next/cache';
+import { redirect } from 'next/navigation';
 import { Resend } from 'resend';
 import { z } from 'zod';
 import { authActionClientWithoutOrg } from '../safe-action';
@@ -88,13 +89,8 @@ export const completeInvitation = authActionClientWithoutOrg
             },
           });
 
-          return {
-            success: true,
-            data: {
-              accepted: true,
-              organizationId: invitation.organizationId,
-            },
-          };
+          // Server redirect to the organization's root
+          redirect(`/${invitation.organizationId}/`);
         }
 
         if (!invitation.role) {
@@ -148,13 +144,8 @@ export const completeInvitation = authActionClientWithoutOrg
         revalidatePath(`/${invitation.organization.id}/settings/users`);
         revalidateTag(`user_${user.id}`);
 
-        return {
-          success: true,
-          data: {
-            accepted: true,
-            organizationId: invitation.organizationId,
-          },
-        };
+        // Server redirect to the organization's root
+        redirect(`/${invitation.organizationId}/`);
       } catch (error) {
         console.error('Error accepting invitation:', error);
         throw new Error(error as string);

apps/app/src/app/(app)/invite/[code]/page.tsx

@@ -1,6 +1,7 @@
-import { getOrganizations } from '@/data/getOrganizations';
+import { OnboardingLayout } from '@/components/onboarding/OnboardingLayout';
+import { SignOut } from '@/components/sign-out';
 import { auth } from '@/utils/auth';
-import type { Organization } from '@db';
+import { Icons } from '@comp/ui/icons';
 import { db } from '@db';
 import { headers } from 'next/headers';
 import { notFound, redirect } from 'next/navigation';
@@ -10,6 +11,21 @@ interface InvitePageProps {
   params: Promise<{ code: string }>;
 }
 
+const maskEmail = (email: string) => {
+  const [local, domain] = email.split('@');
+  if (!domain) return email;
+  const maskedLocal =
+    local.length <= 2 ? `${local[0] ?? ''}***` : `${local[0]}***${local.slice(-1)}`;
+
+  const domainParts = domain.split('.');
+  if (domainParts.length === 0) return `${maskedLocal}@***`;
+  const tld = domainParts[domainParts.length - 1];
+  const secondLevel = domainParts.length >= 2 ? domainParts[domainParts.length - 2] : '';
+  const maskedSecondLevel = secondLevel ? `${secondLevel[0]}***` : '***';
+
+  return `${maskedLocal}@${maskedSecondLevel}.${tld}`;
+};
+
 export default async function InvitePage({ params }: InvitePageProps) {
   const { code } = await params;
   const session = await auth.api.getSession({
@@ -21,21 +37,10 @@ export default async function InvitePage({ params }: InvitePageProps) {
     return redirect(`/auth?inviteCode=${code}`);
   }
 
-  // Fetch existing organizations
-  let organizations: Organization[] = [];
-  try {
-    const result = await getOrganizations();
-    organizations = result.organizations;
-  } catch (error) {
-    // If user has no organizations, continue with empty array
-    console.error('Failed to fetch organizations:', error);
-  }
-
-  // Check if this invitation exists and is valid for this user
+  // Load invite by code, then verify email after
   const invitation = await db.invitation.findFirst({
     where: {
       id: code,
-      email: session.user.email,
       status: 'pending',
     },
     include: {
@@ -48,16 +53,53 @@ export default async function InvitePage({ params }: InvitePageProps) {
   });
 
   if (!invitation) {
-    // Either invitation doesn't exist, already accepted, or not for this user
     notFound();
   }
 
+  // If signed-in user email doesn't match the invited email, prompt to switch accounts
+  if (invitation.email !== session.user.email) {
+    return (
+      <OnboardingLayout variant="setup" currentOrganization={null}>
+        <div className="flex min-h-[calc(100dvh-80px)] w-full items-center justify-center p-4">
+          <div className="bg-card relative w-full max-w-[480px] rounded-sm border p-10 shadow-lg">
+            <div className="flex flex-col items-center gap-6 text-center">
+              <Icons.Logo />
+              <h1 className="text-2xl font-semibold tracking-tight">Wrong account</h1>
+              <div className="mx-auto max-w-[42ch] text-muted-foreground leading-relaxed flex flex-col gap-4">
+                <div className="space-y-2 text-sm">
+                  <p>
+                    You are signed in as
+                    <span className="mx-1 inline-flex items-center rounded-xs border border-muted bg-muted/40 px-2 py-0.5 text-sm">
+                      {session.user.email}
+                    </span>
+                  </p>
+                  <p>
+                    This invite is for
+                    <span className="mx-1 inline-flex items-center rounded-xs border border-muted bg-muted/40 px-2 py-0.5 text-sm">
+                      {maskEmail(invitation.email)}
+                    </span>
+                  </p>
+                </div>
+                <p className="text-base font-medium">
+                  To accept, sign out and sign back in with the invited email.
+                </p>
+              </div>
+              <SignOut asButton className="w-full" />
+            </div>
+          </div>
+        </div>
+      </OnboardingLayout>
+    );
+  }
+
   return (
-    <div className="flex flex-1 items-center justify-center p-4">
-      <AcceptInvite
-        inviteCode={invitation.id}
-        organizationName={invitation.organization.name || ''}
-      />
-    </div>
+    <OnboardingLayout variant="setup" currentOrganization={null}>
+      <div className="flex min-h-[calc(100dvh-80px)] w-full items-center justify-center p-4">
+        <AcceptInvite
+          inviteCode={invitation.id}
+          organizationName={invitation.organization.name || ''}
+        />
+      </div>
+    </OnboardingLayout>
   );
 }

apps/app/src/app/(app)/layout.tsx

@@ -0,0 +1,34 @@
+import { auth } from '@/utils/auth';
+import { db } from '@db';
+import { headers } from 'next/headers';
+import { redirect } from 'next/navigation';
+
+export default async function Layout({ children }: { children: React.ReactNode }) {
+  const hdrs = await headers();
+  const session = await auth.api.getSession({
+    headers: hdrs,
+  });
+
+  if (!session) {
+    return redirect('/auth');
+  }
+
+  const pendingInvite = await db.invitation.findFirst({
+    where: {
+      email: session.user.email,
+      status: 'pending',
+    },
+  });
+
+  if (pendingInvite) {
+    let path = hdrs.get('x-pathname') || hdrs.get('referer') || '';
+    // normalize potential locale prefix
+    path = path.replace(/\/([a-z]{2})\//, '/');
+    const target = `/invite/${pendingInvite.id}`;
+    if (!path.startsWith(target)) {
+      return redirect(target);
+    }
+  }
+
+  return <>{children}</>;
+}

apps/app/src/app/(app)/setup/components/accept-invite.tsx

@@ -7,7 +7,7 @@ import { Icons } from '@comp/ui/icons';
 import { Loader2 } from 'lucide-react';
 import { useAction } from 'next-safe-action/hooks';
 import Link from 'next/link';
-import { useRouter } from 'next/navigation';
+import { redirect } from 'next/navigation';
 import { toast } from 'sonner';
 
 export function AcceptInvite({
@@ -17,7 +17,7 @@ export function AcceptInvite({
   inviteCode: string;
   organizationName: string;
 }) {
-  const router = useRouter();
+  // Using next/navigation redirect to avoid showing invite page after accept
 
   const { execute, isPending } = useAction(completeInvitation, {
     onSuccess: async (result) => {
@@ -26,36 +26,39 @@ export function AcceptInvite({
         await authClient.organization.setActive({
           organizationId: result.data.data.organizationId,
         });
-        // Redirect to the organization's root path
-        router.push(`/${result.data.data.organizationId}/`);
       }
     },
     onError: (error) => {
       toast.error('Failed to accept invitation');
     },
   });
 
-  const handleAccept = () => {
-    execute({ inviteCode });
+  const handleAccept = async () => {
+    await execute({ inviteCode });
+    redirect(`/`);
   };
 
   return (
     <div className="bg-card relative w-full max-w-[440px] rounded-sm border p-8 shadow-lg">
-      <div className="mb-8 flex justify-between">
+      <div className="mb-8 flex justify-center">
         <Link href="/">
           <Icons.Logo />
         </Link>
       </div>
 
-      <div className="mb-8 space-y-2">
-        <h1 className="text-2xl font-semibold tracking-tight">You have been invited to join</h1>
-        <p className="text-xl font-medium line-clamp-1">{organizationName || 'an organization'}</p>
+      <div className="mb-8 space-y-1.5 text-center">
+        <h1 className="text-xs font-medium uppercase tracking-wider text-muted-foreground">
+          You have been invited to join
+        </h1>
+        <p className="text-2xl font-semibold tracking-tight line-clamp-1">
+          {organizationName || 'an organization'}
+        </p>
         <p className="text-muted-foreground text-sm">
           Please accept the invitation to join the organization.
         </p>
       </div>
 
-      <Button onClick={handleAccept} className="w-full" disabled={isPending}>
+      <Button onClick={handleAccept} className="w-full" size="sm" disabled={isPending}>
         {isPending ? (
           <>
             <Loader2 className="mr-2 h-4 w-4 animate-spin" />

apps/app/src/app/page.tsx

@@ -36,6 +36,18 @@ export default async function RootPage({
   const orgId = session.session.activeOrganizationId;
 
   if (!orgId) {
+    // If the user has no active org, check for pending invitations for this email
+    const pendingInvite = await db.invitation.findFirst({
+      where: {
+        email: session.user.email,
+        status: 'pending',
+      },
+    });
+
+    if (pendingInvite) {
+      return redirect(await buildUrlWithParams(`/invite/${pendingInvite.id}`));
+    }
+
     return redirect(await buildUrlWithParams('/setup'));
   }
 

apps/app/src/components/onboarding/OnboardingLayout.tsx

@@ -39,7 +39,7 @@ export async function OnboardingLayout({
         currentOrganization={currentOrganization}
         variant={variant}
       />
-      <div>{children}</div>
+      <div className="flex flex-1">{children}</div>
     </main>
   );
 }

apps/app/src/components/sign-out.tsx

@@ -6,7 +6,15 @@ import { DropdownMenuItem } from '@comp/ui/dropdown-menu';
 import { useRouter } from 'next/navigation';
 import { useState } from 'react';
 
-export function SignOut({ asButton = false }: { asButton?: boolean }) {
+export function SignOut({
+  asButton = false,
+  className = '',
+  size = 'sm',
+}: {
+  asButton?: boolean;
+  className?: string;
+  size?: 'default' | 'sm' | 'lg' | 'icon';
+}) {
   const router = useRouter();
   const [isLoading, setLoading] = useState(false);
 
@@ -22,7 +30,11 @@ export function SignOut({ asButton = false }: { asButton?: boolean }) {
   };
 
   if (asButton) {
-    return <Button onClick={handleSignOut}>{isLoading ? 'Loading...' : 'Sign out'}</Button>;
+    return (
+      <Button onClick={handleSignOut} className={className} size={size}>
+        {isLoading ? 'Loading...' : 'Sign out'}
+      </Button>
+    );
   }
 
   return (