refactor(console): save console redirect after sign-in (#4180)

* refactor(console): save console redirect after sign-in

* test: add integration test for saved redirect

packages/console/src/components/SessionExpired/index.tsx

@@ -2,16 +2,20 @@ import { useLogto } from '@logto/react';
 import { useEffect } from 'react';
 
 import { getCallbackUrl, getUserTenantId } from '@/consts';
+import { saveRedirect } from '@/utils/storage';
 
 import AppLoading from '../AppLoading';
 
 /** This component shows a loading indicator and tries to sign in again. */
 function SessionExpired() {
-  const { signIn } = useLogto();
+  const { signIn, isLoading } = useLogto();
 
   useEffect(() => {
-    void signIn(getCallbackUrl(getUserTenantId()).href);
-  }, [signIn]);
+    if (!isLoading) {
+      saveRedirect();
+      void signIn(getCallbackUrl(getUserTenantId()).href);
+    }
+  }, [signIn, isLoading]);
 
   return <AppLoading />;
 }

packages/console/src/consts/storage.ts

@@ -2,7 +2,11 @@ export type CamelCase<T> = T extends `${infer A}_${infer B}`
   ? `${A}${Capitalize<CamelCase<B>>}`
   : T;
 
-export type StorageType = 'appearance_mode' | 'linking_social_connector' | 'checkout_session';
+export type StorageType =
+  | 'appearance_mode'
+  | 'linking_social_connector'
+  | 'checkout_session'
+  | 'redirect_after_sign_in';
 
 export const getStorageKey = <T extends StorageType>(forType: T) =>
   `logto:admin_console:${forType}` as const;
@@ -11,4 +15,6 @@ export const storageKeys = Object.freeze({
   appearanceMode: getStorageKey('appearance_mode'),
   linkingSocialConnector: getStorageKey('linking_social_connector'),
   checkoutSession: getStorageKey('checkout_session'),
+  /** The react-router redirect location after sign in. The value should be a stringified Location object. */
+  redirectAfterSignIn: getStorageKey('redirect_after_sign_in'),
 } satisfies Record<CamelCase<StorageType>, string>);

packages/console/src/containers/ProtectedRoutes/index.tsx

@@ -7,6 +7,7 @@ import { useCloudApi } from '@/cloud/hooks/use-cloud-api';
 import AppLoading from '@/components/AppLoading';
 import { searchKeys, getCallbackUrl } from '@/consts';
 import { TenantsContext } from '@/contexts/TenantsProvider';
