feat(core): support signup set password before goto stripe payment url (#4892)

packages/frontend/core/src/components/affine/auth/subscription-redirect.tsx

@@ -1,13 +1,19 @@
+import { SignUpPage } from '@affine/component/auth-components';
 import { AffineShapeIcon } from '@affine/component/page-list';
 import type { SubscriptionRecurring } from '@affine/graphql';
-import { checkoutMutation, subscriptionQuery } from '@affine/graphql';
+import {
+  changePasswordMutation,
+  checkoutMutation,
+  subscriptionQuery,
+} from '@affine/graphql';
 import { useAFFiNEI18N } from '@affine/i18n/hooks';
 import { useMutation, useQuery } from '@affine/workspace/affine/gql';
 import { Button } from '@toeverything/components/button';
 import { Loading } from '@toeverything/components/loading';
 import { nanoid } from 'nanoid';
-import { type FC, Suspense, useCallback, useEffect, useMemo } from 'react';
+import { Suspense, useCallback, useEffect, useMemo } from 'react';
 
+import { useCurrentUser } from '../../../hooks/affine/use-current-user';
 import {
   RouteLogic,
   useNavigateHelper,
@@ -15,6 +21,27 @@ import {
 import * as styles from './subscription-redirect.css';
 import { useSubscriptionSearch } from './use-subscription';
 
+const usePaymentRedirect = () => {
+  const searchData = useSubscriptionSearch();
+  if (!searchData?.recurring) {
+    throw new Error('Invalid recurring data.');
+  }
+
+  const recurring = searchData.recurring as SubscriptionRecurring;
+  const idempotencyKey = useMemo(() => nanoid(), []);
+  const { trigger: checkoutSubscription } = useMutation({
+    mutation: checkoutMutation,
+  });
+
+  return useCallback(() => {
+    checkoutSubscription({ recurring, idempotencyKey })
+      .then(({ checkout }) => {
+        window.open(checkout, '_self', 'norefferer');
+      })
+      .catch(e => console.error(e));
+  }, [recurring, idempotencyKey, checkoutSubscription]);
+};
+
 const CenterLoading = () => {
   return (
     <div className={styles.loadingContainer}>
@@ -51,54 +78,65 @@ const SubscriptionExisting = () => {
   );
 };
 
-const SubscriptionRedirectInner: FC = () => {
-  const subscriptionData = useSubscriptionSearch();
-  const idempotencyKey = useMemo(() => nanoid(), []);
-  const { data } = useQuery({
-    query: subscriptionQuery,
-  });
-  const { trigger: checkoutSubscription } = useMutation({
-    mutation: checkoutMutation,
-  });
-
+const SubscriptionRedirection = ({ redirect }: { redirect: () => void }) => {
   useEffect(() => {
-    if (!subscriptionData) {
-      throw new Error('No subscription data found');
-    }
-
-    if (data.currentUser?.subscription) {
-      return;
-    }
-
-    // This component will be render multiple times, use timeout to avoid multiple effect.
     const timeoutId = setTimeout(() => {
-      const recurring = subscriptionData.recurring as SubscriptionRecurring;
-      checkoutSubscription({ recurring, idempotencyKey }).then(
-        ({ checkout }) => {
-          window.open(checkout, '_self', 'norefferer');
-        }
-      );
+      redirect();
     }, 100);
 
     return () => {
       clearTimeout(timeoutId);
     };
+  }, [redirect]);
 
-    // Just run this once, do not react to changes
-    // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, []);
+  return <CenterLoading />;
+};
 
-  if (data.currentUser?.subscription) {
+const SubscriptionRedirectWithData = () => {
+  const t = useAFFiNEI18N();
+  const user = useCurrentUser();
+  const searchData = useSubscriptionSearch();
+  const openPaymentUrl = usePaymentRedirect();
+
+  const { trigger: changePassword } = useMutation({
+    mutation: changePasswordMutation,
+  });
+  const { data: subscriptionData } = useQuery({
+    query: subscriptionQuery,
+  });
+
+  const onSetPassword = useCallback(
+    async (password: string) => {
+      await changePassword({
+        token: searchData?.passwordToken ?? '',
+        newPassword: password,
+      });
+    },
+    [changePassword, searchData]
+  );
+
+  if (searchData?.withSignUp) {
+    return (
+      <SignUpPage
+        user={user}
+        onSetPassword={onSetPassword}
+        onOpenAffine={openPaymentUrl}
+        openButtonText={t['com.affine.payment.subscription.go-to-subscribe']()}
+      />
+    );
+  }
+
+  if (subscriptionData.currentUser?.subscription) {
     return <SubscriptionExisting />;
   }
 
-  return <CenterLoading />;
+  return <SubscriptionRedirection redirect={openPaymentUrl} />;
 };
 
 export const SubscriptionRedirect = () => {
   return (
     <Suspense fallback={<CenterLoading />}>
-      <SubscriptionRedirectInner />
+      <SubscriptionRedirectWithData />
     </Suspense>
   );
 };

packages/frontend/core/src/components/affine/auth/use-auth.ts

@@ -78,7 +78,7 @@ export const useAuth = () => {
         {
           email: email,
           callbackUrl: subscriptionData
-            ? subscriptionData.redirectUrl
+            ? subscriptionData.getRedirectUrl(false)
             : '/auth/signIn',
           redirect: false,
         },
@@ -119,7 +119,7 @@ export const useAuth = () => {
         {
           email: email,
           callbackUrl: subscriptionData
-            ? subscriptionData.redirectUrl
+            ? subscriptionData.getRedirectUrl(true)
             : '/auth/signUp',
           redirect: false,
         },

packages/frontend/core/src/components/affine/auth/use-subscription.ts

@@ -4,6 +4,8 @@ import { useSearchParams } from 'react-router-dom';
 enum SubscriptionKey {
   Recurring = 'subscription_recurring',
   Plan = 'subscription_plan',
+  SignUp = 'sign_up', // A new user with subscription journey: signup > set password > pay in stripe > go to app
+  Token = 'token', // When signup, there should have a token to set password
 }
 
 export function useSubscriptionSearch() {
@@ -20,14 +22,23 @@ export function useSubscriptionSearch() {
 
     const recurring = searchParams.get(SubscriptionKey.Recurring);
     const plan = searchParams.get(SubscriptionKey.Plan);
+    const withSignUp = searchParams.get(SubscriptionKey.SignUp) === '1';
+    const passwordToken = searchParams.get(SubscriptionKey.Token);
     return {
       recurring,
       plan,
-      get redirectUrl() {
+      withSignUp,
+      passwordToken,
+      getRedirectUrl(signUp?: boolean) {
         const paymentParams = new URLSearchParams([
           [SubscriptionKey.Recurring, recurring ?? ''],
           [SubscriptionKey.Plan, plan ?? ''],
         ]);
+
+        if (signUp) {
+          paymentParams.set(SubscriptionKey.SignUp, '1');
+        }
+
         return `/auth/subscription-redirect?${paymentParams.toString()}`;
       },
     };

packages/frontend/core/src/pages/auth.tsx

@@ -80,11 +80,11 @@ export const AuthPage = (): ReactElement | null => {
   );
 
   const onSetPassword = useCallback(
-    (password: string) => {
-      changePassword({
+    async (password: string) => {
+      await changePassword({
         token: searchParams.get('token') || '',
         newPassword: password,
-      }).catch(console.error);
+      });
     },
     [changePassword, searchParams]
   );