From 952e60f08afac6a946c4e60072f80629a15d239e Mon Sep 17 00:00:00 2001
From: x032205 <x032205@gmail.com>
Date: Wed, 4 Jun 2025 16:54:14 -0400
Subject: [PATCH] Select organization checkpoint

---
 frontend/src/pages/auth/SelectOrgPage/route.tsx            | 3 ++-
 .../src/pages/auth/SignUpInvitePage/SignUpInvitePage.tsx   | 2 +-
 frontend/src/pages/middlewares/restrict-login-signup.tsx   | 7 +++++--
 3 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/frontend/src/pages/auth/SelectOrgPage/route.tsx b/frontend/src/pages/auth/SelectOrgPage/route.tsx
index 445479b3f..27ad4bb94 100644
--- a/frontend/src/pages/auth/SelectOrgPage/route.tsx
+++ b/frontend/src/pages/auth/SelectOrgPage/route.tsx
@@ -7,7 +7,8 @@ import { SelectOrganizationPage } from "./SelectOrgPage";
 export const SelectOrganizationPageQueryParams = z.object({
   org_id: z.string().optional().catch(""),
   callback_port: z.coerce.number().optional().catch(undefined),
-  is_admin_login: z.boolean().optional().catch(false)
+  is_admin_login: z.boolean().optional().catch(false),
+  force: z.boolean().optional()
 });
 
 export const Route = createFileRoute("/_restrict-login-signup/login/select-organization")({
diff --git a/frontend/src/pages/auth/SignUpInvitePage/SignUpInvitePage.tsx b/frontend/src/pages/auth/SignUpInvitePage/SignUpInvitePage.tsx
index 19a2694f8..2db4f7901 100644
--- a/frontend/src/pages/auth/SignUpInvitePage/SignUpInvitePage.tsx
+++ b/frontend/src/pages/auth/SignUpInvitePage/SignUpInvitePage.tsx
@@ -243,7 +243,7 @@ export const SignupInvitePage = () => {
                   SecurityClient.setSignupToken(response.token);
                   setStep(2);
                 } else {
-                  navigate({ to: "/login" });
+                  navigate({ to: "/login/select-organization", search: { force: true } });
 
                   // TODO: Redirect user to the newly joined org. The commented code does not work due to auth issues.
                   // const redirectExistingUser = async () => {
diff --git a/frontend/src/pages/middlewares/restrict-login-signup.tsx b/frontend/src/pages/middlewares/restrict-login-signup.tsx
index 31434458f..210f848cc 100644
--- a/frontend/src/pages/middlewares/restrict-login-signup.tsx
+++ b/frontend/src/pages/middlewares/restrict-login-signup.tsx
@@ -15,7 +15,8 @@ import { setAuthToken } from "@app/hooks/api/reactQuery";
 import { ProjectType } from "@app/hooks/api/workspace/types";
 
 const QueryParamsSchema = z.object({
-  callback_port: z.coerce.number().optional().catch(undefined)
+  callback_port: z.coerce.number().optional().catch(undefined),
+  force: z.boolean().optional()
 });
 
 export const AuthConsentWrapper = () => {
@@ -71,7 +72,7 @@ export const AuthConsentWrapper = () => {
 export const Route = createFileRoute("/_restrict-login-signup")({
   validateSearch: zodValidator(QueryParamsSchema),
   search: {
-    middlewares: [stripSearchParams({ callback_port: undefined })]
+    middlewares: [stripSearchParams({ callback_port: undefined, force: undefined })]
   },
   beforeLoad: async ({ context, location, search }) => {
     if (location.pathname === "/signupinvite") return;
@@ -93,6 +94,8 @@ export const Route = createFileRoute("/_restrict-login-signup")({
 
     setAuthToken(data.token);
 
+    if (location.pathname.endsWith("select-organization") && search?.force === true) return;
+
     // to do cli login
     if (search?.callback_port) {
       if (location.pathname.endsWith("select-organization") || location.pathname.endsWith("login"))