+import { saveRedirect } from '@/utils/storage';
 
 /**
  * The container for all protected routes. It renders `<AppLoading />` when the user is not
@@ -35,6 +36,7 @@ export default function ProtectedRoutes() {
 
   useEffect(() => {
     if (!isLoading && !isAuthenticated) {
+      saveRedirect();
       const isSignUpMode = yes(searchParameters.get(searchKeys.signUp));
       void signIn(getCallbackUrl(currentTenantId).href, conditional(isSignUpMode && 'signUp'));
     }

packages/console/src/containers/TenantAccess/index.tsx

@@ -53,7 +53,7 @@ export default function TenantAccess() {
       // Test fetching an access token for the current Tenant ID.
       // If failed, it means the user finishes the first auth, ands still needs to auth again to
       // fetch the full-scoped (with all available tenants) token.
-      if (await trySafe(async () => getAccessToken(indicator))) {
+      if (await trySafe(getAccessToken(indicator))) {
         setCurrentTenantStatus('validated');
       }
       // If failed, it will be treated as a session expired error, and will be handled by the

packages/console/src/pages/Callback/index.tsx

@@ -3,20 +3,35 @@ import { useNavigate } from 'react-router-dom';
 
 import AppLoading from '@/components/AppLoading';
 import { getUserTenantId } from '@/consts';
+import { consumeSavedRedirect } from '@/utils/storage';
 import { isInFirstLevelCallback } from '@/utils/url';
 
+/** The global callback page for all sign-in redirects from Logto main flow. */
 function Callback() {
   const navigate = useNavigate();
 
   useHandleSignInCallback(() => {
+    const saved = consumeSavedRedirect();
+
+    if (saved) {
+      const { pathname, ...rest } = saved;
+      // Remove the first two path segments since we are using `basename` for tenant-specific
+      // routes (the first segment is empty because of the leading slash).
+      // For example, `/:tenantId/applications` will be `/applications` after removing.
+      // Once we merge all the routes into one router, we can remove this implementation.
+      const [_, __, ...segments] = pathname?.split('/') ?? [];
+      navigate({ ...rest, pathname: '/' + segments.join('/') });
+      return;
+    }
+
     /**
      * The first level callback check is due to the usage of `basename`
      * for tenant-specific routes, e.g., `/:tenantId/applications`.
      * Once we merge all the routes into one router, we can remove this check.
      */
-    navigate(isInFirstLevelCallback() ? `/${getUserTenantId()}` : '/', {
-      replace: true,
-    });
+    const defaultTo = isInFirstLevelCallback() ? `/${getUserTenantId()}` : '/';
+
+    navigate(defaultTo, { replace: true });
   });
 
   return <AppLoading />;

packages/console/src/pages/Welcome/index.tsx

@@ -18,7 +18,7 @@ function Welcome() {
   const theme = useTheme();
 
   useEffect(() => {
-    // If Authenticated, navigate to the Admin Console root page. directly
+    // If authenticated, navigate to the console root page directly
     if (isAuthenticated) {
       navigate('/');
     }

packages/console/src/utils/storage.ts

@@ -0,0 +1,36 @@
+import { trySafe } from '@silverhand/essentials';
+import { type Path } from 'react-router-dom';
+import { z } from 'zod';
+
+import { storageKeys } from '@/consts';
+
+/** Save the current url to session storage so that we can redirect to it after sign in. */
+export const saveRedirect = () => {
+  const { pathname, search, hash } = window.location;
+  sessionStorage.setItem(
+    storageKeys.redirectAfterSignIn,
+    JSON.stringify({ pathname, search, hash })
+  );
+};
+
+const partialPathGuard = z
+  .object({
+    pathname: z.string(),
+    search: z.string(),
+    hash: z.string(),
+  })
+  .partial() satisfies z.ZodType<Partial<Path>>;
+
+/** Get the saved redirect url from session storage and remove it. */
+export const consumeSavedRedirect = () => {
+  const saved = trySafe(() =>
+    partialPathGuard.parse(
+      JSON.parse(sessionStorage.getItem(storageKeys.redirectAfterSignIn) ?? '')
+    )
+  );
+
+  if (saved) {
+    sessionStorage.removeItem(storageKeys.redirectAfterSignIn);
+  }
+  return saved;
+};

packages/integration-tests/src/tests/ui/bootstrap.test.ts

@@ -1,3 +1,4 @@
+import { appendPath } from '@silverhand/essentials';
 import { setDefaultOptions } from 'expect-puppeteer';
 
 import {
@@ -62,8 +63,10 @@ describe('smoke testing for console admin account creation and sign-in', () => {
     expect(page.url()).toBe(new URL('sign-in', logtoConsoleUrl).href);
   });
 
-  it('can sign in to admin console', async () => {
-    await page.goto(logtoConsoleUrl.href);
+  it('can sign in to admin console again', async () => {
+    const initialHref = appendPath(logtoConsoleUrl, 'console', 'applications').href;
+    // Should be able to redirect back after sign-in
+    await page.goto(initialHref);
     await page.waitForNavigation({ waitUntil: 'networkidle0' });
     await expect(page).toFillForm('form', {
       identifier: consoleUsername,
@@ -72,7 +75,7 @@ describe('smoke testing for console admin account creation and sign-in', () => {
     await expect(page).toClick('button[name=submit]');
     await page.waitForNavigation({ waitUntil: 'networkidle0' });
 
-    expect(page.url()).toBe(new URL('console/get-started', logtoConsoleUrl).href);
+    expect(page.url()).toBe(initialHref);
 
     await expect(page).toClick('div[class$=topbar] > div:last-child');