feat(core): support subscribe plan after login (#4788)

packages/frontend/core/src/components/affine/auth/after-sign-in-send-email.tsx

@@ -7,13 +7,14 @@ import {
 import { Trans } from '@affine/i18n';
 import { useAFFiNEI18N } from '@affine/i18n/hooks';
 import { Button } from '@toeverything/components/button';
-import { useCallback } from 'react';
+import React, { useCallback } from 'react';
 
 import { useCurrentLoginStatus } from '../../../hooks/affine/use-current-login-status';
 import type { AuthPanelProps } from './index';
 import * as style from './style.css';
 import { useAuth } from './use-auth';
 import { Captcha, useCaptcha } from './use-captcha';
+import { useSubscriptionSearch } from './use-subscription';
 
 export const AfterSignInSendEmail = ({
   setAuthState,
@@ -23,6 +24,7 @@ export const AfterSignInSendEmail = ({
   const t = useAFFiNEI18N();
   const loginStatus = useCurrentLoginStatus();
   const [verifyToken, challenge] = useCaptcha();
+  const subscriptionData = useSubscriptionSearch();
 
   const { resendCountDown, allowSendEmail, signIn } = useAuth();
   if (loginStatus === 'authenticated') {
@@ -35,6 +37,14 @@ export const AfterSignInSendEmail = ({
     }
   }, [challenge, email, signIn, verifyToken]);
 
+  const onSignInWithPasswordClick = useCallback(() => {
+    setAuthState('signInWithPassword');
+  }, [setAuthState]);
+
+  const onBackBottomClick = useCallback(() => {
+    setAuthState('signIn');
+  }, [setAuthState]);
+
   return (
     <>
       <ModalHeader
@@ -75,26 +85,25 @@ export const AfterSignInSendEmail = ({
       </div>
 
       <div className={style.authMessage} style={{ marginTop: 20 }}>
-        {/*prettier-ignore*/}
-        <Trans i18nKey="com.affine.auth.sign.auth.code.message.password">
-          If you haven&apos;t received the email, please check your spam folder.
-          Or <span
-            className="link"
-            data-testid='sign-in-with-password'
-            onClick={useCallback(() => {
-              setAuthState('signInWithPassword');
-            }, [setAuthState])}
-          >
-            sign in with password
-          </span> instead.
-        </Trans>
+        {t['com.affine.auth.sign.auth.code.message']()}
+        {subscriptionData ? null : ( // If with payment, just support email sign in to avoid duplicate redirect to the same stripe url.
+          <React.Fragment>
+            &nbsp;
+            {/*prettier-ignore*/}
+            <Trans i18nKey="com.affine.auth.sign.auth.code.message.password">
+              Or <span
+                className="link"
+                data-testid='sign-in-with-password'
+                onClick={onSignInWithPasswordClick}
+              >
+                sign in with password
+              </span> instead.
+            </Trans>
+          </React.Fragment>
+        )}
       </div>
 
-      <BackButton
-        onClick={useCallback(() => {
-          setAuthState('signIn');
-        }, [setAuthState])}
-      />
+      <BackButton onClick={onBackBottomClick} />
     </>
   );
 };

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

@@ -0,0 +1,9 @@
+import { style } from '@vanilla-extract/css';
+
+export const loadingContainer = style({
+  display: 'flex',
+  width: '100vw',
+  height: '60vh',
+  justifyContent: 'center',
+  alignItems: 'center',
+});

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

@@ -0,0 +1,44 @@
+import type { SubscriptionRecurring } from '@affine/graphql';
+import { checkoutMutation } from '@affine/graphql';
+import { useMutation } from '@affine/workspace/affine/gql';
+import { Loading } from '@toeverything/components/loading';
+import { nanoid } from 'nanoid';
+import { type FC, useEffect, useMemo } from 'react';
+
+import * as styles from './subscription-redirect.css';
+import { useSubscriptionSearch } from './use-subscription';
+
+export const SubscriptionRedirect: FC = () => {
+  const subscriptionData = useSubscriptionSearch();
+  const idempotencyKey = useMemo(() => nanoid(), []);
+  const { trigger: checkoutSubscription } = useMutation({
+    mutation: checkoutMutation,
+  });
+
+  useEffect(() => {
+    if (!subscriptionData) {
+      throw new Error('No subscription data found');
+    }
+
+    // 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(data => {
+        window.open(data.checkout, '_self', 'norefferer');
+      });
+    }, 100);
+
+    return () => {
+      clearTimeout(timeoutId);
+    };
+
+    // Just run this once, do not react to changes
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, []);
+
+  return (
+    <div className={styles.loadingContainer}>
+      <Loading size={40} />
+    </div>
+  );
+};

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

@@ -5,6 +5,7 @@ import { type SignInResponse } from 'next-auth/react';
 import { useCallback } from 'react';
 
 import { signInCloud } from '../../../utils/cloud-utils';
+import { useSubscriptionSearch } from './use-subscription';
 
 const COUNT_DOWN_TIME = 60;
 export const INTERNAL_BETA_URL = `https://community.affine.pro/c/insider-general/`;
@@ -58,6 +59,7 @@ const countDownAtom = atom(
 );
 
 export const useAuth = () => {
+  const subscriptionData = useSubscriptionSearch();
   const pushNotification = useSetAtom(pushNotificationAtom);
   const [authStore, setAuthStore] = useAtom(authStoreAtom);
   const startResendCountDown = useSetAtom(countDownAtom);
@@ -75,7 +77,9 @@ export const useAuth = () => {
         'email',
         {
           email: email,
-          callbackUrl: '/auth/signIn',
+          callbackUrl: subscriptionData
+            ? subscriptionData.redirectUrl
+            : '/auth/signIn',
           redirect: false,
         },
         challenge
@@ -98,7 +102,7 @@ export const useAuth = () => {
 
       return res;
     },
-    [pushNotification, setAuthStore, startResendCountDown]
+    [pushNotification, setAuthStore, startResendCountDown, subscriptionData]
   );
 
   const signUp = useCallback(
@@ -114,7 +118,9 @@ export const useAuth = () => {
         'email',
         {
           email: email,
-          callbackUrl: '/auth/signUp',
+          callbackUrl: subscriptionData
+            ? subscriptionData.redirectUrl
+            : '/auth/signUp',
           redirect: false,
         },
         challenge
@@ -137,7 +143,7 @@ export const useAuth = () => {
 
       return res;
     },
-    [pushNotification, setAuthStore, startResendCountDown]
+    [pushNotification, setAuthStore, startResendCountDown, subscriptionData]
   );
 
   const signInWithGoogle = useCallback(() => {

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

@@ -0,0 +1,35 @@
+import { useMemo } from 'react';
+import { useSearchParams } from 'react-router-dom';
+
+enum SubscriptionKey {
+  Recurring = 'subscription_recurring',
+  Plan = 'subscription_plan',
+}
+
+export function useSubscriptionSearch() {
+  const [searchParams] = useSearchParams();
+
+  return useMemo(() => {
+    const withPayment =
+      searchParams.has(SubscriptionKey.Recurring) &&
+      searchParams.has(SubscriptionKey.Plan);
+
+    if (!withPayment) {
+      return null;
+    }
+
+    const recurring = searchParams.get(SubscriptionKey.Recurring);
+    const plan = searchParams.get(SubscriptionKey.Plan);
+    return {
+      recurring,
+      plan,
+      get redirectUrl() {
+        const paymentParams = new URLSearchParams([
+          [SubscriptionKey.Recurring, recurring ?? ''],
+          [SubscriptionKey.Plan, plan ?? ''],
+        ]);
+        return `/auth/subscription-redirect?${paymentParams.toString()}`;
+      },
+    };
+  }, [searchParams]);
+}

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

@@ -25,6 +25,7 @@ import {
 } from 'react-router-dom';
 import { z } from 'zod';
 
+import { SubscriptionRedirect } from '../components/affine/auth/subscription-redirect';
 import { useCurrentLoginStatus } from '../hooks/affine/use-current-login-status';
 import { useCurrentUser } from '../hooks/affine/use-current-user';
 import { RouteLogic, useNavigateHelper } from '../hooks/use-navigate-helper';
@@ -36,6 +37,7 @@ const authTypeSchema = z.enum([
   'signUp',
   'changeEmail',
   'confirm-change-email',
+  'subscription-redirect',
 ]);
 
 export const AuthPage = (): ReactElement | null => {
@@ -132,6 +134,9 @@ export const AuthPage = (): ReactElement | null => {
     case 'confirm-change-email': {
       return <ConfirmChangeEmail onOpenAffine={onOpenAffine} />;
     }
+    case 'subscription-redirect': {
+      return <SubscriptionRedirect />;
+    }
   }
   return null;
 };