add redirect persistence for SSO

Author: Cazden

src/components/auth/form/login.tsx

@@ -20,6 +20,7 @@ import { useSession } from "@/hooks/use-session"
 import { cn } from "@/lib/utils"
 import { toast } from "@/lib/toast"
 import { truncator } from "@/lib/truncate"
+import { setPendingRedirect } from "@/lib/auth"
 import { getRecentAccounts } from "@/lib/accounts"
 
 import * as Auth from "@/components/auth"
@@ -591,11 +592,13 @@ function SsoStep({
   onCancel: () => void
 }) {
   const [loading, setLoading] = useState(false)
+  const { redirect } = useSearch({ from: "/$accountId/auth/login" })
 
   function onContinue() {
     if (!redirectUrl) return
 
     setLoading(true)
+    if (redirect) setPendingRedirect(redirect)
     window.location.href = redirectUrl
   }
 

src/lib/auth.ts

@@ -19,12 +19,27 @@ interface ResetToken {
   token: string
 }
 
+const REDIRECT_STORAGE_KEY = "keygen.auth.redirect"
+
 export function parseRedirect(value: unknown): string | undefined {
   return typeof value === "string" && /^\/goto(?:[/?#]|$)/.test(value)
     ? value
     : undefined
 }
 
+export function setPendingRedirect(value: string): void {
+  window.sessionStorage.setItem(REDIRECT_STORAGE_KEY, value)
+}
+
+export function takePendingRedirect(): string | undefined {
+  const value = parseRedirect(
+    window.sessionStorage.getItem(REDIRECT_STORAGE_KEY),
+  )
+  window.sessionStorage.removeItem(REDIRECT_STORAGE_KEY)
+
+  return value
+}
+
 // parse a reset token from the URL query string
 export function parseResetToken(
   value: string | undefined | null,

src/routes/$accountId/index.tsx

@@ -1,17 +1,25 @@
 import { createFileRoute, redirect } from "@tanstack/react-router"
 
 import { restoreSession } from "@/keygen/session"
+import { takePendingRedirect } from "@/lib/auth"
 
 export const Route = createFileRoute("/$accountId/")({
   loader: async ({ params }) => {
-    // SSO returns the browser to the slug-addressed account root. Resolve the
+    // SSO returns the browser to the slug-addressed account root; resolve the
     // session so we can land on the account's stable id rather than the slug.
     const { accountId } = await restoreSession()
 
-    return redirect({
+    const pendingRedirect = takePendingRedirect()
+    if (pendingRedirect) {
+      redirect({ href: pendingRedirect, replace: true, throw: true })
+      return
+    }
+
+    redirect({
       to: "/$accountId/app",
       params: { accountId: accountId ?? params.accountId },
       replace: true,
+      throw: true,
     })
   },
   component: () => null,