Implement redirect path for saml

Signed-off-by: Émile Ré <emile@getprobo.com>

apps/console/src/__generated__/iam/AssumePageMutation.graphql.ts

@@ -1,5 +1,5 @@
 /**
- * @generated SignedSource<<096c8befe586a961b8ebf410eda0e019>>
+ * @generated SignedSource<<8717e91a7e1602c06f6c67fa76580ecf>>
  * @lightSyntaxTransform
  * @nogrep
  */
@@ -12,6 +12,7 @@ import { ConcreteRequest } from 'relay-runtime';
 export type ReauthenticationReason = "POLICY_REQUIREMENT" | "SENSITIVE_ACTION" | "SESSION_EXPIRED";
 export type AssumeOrganizationSessionInput = {
   organizationId: string;
+  redirectPath: string;
 };
 export type AssumePageMutation$variables = {
   input: AssumeOrganizationSessionInput;

apps/console/src/__generated__/iam/MembershipsDropdownMenuItem_assumeMutation.graphql.ts

@@ -1,5 +1,5 @@
 /**
- * @generated SignedSource<<4beff809268fb32711cc02699683b623>>
+ * @generated SignedSource<<f829648ef4b9a5c0b8c4d62ac844bca0>>
  * @lightSyntaxTransform
  * @nogrep
  */
@@ -12,6 +12,7 @@ import { ConcreteRequest } from 'relay-runtime';
 export type ReauthenticationReason = "POLICY_REQUIREMENT" | "SENSITIVE_ACTION" | "SESSION_EXPIRED";
 export type AssumeOrganizationSessionInput = {
   organizationId: string;
+  redirectPath: string;
 };
 export type MembershipsDropdownMenuItem_assumeMutation$variables = {
   input: AssumeOrganizationSessionInput;

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, useSearchParams } from "react-router";
+import { Link, useLocation, useSearchParams } from "react-router";
 import { graphql } from "relay-runtime";
 
 import type { PasswordSignInPageMutation } from "#/__generated__/iam/PasswordSignInPageMutation.graphql";
@@ -19,6 +19,7 @@ const signInMutation = graphql`
 `;
 
 export default function PasswordSignInPage() {
+  const location = useLocation();
   const [searchParams] = useSearchParams();
 
   const { __ } = useTranslate();
@@ -72,7 +73,7 @@ export default function PasswordSignInPage() {
   return (
     <form className="space-y-6 w-full max-w-md mx-auto pt-4" onSubmit={handlePasswordLogin}>
       <Link
-        to="/auth/login"
+        to={{ pathname: "/auth/login", search: location.search }}
         className="flex items-center gap-2 text-txt-secondary hover:text-txt-primary transition-colors mb-4"
       >
         <IconChevronLeft size={20} />
@@ -112,7 +113,7 @@ export default function PasswordSignInPage() {
       <div className="text-center mt-6 text-sm text-txt-secondary">
         {__("Don't have an account ?")}
         {" "}
-        <Link to="/auth/register" className="underline hover:text-txt-primary">
+        <Link to={{ pathname: "/auth/register", search: location.search }} className="underline hover:text-txt-primary">
           {__("Register")}
         </Link>
       </div>

apps/console/src/pages/iam/auth/sign-in/SSOSignInPage.tsx

@@ -6,7 +6,7 @@ import {
   usePreloadedQuery,
   useQueryLoader,
 } from "react-relay";
-import { Link, useNavigate } from "react-router";
+import { Link, useLocation, useNavigate, useSearchParams } from "react-router";
 import { graphql } from "relay-runtime";
 
 import type { SSOSignInPageQuery } from "#/__generated__/iam/SSOSignInPageQuery.graphql";
@@ -18,6 +18,7 @@ const ssoAvailabilityQuery = graphql`
 `;
 
 export default function SSOSignInPage() {
+  const location = useLocation();
   const { __ } = useTranslate();
 
   const [queryRef, loadQuery]
@@ -39,7 +40,7 @@ export default function SSOSignInPage() {
     <>
       <form className="space-y-6 w-full max-w-md mx-auto pt-4" onSubmit={handleSSOCheck}>
         <Link
-          to="/auth/login"
+          to={{ pathname: "/auth/register", search: location.search }}
           className="flex items-center gap-2 text-txt-secondary hover:text-txt-primary transition-colors mb-4"
         >
           <IconChevronLeft size={20} />
@@ -70,7 +71,7 @@ export default function SSOSignInPage() {
           {__("Don't have an account ?")}
           {" "}
           <Link
-            to="/auth/register"
+            to={{ pathname: "/auth/register", search: location.search }}
             className="underline hover:text-txt-primary"
           >
             {__("Register")}
@@ -96,6 +97,7 @@ function NavigateToSSOLoginURL(props: {
 
   const { __ } = useTranslate();
   const { toast } = useToast();
+  const [searchParams] = useSearchParams();
   const navigate = useNavigate();
 
   const { ssoLoginURL } = usePreloadedQuery<SSOSignInPageQuery>(
@@ -127,8 +129,11 @@ function NavigateToSSOLoginURL(props: {
       return;
     }
 
-    window.location.href = ssoLoginURL.value;
-  }, [__, navigate, ssoLoginURL, toast]);
+    const url = new URL(ssoLoginURL.value);
+    url.search = "?" + searchParams.toString();
+
+    window.location.href = url.toString();
+  }, [__, navigate, ssoLoginURL, toast, searchParams]);
 
   return null;
 }

apps/console/src/pages/iam/organizations/AssumePage.tsx

@@ -38,19 +38,18 @@ function AssumePageInner() {
 
   const [assumeOrganizationSession] = useMutation<AssumePageMutation>(assumeMutation);
 
+  const redirectPath = searchParams.get("redirect-path") ?? `/organizations/${organizationId}`;
+
   useEffect(() => {
     assumeOrganizationSession({
       variables: {
-        input: { organizationId },
+        input: { organizationId, redirectPath },
       },
       onError: (error) => {
         if (error instanceof UnAuthenticatedError) {
           const search = new URLSearchParams([
             ["organization-id", organizationId],
-            [
-              "redirect-path",
-              searchParams.get("redirect-path") ?? window.location.pathname + window.location.search,
-            ],
+            ["redirect-path", redirectPath],
           ]);
 
           void navigate({ pathname: "/auth/login", search: "?" + search.toString() });
@@ -64,26 +63,27 @@ function AssumePageInner() {
 
         const { result } = assumeOrganizationSession;
         const search = new URLSearchParams();
+        let samlSSOLoginURL: URL;
 
         switch (result.__typename) {
           case "PasswordRequired":
             search.set("organization-id", organizationId);
-            search.set(
-              "redirect-path",
-              searchParams.get("redirect-path") ?? window.location.pathname + window.location.search,
-            );
+            search.set("redirect-path", redirectPath);
 
             void navigate({ pathname: "/auth/passord-login", search: "?" + search.toString() });
             break;
           case "SAMLAuthenticationRequired":
-            window.location.href = result.redirectUrl;
+            samlSSOLoginURL = new URL(result.redirectUrl);
+            samlSSOLoginURL.search = "?" + searchParams.toString();
+
+            window.location.href = samlSSOLoginURL.toString();
             break;
           default:
-            void navigate(searchParams.get("redirect-path") ?? window.location.pathname + window.location.search);
+            void navigate(redirectPath);
         }
       },
     });
-  }, [organizationId, navigate, assumeOrganizationSession, searchParams]);
+  }, [organizationId, navigate, assumeOrganizationSession, redirectPath, searchParams]);
 
   return (
     <AuthLayout>