From 09832dc940176518ff08a8805f9dea82ffcafad5 Mon Sep 17 00:00:00 2001
From: EYHN <cneyhn@gmail.com>
Date: Thu, 18 Apr 2024 13:28:32 +0000
Subject: [PATCH] feat(core): add subscribe link (#6610)

---
 .../core/src/components/affine/auth/oauth.tsx |  20 ++-
 .../src/components/affine/auth/sign-in.tsx    |  29 ++++-
 .../core/src/hooks/use-navigate-helper.ts     |  15 ++-
 .../src/modules/cloud/entities/session.ts     |   7 +-
 .../modules/cloud/entities/subscription.ts    |   7 +-
 .../core/src/modules/cloud/services/auth.ts   |  16 +--
 packages/frontend/core/src/pages/invite.tsx   |   9 +-
 packages/frontend/core/src/pages/sign-in.tsx  |  22 +---
 .../frontend/core/src/pages/subscribe.css.ts  |  13 ++
 .../frontend/core/src/pages/subscribe.tsx     | 122 ++++++++++++++++++
 packages/frontend/core/src/router.tsx         |   4 +
 11 files changed, 216 insertions(+), 48 deletions(-)
 create mode 100644 packages/frontend/core/src/pages/subscribe.css.ts
 create mode 100644 packages/frontend/core/src/pages/subscribe.tsx

diff --git a/packages/frontend/core/src/components/affine/auth/oauth.tsx b/packages/frontend/core/src/components/affine/auth/oauth.tsx
index c5c8513021..9ca5f84e77 100644
--- a/packages/frontend/core/src/components/affine/auth/oauth.tsx
+++ b/packages/frontend/core/src/components/affine/auth/oauth.tsx
@@ -25,7 +25,7 @@ const OAuthProviderMap: Record<
   },
 };
 
