From b100c6274ea8aafe080481ce5bc45bdaa59480ed Mon Sep 17 00:00:00 2001
From: Luke Vella <me@lukevella.com>
Date: Thu, 24 Apr 2025 11:49:24 +0100
Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9B=20Fix=20login/registration=20flow?=
 =?UTF-8?q?=20(#1688)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../register/components/register-name-form.tsx       |  8 ++++++--
 apps/web/src/components/poll/poll-header.tsx         | 12 ++++++++++--
 apps/web/tests/guest-to-user-migration.spec.ts       |  4 ----
 3 files changed, 16 insertions(+), 8 deletions(-)

diff --git a/apps/web/src/app/[locale]/(auth)/register/components/register-name-form.tsx b/apps/web/src/app/[locale]/(auth)/register/components/register-name-form.tsx
index d6aeeaddb..a834f9a84 100644
--- a/apps/web/src/app/[locale]/(auth)/register/components/register-name-form.tsx
+++ b/apps/web/src/app/[locale]/(auth)/register/components/register-name-form.tsx
@@ -10,7 +10,7 @@ import {
   FormMessage,
 } from "@rallly/ui/form";
 import { Input } from "@rallly/ui/input";
-import { useRouter } from "next/navigation";
+import { useRouter, useSearchParams } from "next/navigation";
 import { useForm } from "react-hook-form";
 import type { z } from "zod";
 
@@ -25,6 +25,7 @@ type RegisterNameFormValues = z.infer<typeof registerNameFormSchema>;
 
 export function RegisterNameForm() {
   const { t } = useTranslation();
+  const searchParams = useSearchParams();
   const form = useForm<RegisterNameFormValues>({
     defaultValues: {
       name: "",
@@ -43,7 +44,10 @@ export function RegisterNameForm() {
 
           if (res.ok) {
             await setToken(res.token);
-            router.push("/register/verify");
+            const redirectTo = searchParams.get("redirectTo");
+            router.push(
+              `/register/verify${redirectTo ? `?redirectTo=${redirectTo}` : ""}`,
+            );
           } else {
             switch (res.reason) {
               case "emailNotAllowed":
diff --git a/apps/web/src/components/poll/poll-header.tsx b/apps/web/src/components/poll/poll-header.tsx
index fc0344c97..32748ec48 100644
--- a/apps/web/src/components/poll/poll-header.tsx
+++ b/apps/web/src/components/poll/poll-header.tsx
@@ -1,13 +1,17 @@
 import { Button } from "@rallly/ui/button";
 import Image from "next/image";
 import Link from "next/link";
+import { usePathname } from "next/navigation";
 
+import { Trans } from "@/components/trans";
 import { UserDropdown } from "@/components/user-dropdown";
 
 import { useUser } from "../user-provider";
 
 export function PollHeader() {
   const { user } = useUser();
+  const pathname = usePathname();
+  const redirectTo = `?redirectTo=${encodeURIComponent(pathname)}`;
   return (
     <div className="sticky top-0 z-40 border-b bg-gray-100/90 p-3 backdrop-blur-md">
       <div className="mx-auto flex max-w-4xl items-center justify-between">
@@ -33,10 +37,14 @@ export function PollHeader() {
           ) : (
             <>
               <Button variant="ghost" asChild>
-                <Link href="/login">Login</Link>
+                <Link href={`/login${redirectTo}`}>
+                  <Trans i18nKey="login" defaults="Login" />
+                </Link>
               </Button>
               <Button variant="primary" asChild>
-                <Link href="/register">Sign Up</Link>
+                <Link href={`/register${redirectTo}`}>
+                  <Trans i18nKey="signUp" defaults="Sign Up" />
+                </Link>
               </Button>
             </>
           )}
diff --git a/apps/web/tests/guest-to-user-migration.spec.ts b/apps/web/tests/guest-to-user-migration.spec.ts
index 139f1fa0d..313fc0b50 100644
--- a/apps/web/tests/guest-to-user-migration.spec.ts
+++ b/apps/web/tests/guest-to-user-migration.spec.ts
@@ -41,10 +41,6 @@ test.describe.serial(() => {
       email: TEST_USER_EMAIL,
     });
 
-    // Step 4: Navigate back to the poll
-    await page.getByRole("main").getByRole("link", { name: "Polls" }).click();
-    await expect(page).toHaveURL(/polls/);
-    await page.click("text=Monthly Meetup");
     await expect(page.getByTestId("poll-title")).toHaveText("Monthly Meetup");
   });