From a4cc5d69526c2584eb2e07393ea8218c8f2e49a3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=89mile=20R=C3=A9?= <emile@getprobo.com>
Date: Thu, 12 Feb 2026 10:08:24 +0400
Subject: [PATCH] Do not set continue when not needed + infer organizationId to
 assume after pw sign in
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Émile Ré <emile@getprobo.com>
---
 .../components/OrganizationErrorBoundary.tsx    | 13 +++++++------
 .../src/components/RootErrorBoundary.tsx        | 10 +++++++---
 .../iam/auth/sign-in/PasswordSignInPage.tsx     | 17 +++++++++++------
 3 files changed, 25 insertions(+), 15 deletions(-)

diff --git a/apps/console/src/components/OrganizationErrorBoundary.tsx b/apps/console/src/components/OrganizationErrorBoundary.tsx
index 9d6e0b6a9..ee023d349 100644
--- a/apps/console/src/components/OrganizationErrorBoundary.tsx
+++ b/apps/console/src/components/OrganizationErrorBoundary.tsx
@@ -1,5 +1,5 @@
 import { AssumptionRequiredError, UnAuthenticatedError } from "@probo/relay";
-import { Navigate, useRouteError } from "react-router";
+import { Navigate, useLocation, useRouteError } from "react-router";
 
 import { useOrganizationId } from "#/hooks/useOrganizationId";
 
@@ -8,17 +8,18 @@ import { PageError } from "./PageError";
 export function OrganizationErrorBoundary() {
   const error = useRouteError();
   const organizationId = useOrganizationId();
+  const location = useLocation();
 
-  const search = new URLSearchParams([
-    ["organization-id", organizationId],
-  ]);
+  const search = new URLSearchParams();
 
-  if (window.location.href !== window.location.origin) {
+  if (location.pathname !== "/" || location.search !== "") {
     search.set("continue", window.location.href);
   }
 
+  const queryString = search.toString();
+
   if (error instanceof UnAuthenticatedError) {
-    return <Navigate to={{ pathname: "/auth/login", search: "?" + search.toString() }} />;
+    return <Navigate to={{ pathname: "/auth/login", search: queryString ? "?" + queryString : "" }} />;
   }
 
   if (error instanceof AssumptionRequiredError) {
diff --git a/apps/console/src/components/RootErrorBoundary.tsx b/apps/console/src/components/RootErrorBoundary.tsx
index 1f2354559..727a9b210 100644
--- a/apps/console/src/components/RootErrorBoundary.tsx
+++ b/apps/console/src/components/RootErrorBoundary.tsx
@@ -1,21 +1,25 @@
 import { UnAuthenticatedError } from "@probo/relay";
-import { Navigate, useRouteError } from "react-router";
+import { Navigate, useLocation, useRouteError } from "react-router";
 
 import { PageError } from "./PageError";
 
 export function RootErrorBoundary() {
   const error = useRouteError();
+  const location = useLocation();
 
   const search = new URLSearchParams();
-  if (window.location.href !== window.location.origin) {
+
+  if (location.pathname !== "/" || location.search !== "") {
     search.set("continue", window.location.href);
   }
 
+  const queryString = search.toString();
+
   if (error instanceof UnAuthenticatedError) {
     return (
       <Navigate to={{
         pathname: "/auth/login",
-        search: search.toString() ? "?" + search.toString() : "",
+        search: queryString ? "?" + queryString : "",
       }}
       />
     );
diff --git a/apps/console/src/pages/iam/auth/sign-in/PasswordSignInPage.tsx b/apps/console/src/pages/iam/auth/sign-in/PasswordSignInPage.tsx
index ee561f996..651ff8333 100644
--- a/apps/console/src/pages/iam/auth/sign-in/PasswordSignInPage.tsx
+++ b/apps/console/src/pages/iam/auth/sign-in/PasswordSignInPage.tsx
@@ -3,7 +3,7 @@ import { useTranslate } from "@probo/i18n";
 import { Button, Field, IconChevronLeft, useToast } from "@probo/ui";
 import type { FormEventHandler } from "react";
 import { useMutation } from "react-relay";
-import { Link, useLocation, useSearchParams } from "react-router";
+import { Link, matchPath, useLocation, useSearchParams } from "react-router";
 import { graphql } from "relay-runtime";
 
 import type { PasswordSignInPageMutation } from "#/__generated__/iam/PasswordSignInPageMutation.graphql";
@@ -37,7 +37,7 @@ export default function PasswordSignInPage() {
     if (!emailValue || !passwordValue) return;
 
     const continueUrlParam = searchParams.get("continue");
-    let safeContinueUrl: string;
+    let safeContinueUrl: URL;
     if (continueUrlParam) {
       let continueUrl: URL;
       try {
@@ -45,18 +45,23 @@ export default function PasswordSignInPage() {
       } catch {
         continueUrl = new URL(window.location.origin);
       }
-      safeContinueUrl = window.location.origin + continueUrl.pathname + continueUrl.search;
+      safeContinueUrl = new URL(continueUrl.pathname + continueUrl.search, window.location.origin);
     } else {
-      safeContinueUrl = window.location.origin;
+      safeContinueUrl = new URL(window.location.origin);
     }
 
+    const match = matchPath(
+      { path: "/organizations/:organizationId", caseSensitive: false, end: false },
+      safeContinueUrl.pathname,
+    );
+
     signIn({
       variables: {
         input: {
           email: emailValue,
           password: passwordValue,
           // Assume when signing in
-          organizationId: searchParams.get("organization-id"),
+          organizationId: match && match.params.organizationId,
         },
       },
       onCompleted: (_, error) => {
@@ -73,7 +78,7 @@ export default function PasswordSignInPage() {
           return;
         }
 
-        window.location.href = safeContinueUrl;
+        window.location.href = safeContinueUrl.href;
       },
       onError: (e) => {
         toast({