From 77e3a1fcecd84c6a676ba22af25c009e67d0e7dc Mon Sep 17 00:00:00 2001
From: James Bhattarai <jamesbhattarai14@gmail.com>
Date: Fri, 6 Mar 2026 09:39:55 +0545
Subject: [PATCH] Fix: 2FA/passkey page redirect + org Setup context

---
 src/pages/auth/LoginPage.tsx    | 6 +++++-
 src/pages/auth/OrgSetupPage.tsx | 3 +++
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/src/pages/auth/LoginPage.tsx b/src/pages/auth/LoginPage.tsx
index 27b1c9a..d95d7d1 100644
--- a/src/pages/auth/LoginPage.tsx
+++ b/src/pages/auth/LoginPage.tsx
@@ -49,7 +49,7 @@ async function completeOidcFlow(oidcSessionId: string, token: string): Promise<s
 }
 
 export default function LoginPage() {
-  const { login, verifyTotp, refreshUser, user, isLoading: authLoading } = useAuth();
+  const { login, verifyTotp, refreshUser, user, isLoading: authLoading, checkOrgAdmin } = useAuth();
   const navigate = useNavigate();
   const { toast } = useToast();
   const [searchParams] = useSearchParams();
@@ -230,6 +230,7 @@ export default function LoginPage() {
           finishCliFlow(response.token);
         } else {
           await refreshUser();
+          await checkOrgAdmin();
           navigate('/profile');
         }
       } else {
@@ -347,6 +348,8 @@ export default function LoginPage() {
 
       // Token is stored by completeLogin, refresh user and navigate
       await refreshUser();
+      await checkOrgAdmin();
+      
       if (oidcSessionId) {
         const token = tokenManager.getToken();
         if (token) await finishOidcFlow(token);
@@ -430,6 +433,7 @@ export default function LoginPage() {
         if (token) finishCliFlow(token);
       } else {
         await refreshUser();
+        await checkOrgAdmin();
         navigate('/profile');
         toast({
           title: "Welcome back",
diff --git a/src/pages/auth/OrgSetupPage.tsx b/src/pages/auth/OrgSetupPage.tsx
index b303442..1bf3060 100644
--- a/src/pages/auth/OrgSetupPage.tsx
+++ b/src/pages/auth/OrgSetupPage.tsx
@@ -8,6 +8,7 @@
  */
 import { useState, useEffect } from "react";
 import { useNavigate, useLocation } from "react-router-dom";
+import { useQueryClient } from "@tanstack/react-query";
 import { Building2, Plus, ArrowRight, Loader2, Mail, ChevronDown, ChevronUp } from "lucide-react";
 import { Button } from "@/components/ui/button";
 import { Input } from "@/components/ui/input";
@@ -33,6 +34,7 @@ interface LocationState {
 export default function OrgSetupPage() {
   const navigate = useNavigate();
   const location = useLocation();
+  const queryClient = useQueryClient();
   const { refreshUser, checkOrgAdmin, isOrgMember, isLoading } = useAuth();
 
   // If the user already belongs to an org (e.g. they bookmarked /org-setup),
@@ -100,6 +102,7 @@ export default function OrgSetupPage() {
   const done = async () => {
     await refreshUser();
     await checkOrgAdmin();
+    queryClient.invalidateQueries({ queryKey: ['organizations'] });
     navigate("/profile", { replace: true });
   };