From ad08ed84713908c5b11ba220c7fd130b3a0e8888 Mon Sep 17 00:00:00 2001
From: Piyush Gupta <piyushguptaa2z123@gmail.com>
Date: Mon, 24 Nov 2025 20:12:34 +0530
Subject: [PATCH] fix: layout token redirect

---
 backend/src/server/routes/v1/auth-router.ts                   | 4 ++--
 frontend/src/hooks/api/auth/types.ts                          | 1 +
 .../layouts/OrganizationLayout/components/NavBar/Navbar.tsx   | 2 ++
 frontend/src/pages/middlewares/restrict-login-signup.tsx      | 3 ++-
 4 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/backend/src/server/routes/v1/auth-router.ts b/backend/src/server/routes/v1/auth-router.ts
index 912bee642..16290f1e3 100644
--- a/backend/src/server/routes/v1/auth-router.ts
+++ b/backend/src/server/routes/v1/auth-router.ts
@@ -81,7 +81,8 @@ export const registerAuthRoutes = async (server: FastifyZodProvider) => {
       response: {
         200: z.object({
           token: z.string(),
-          organizationId: z.string().optional()
+          organizationId: z.string().optional(),
+          subOrganizationId: z.string().optional()
         })
       }
     },
@@ -131,7 +132,6 @@ export const registerAuthRoutes = async (server: FastifyZodProvider) => {
         appCfg.AUTH_SECRET,
         { expiresIn }
       );
-
       return { token, organizationId: decodedToken.organizationId, subOrganizationId: decodedToken.subOrganizationId };
     }
   });
diff --git a/frontend/src/hooks/api/auth/types.ts b/frontend/src/hooks/api/auth/types.ts
index fd9d6fb5c..d0ad86526 100644
--- a/frontend/src/hooks/api/auth/types.ts
+++ b/frontend/src/hooks/api/auth/types.ts
@@ -1,6 +1,7 @@
 export type GetAuthTokenAPI = {
   token: string;
   organizationId?: string;
+  subOrganizationId?: string;
 };
 
 export enum UserEncryptionVersion {
diff --git a/frontend/src/layouts/OrganizationLayout/components/NavBar/Navbar.tsx b/frontend/src/layouts/OrganizationLayout/components/NavBar/Navbar.tsx
index 25e36edef..49e5a66eb 100644
--- a/frontend/src/layouts/OrganizationLayout/components/NavBar/Navbar.tsx
+++ b/frontend/src/layouts/OrganizationLayout/components/NavBar/Navbar.tsx
@@ -250,6 +250,8 @@ export const Navbar = () => {
       subOrganizationId: subOrgId
     });
 
+    localStorage.setItem("orgData.id", subOrgId);
+
     if (isMfaEnabled) {
       SecurityClient.setMfaToken(token);
       if (mfaMethod) {
diff --git a/frontend/src/pages/middlewares/restrict-login-signup.tsx b/frontend/src/pages/middlewares/restrict-login-signup.tsx
index ec1958d14..582b6331b 100644
--- a/frontend/src/pages/middlewares/restrict-login-signup.tsx
+++ b/frontend/src/pages/middlewares/restrict-login-signup.tsx
@@ -117,9 +117,10 @@ export const Route = createFileRoute("/_restrict-login-signup")({
         return;
       throw redirect({ to: "/login/select-organization" });
     }
+    const orgId = data.subOrganizationId || data.organizationId;
     throw redirect({
       to: "/organizations/$orgId/projects",
-      params: { orgId: data.organizationId }
+      params: { orgId }
     });
   },
   component: AuthConsentWrapper