From ac34d6ee47733eb0baf27f155bb338748d23b4e6 Mon Sep 17 00:00:00 2001
From: Remington Breeze <remington@breeze.software>
Date: Thu, 9 May 2024 10:27:00 -0700
Subject: [PATCH] feat(ui): remember previous location when token refresh logs
 you out (#1986)

Signed-off-by: Remington Breeze <remington@breeze.software>
---
 ui/src/config/transport.ts   | 2 +-
 ui/src/pages/login/login.tsx | 7 +++++--
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/ui/src/config/transport.ts b/ui/src/config/transport.ts
index 9e815304a..5abdb7d29 100644
--- a/ui/src/config/transport.ts
+++ b/ui/src/config/transport.ts
@@ -7,7 +7,7 @@ import { paths } from './paths';
 
 const logout = () => {
   localStorage.removeItem(authTokenKey);
-  window.location.replace(paths.login);
+  window.location.replace(`${paths.login}?${redirectToQueryParam}=${window.location.pathname}`);
 };
 
 const renewToken = () => {
diff --git a/ui/src/pages/login/login.tsx b/ui/src/pages/login/login.tsx
index 34be77e6a..7464bbf62 100644
--- a/ui/src/pages/login/login.tsx
+++ b/ui/src/pages/login/login.tsx
@@ -1,7 +1,8 @@
 import { useQuery } from '@connectrpc/connect-query';
 import { Divider, Typography } from 'antd';
-import { Navigate } from 'react-router-dom';
+import { Navigate, generatePath, useSearchParams } from 'react-router-dom';
 
+import { redirectToQueryParam } from '@ui/config/auth';
 import { paths } from '@ui/config/paths';
 import { AdminLogin } from '@ui/features/auth/admin-login';
 import { useAuthContext } from '@ui/features/auth/context/use-auth-context';
@@ -13,10 +14,12 @@ import * as styles from './login.module.less';
 
 export const Login = () => {
   const { data, isLoading } = useQuery(getPublicConfig);
+  const [params] = useSearchParams();
   const { isLoggedIn } = useAuthContext();
+  const redirectTo = params.get(redirectToQueryParam);
 
   if (isLoggedIn) {
-    return <Navigate to={paths.home} replace />;
+    return <Navigate to={redirectTo ? generatePath(redirectTo) : paths.home} replace />;
   }
 
   return (