-export function OAuth() {
+export function OAuth({ redirectUri }: { redirectUri?: string | null }) {
   const serverConfig = useService(ServerConfigService).serverConfig;
   const oauth = useLiveData(serverConfig.features$.map(r => r?.oauth));
   const oauthProviders = useLiveData(
@@ -42,11 +42,21 @@ export function OAuth() {
   }
 
   return oauthProviders?.map(provider => (
-    <OAuthProvider key={provider} provider={provider} />
+    <OAuthProvider
+      key={provider}
+      provider={provider}
+      redirectUri={redirectUri}
+    />
   ));
 }
 
-function OAuthProvider({ provider }: { provider: OAuthProviderType }) {
+function OAuthProvider({
+  provider,
+  redirectUri,
+}: {
+  provider: OAuthProviderType;
+  redirectUri?: string | null;
+}) {
   const { icon } = OAuthProviderMap[provider];
   const authService = useService(AuthService);
   const [isConnecting, setIsConnecting] = useState(false);
@@ -54,7 +64,7 @@ function OAuthProvider({ provider }: { provider: OAuthProviderType }) {
   const onClick = useAsyncCallback(async () => {
     try {
       setIsConnecting(true);
-      await authService.signInOauth(provider);
+      await authService.signInOauth(provider, redirectUri);
     } catch (err) {
       console.error(err);
       notify.error({ message: 'Failed to sign in, please try again.' });
@@ -62,7 +72,7 @@ function OAuthProvider({ provider }: { provider: OAuthProviderType }) {
       setIsConnecting(false);
       mixpanel.track('OAuth', { provider });
     }
-  }, [authService, provider]);
+  }, [authService, provider, redirectUri]);
 
   return (
     <Button
diff --git a/packages/frontend/core/src/components/affine/auth/sign-in.tsx b/packages/frontend/core/src/components/affine/auth/sign-in.tsx
index 926e7bfb9e..6e05efd048 100644
--- a/packages/frontend/core/src/components/affine/auth/sign-in.tsx
+++ b/packages/frontend/core/src/components/affine/auth/sign-in.tsx
@@ -8,6 +8,7 @@ import { ArrowDownBigIcon } from '@blocksuite/icons';
 import { useLiveData, useService } from '@toeverything/infra';
 import type { FC } from 'react';
 import { useCallback, useEffect, useState } from 'react';
+import { useSearchParams } from 'react-router-dom';
 
 import { AuthService } from '../../../modules/cloud';
 import { mixpanel } from '../../../utils';
@@ -29,7 +30,7 @@ export const SignIn: FC<AuthPanelProps> = ({
 }) => {
   const t = useAFFiNEI18N();
   const authService = useService(AuthService);
-
+  const [searchParams] = useSearchParams();
   const [isMutating, setIsMutating] = useState(false);
   const [verifyToken, challenge] = useCaptcha();
 
@@ -74,11 +75,21 @@ export const SignIn: FC<AuthPanelProps> = ({
             mixpanel.track_forms('SignIn', 'Email', {
               email,
             });
-            await authService.sendEmailMagicLink(email, verifyToken, challenge);
+            await authService.sendEmailMagicLink(
+              email,
+              verifyToken,
+              challenge,
+              searchParams.get('redirect_uri')
+            );
             setAuthState('afterSignInSendEmail');
           }
         } else {
-          await authService.sendEmailMagicLink(email, verifyToken, challenge);
+          await authService.sendEmailMagicLink(
+            email,
+            verifyToken,
+            challenge,
+            searchParams.get('redirect_uri')
+          );
           mixpanel.track_forms('SignUp', 'Email', {
             email,
           });
@@ -95,7 +106,15 @@ export const SignIn: FC<AuthPanelProps> = ({
     }
 
     setIsMutating(false);
-  }, [authService, challenge, email, setAuthEmail, setAuthState, verifyToken]);
+  }, [
+    authService,
+    challenge,
+    email,
+    searchParams,
+    setAuthEmail,
+    setAuthState,
+    verifyToken,
+  ]);
 
   return (
     <>
@@ -104,7 +123,7 @@ export const SignIn: FC<AuthPanelProps> = ({
         subTitle={t['com.affine.brand.affineCloud']()}
       />
 
-      <OAuth />
+      <OAuth redirectUri={searchParams.get('redirect_uri')} />
 
       <div className={style.authModalContent}>
         <AuthInput
diff --git a/packages/frontend/core/src/hooks/use-navigate-helper.ts b/packages/frontend/core/src/hooks/use-navigate-helper.ts
index c2b24be04b..5c86cdd4af 100644
--- a/packages/frontend/core/src/hooks/use-navigate-helper.ts
+++ b/packages/frontend/core/src/hooks/use-navigate-helper.ts
@@ -137,13 +137,20 @@ export function useNavigateHelper() {
   );
   const jumpToSignIn = useCallback(
     (
+      redirectUri?: string,
       logic: RouteLogic = RouteLogic.PUSH,
       otherOptions?: Omit<NavigateOptions, 'replace'>
     ) => {
-      return navigate('/signIn', {
-        replace: logic === RouteLogic.REPLACE,
-        ...otherOptions,
-      });
+      return navigate(
+        '/signIn' +
+          (redirectUri
+            ? `?redirect_uri=${encodeURIComponent(redirectUri)}`
+            : ''),
+        {
+          replace: logic === RouteLogic.REPLACE,
+          ...otherOptions,
+        }
+      );
     },
     [navigate]
   );
diff --git a/packages/frontend/core/src/modules/cloud/entities/session.ts b/packages/frontend/core/src/modules/cloud/entities/session.ts
index 83c3112364..c5dd7b664d 100644
--- a/packages/frontend/core/src/modules/cloud/entities/session.ts
+++ b/packages/frontend/core/src/modules/cloud/entities/session.ts
@@ -106,9 +106,12 @@ export class AuthSession extends Entity {
     }
   }
 
-  async waitForRevalidation() {
+  async waitForRevalidation(signal?: AbortSignal) {
     this.revalidate();
-    await this.isRevalidating$.waitFor(isRevalidating => !isRevalidating);
+    await this.isRevalidating$.waitFor(
+      isRevalidating => !isRevalidating,
+      signal
+    );
   }
 
   async removeAvatar() {
diff --git a/packages/frontend/core/src/modules/cloud/entities/subscription.ts b/packages/frontend/core/src/modules/cloud/entities/subscription.ts
index 8b3eb61e08..71393ece22 100644
--- a/packages/frontend/core/src/modules/cloud/entities/subscription.ts
+++ b/packages/frontend/core/src/modules/cloud/entities/subscription.ts
@@ -84,9 +84,12 @@ export class Subscription extends Entity {
     await this.waitForRevalidation();
   }
 
-  async waitForRevalidation() {
+  async waitForRevalidation(signal?: AbortSignal) {
     this.revalidate();
-    await this.isRevalidating$.waitFor(isRevalidating => !isRevalidating);
+    await this.isRevalidating$.waitFor(
+      isRevalidating => !isRevalidating,
+      signal
+    );
   }
 
   revalidate = effect(
diff --git a/packages/frontend/core/src/modules/cloud/services/auth.ts b/packages/frontend/core/src/modules/cloud/services/auth.ts
index 8f955c0a3d..cc7ae434b6 100644
--- a/packages/frontend/core/src/modules/cloud/services/auth.ts
+++ b/packages/frontend/core/src/modules/cloud/services/auth.ts
@@ -76,18 +76,18 @@ export class AuthService extends Service {
   async sendEmailMagicLink(
     email: string,
     verifyToken: string,
-    challenge?: string
+    challenge?: string,
+    redirectUri?: string | null
   ) {
     const searchParams = new URLSearchParams();
     if (challenge) {
       searchParams.set('challenge', challenge);
     }
     searchParams.set('token', verifyToken);
-    const redirectUri = new URL(location.href);
-    if (environment.isDesktop) {
-      redirectUri.pathname = this.buildRedirectUri('/open-app/signin-redirect');
-    }
-    searchParams.set('redirect_uri', redirectUri.toString());
+    const redirect = environment.isDesktop
+      ? this.buildRedirectUri('/open-app/signin-redirect')
+      : redirectUri ?? location.href;
+    searchParams.set('redirect_uri', redirect.toString());
 
     const res = await this.fetchService.fetch(
       '/api/auth/sign-in?' + searchParams.toString(),
@@ -104,7 +104,7 @@ export class AuthService extends Service {
     }
   }
 
-  async signInOauth(provider: OAuthProviderType) {
+  async signInOauth(provider: OAuthProviderType, redirectUri?: string | null) {
     if (environment.isDesktop) {
       await apis?.ui.openExternal(
         `${
@@ -117,7 +117,7 @@ export class AuthService extends Service {
       location.href = `${
         runtimeConfig.serverUrlPrefix
       }/oauth/login?provider=${provider}&redirect_uri=${encodeURIComponent(
-        location.pathname
+        redirectUri ?? location.pathname
       )}`;
     }
 
diff --git a/packages/frontend/core/src/pages/invite.tsx b/packages/frontend/core/src/pages/invite.tsx
index 69ae5bc671..8fff63943d 100644
--- a/packages/frontend/core/src/pages/invite.tsx
+++ b/packages/frontend/core/src/pages/invite.tsx
@@ -79,11 +79,10 @@ export const Component = () => {
     if (loginStatus === 'unauthenticated' && !isRevalidating) {
       // We can not pass function to navigate state, so we need to save it in atom
       setOnceSignedInEvent(openWorkspace);
-      jumpToSignIn(RouteLogic.REPLACE, {
-        state: {
-          callbackURL: `/workspace/${inviteInfo.workspace.id}/all`,
-        },
-      });
+      jumpToSignIn(
+        `/workspace/${inviteInfo.workspace.id}/all`,
+        RouteLogic.REPLACE
+      );
     }
   }, [
     inviteInfo.workspace.id,
diff --git a/packages/frontend/core/src/pages/sign-in.tsx b/packages/frontend/core/src/pages/sign-in.tsx
index 7bbf5e2fa1..28b9e5084f 100644
--- a/packages/frontend/core/src/pages/sign-in.tsx
+++ b/packages/frontend/core/src/pages/sign-in.tsx
@@ -5,25 +5,19 @@ import { useLiveData, useService } from '@toeverything/infra';
 import { useAtom } from 'jotai';
 import { useCallback, useEffect } from 'react';
 // eslint-disable-next-line @typescript-eslint/no-restricted-imports
-import { useLocation, useNavigate, useSearchParams } from 'react-router-dom';
+import { useNavigate, useSearchParams } from 'react-router-dom';
 
 import { authAtom } from '../atoms';
 import type { AuthProps } from '../components/affine/auth';
 import { AuthPanel } from '../components/affine/auth';
 import { RouteLogic, useNavigateHelper } from '../hooks/use-navigate-helper';
 
-interface LocationState {
-  state?: {
-    callbackURL?: string;
-  };
-}
 export const SignIn = () => {
   const [{ state, email = '', emailType = 'changePassword' }, setAuthAtom] =
     useAtom(authAtom);
   const session = useService(AuthService).session;
   const status = useLiveData(session.status$);
   const isRevalidating = useLiveData(session.isRevalidating$);
-  const location = useLocation() as LocationState;
   const navigate = useNavigate();
   const { jumpToIndex } = useNavigateHelper();
   const [searchParams] = useSearchParams();
@@ -31,8 +25,9 @@ export const SignIn = () => {
 
   useEffect(() => {
     if (isLoggedIn) {
-      if (location.state?.callbackURL) {
-        navigate(location.state.callbackURL, {
+      const redirectUri = searchParams.get('redirect_uri');
+      if (redirectUri) {
+        navigate(redirectUri, {
           replace: true,
         });
       } else {
@@ -41,14 +36,7 @@ export const SignIn = () => {
         });
       }
     }
-  }, [
-    jumpToIndex,
-    location.state,
-    navigate,
-    setAuthAtom,
-    isLoggedIn,
-    searchParams,
-  ]);
+  }, [jumpToIndex, navigate, setAuthAtom, isLoggedIn, searchParams]);
 
   const onSetEmailType = useCallback(
     (emailType: AuthProps['emailType']) => {
diff --git a/packages/frontend/core/src/pages/subscribe.css.ts b/packages/frontend/core/src/pages/subscribe.css.ts
new file mode 100644
index 0000000000..bcba5554e8
--- /dev/null
+++ b/packages/frontend/core/src/pages/subscribe.css.ts
@@ -0,0 +1,13 @@
+import { cssVar } from '@toeverything/theme';
+import { style } from '@vanilla-extract/css';
+
+export const container = style({
+  display: 'flex',
+  flexDirection: 'column',
+  alignItems: 'center',
+  justifyContent: 'center',
+  height: '100vh',
+  width: '100%',
+  lineHeight: 4,
+  color: cssVar('--affine-text-secondary-color'),
+});
diff --git a/packages/frontend/core/src/pages/subscribe.tsx b/packages/frontend/core/src/pages/subscribe.tsx
new file mode 100644
index 0000000000..162f0651cc
--- /dev/null
+++ b/packages/frontend/core/src/pages/subscribe.tsx
@@ -0,0 +1,122 @@
+import { Button, Loading } from '@affine/component';
+import { SubscriptionPlan, SubscriptionRecurring } from '@affine/graphql';
+import { effect, fromPromise, useServices } from '@toeverything/infra';
+import { nanoid } from 'nanoid';
+import { useEffect, useMemo, useState } from 'react';
+import { useSearchParams } from 'react-router-dom';
+import { EMPTY, mergeMap, switchMap } from 'rxjs';
+
+import { RouteLogic, useNavigateHelper } from '../hooks/use-navigate-helper';
+import { AuthService, SubscriptionService } from '../modules/cloud';
+import { container } from './subscribe.css';
+
+export const Component = () => {
+  const { authService, subscriptionService } = useServices({
+    AuthService,
+    SubscriptionService,
+  });
+  const [searchParams] = useSearchParams();
+  const [message, setMessage] = useState('');
+  const [error, setError] = useState('');
+  const [retryKey, setRetryKey] = useState(0);
+  const { jumpToSignIn, jumpToIndex } = useNavigateHelper();
+  const idempotencyKey = useMemo(() => nanoid(), []);
+
+  const plan = searchParams.get('plan') as string | null;
+  const recurring = searchParams.get('recurring') as string | null;
+
+  useEffect(() => {
+    const call = effect(
+      switchMap(() => {
+        return fromPromise(async signal => {
+          retryKey;
+          // TODO: i18n
+          setMessage('Checking account status...');
+          setError('');
+          await authService.session.waitForRevalidation(signal);
+          const loggedIn =
+            authService.session.status$.value === 'authenticated';
+          if (!loggedIn) {
+            setMessage('Redirecting to sign in...');
+            jumpToSignIn(
+              location.pathname + location.search,
+              RouteLogic.REPLACE
+            );
+            return;
+          }
+          setMessage('Checking subscription status...');
+          await subscriptionService.subscription.waitForRevalidation(signal);
+          const subscribed = !!subscriptionService.subscription.ai$.value;
+          if (!subscribed) {
+            setMessage('Creating checkout...');
+            try {
+              const checkout = await subscriptionService.createCheckoutSession({
+                idempotencyKey,
+                plan:
+                  plan?.toLowerCase() === 'ai'
+                    ? SubscriptionPlan.AI
+                    : SubscriptionPlan.Pro,
+                coupon: null,
+                recurring:
+                  recurring?.toLowerCase() === 'monthly'
+                    ? SubscriptionRecurring.Monthly
+                    : SubscriptionRecurring.Yearly,
+                successCallbackLink: null,
+              });
+              setMessage('Redirecting...');
+              location.href = checkout;
+            } catch (err) {
+              console.error(err);
+              setError('Something went wrong. Please try again.');
+            }
+          } else {
+            setMessage('Your account is already subscribed. Redirecting...');
+            await new Promise(resolve => {
+              setTimeout(resolve, 5000);
+            });
+            jumpToIndex(RouteLogic.REPLACE);
+          }
+        }).pipe(mergeMap(() => EMPTY));
+      })
+    );
+
+    call();
+
+    return () => {
+      call.unsubscribe();
+    };
+  }, [
+    authService,
+    subscriptionService,
+    jumpToSignIn,
+    idempotencyKey,
+    plan,
+    jumpToIndex,
+    recurring,
+    retryKey,
+  ]);
+
+  useEffect(() => {
+    authService.session.revalidate();
+  }, [authService]);
+
+  return (
+    <div className={container}>
+      {!error ? (
+        <>
+          {message}
+          <br />
+          <Loading size={20} />
+        </>
+      ) : (
+        <>
+          {error}
+          <br />
+          <Button type="primary" onClick={() => setRetryKey(i => i + 1)}>
+            Retry
+          </Button>
+        </>
+      )}
+    </div>
+  );
+};
diff --git a/packages/frontend/core/src/router.tsx b/packages/frontend/core/src/router.tsx
index d362881f33..44179d521c 100644
--- a/packages/frontend/core/src/router.tsx
+++ b/packages/frontend/core/src/router.tsx
@@ -96,6 +96,10 @@ export const topLevelRoutes = [
         path: '/redirect-proxy',
         lazy: () => import('./pages/redirect'),
       },
+      {
+        path: '/subscribe',
+        lazy: () => import('./pages/subscribe'),
+      },
       {
         path: '*',
         lazy: () => import('./pages/404'),