feat(core): add workspace billing (#9043)

packages/frontend/core/src/desktop/dialogs/setting/general-setting/plans/actions.tsx

@@ -8,7 +8,11 @@ import { nanoid } from 'nanoid';
 import type { PropsWithChildren } from 'react';
 import { useEffect, useState } from 'react';
 
-import { AuthService, SubscriptionService } from '../../../../../modules/cloud';
+import {
+  AuthService,
+  SubscriptionService,
+  WorkspaceSubscriptionService,
+} from '../../../../../modules/cloud';
 import { ConfirmLoadingModal, DowngradeModal } from './modals';
 
 /**
@@ -101,15 +105,15 @@ export const CancelTeamAction = ({
 } & PropsWithChildren) => {
   const [idempotencyKey, setIdempotencyKey] = useState(nanoid());
   const [isMutating, setIsMutating] = useState(false);
-  const subscription = useService(SubscriptionService).subscription;
-  const teamSubscription = useLiveData(subscription.team$);
+  const subscription = useService(WorkspaceSubscriptionService).subscription;
+  const workspaceSubscription = useLiveData(subscription.subscription$);
   const authService = useService(AuthService);
   const downgradeNotify = useDowngradeNotify();
 
   const downgrade = useAsyncCallback(async () => {
     try {
       const account = authService.session.account$.value;
-      const prevRecurring = teamSubscription?.recurring;
+      const prevRecurring = workspaceSubscription?.recurring;
       setIsMutating(true);
       await subscription.cancelSubscription(idempotencyKey);
       await subscription.waitForRevalidation();
@@ -133,7 +137,7 @@ export const CancelTeamAction = ({
     }
   }, [
     authService.session.account$.value,
-    teamSubscription,
+    workspaceSubscription,
     subscription,
     idempotencyKey,
     onOpenChange,

packages/frontend/core/src/desktop/dialogs/setting/workspace-setting/billing/index.tsx

@@ -10,8 +10,8 @@ import { useAsyncCallback } from '@affine/core/components/hooks/affine-async-hoo
 import { useMutation } from '@affine/core/components/hooks/use-mutation';
 import {
   AuthService,
-  InvoicesService,
-  SubscriptionService,
+  WorkspaceInvoicesService,
+  WorkspaceSubscriptionService,
 } from '@affine/core/modules/cloud';
 import { UrlService } from '@affine/core/modules/url';
 import {
@@ -41,8 +41,10 @@ import * as styles from './styles.css';
 export const WorkspaceSettingBilling = () => {
   const t = useI18n();
   const workspace = useService(WorkspaceService).workspace;
-  const subscriptionService = useService(SubscriptionService);
-  const team = useLiveData(subscriptionService.subscription.team$);
+  const subscriptionService = useService(WorkspaceSubscriptionService);
+  const subscription = useLiveData(
+    subscriptionService.subscription.subscription$
+  );
   const title = useLiveData(workspace.name$) || 'untitled';
 
   if (workspace === null) {
@@ -51,7 +53,7 @@ export const WorkspaceSettingBilling = () => {
     return null;
   }
 
-  if (!team) {
+  if (!subscription) {
     return <Loading />;
   }
 
@@ -67,8 +69,8 @@ export const WorkspaceSettingBilling = () => {
         <TeamCard />
         <TypeFormLink />
         <PaymentMethodUpdater />
-        {team.end && team.canceledAt ? (
-          <ResumeSubscription expirationDate={team.end} />
+        {subscription?.end && subscription.canceledAt ? (
+          <ResumeSubscription expirationDate={subscription.end} />
         ) : null}
       </SettingWrapper>
 
@@ -81,8 +83,10 @@ export const WorkspaceSettingBilling = () => {
 
 const TeamCard = () => {
   const t = useI18n();
-  const subscriptionService = useService(SubscriptionService);
-  const teamSubscription = useLiveData(subscriptionService.subscription.team$);
+  const subscriptionService = useService(WorkspaceSubscriptionService);
+  const teamSubscription = useLiveData(
+    subscriptionService.subscription.subscription$
+  );
   const teamPrices = useLiveData(subscriptionService.prices.teamPrice$);
 
   const [openCancelModal, setOpenCancelModal] = useState(false);
@@ -198,23 +202,22 @@ const ResumeSubscription = ({ expirationDate }: { expirationDate: string }) => {
 
 const TypeFormLink = () => {
   const t = useI18n();
-  const subscriptionService = useService(SubscriptionService);
+  const workspaceSubscriptionService = useService(WorkspaceSubscriptionService);
   const authService = useService(AuthService);
 
-  const team = useLiveData(subscriptionService.subscription.team$);
+  const workspaceSubscription = useLiveData(
+    workspaceSubscriptionService.subscription.subscription$
+  );
   const account = useLiveData(authService.session.account$);
 
   if (!account) return null;
 
-  const plan = [];
-  if (team) plan.push(SubscriptionPlan.Team);
-
   const link = getUpgradeQuestionnaireLink({
     name: account.info?.name,
     id: account.id,
     email: account.email,
-    recurring: team?.recurring ?? SubscriptionRecurring.Yearly,
-    plan,
+    recurring: workspaceSubscription?.recurring ?? SubscriptionRecurring.Yearly,
+    plan: SubscriptionPlan.Team,
   });
 
   return (
@@ -263,7 +266,7 @@ const PaymentMethodUpdater = () => {
 const BillingHistory = () => {
   const t = useI18n();
 
-  const invoicesService = useService(InvoicesService);
+  const invoicesService = useService(WorkspaceInvoicesService);
   const pageInvoices = useLiveData(invoicesService.invoices.pageInvoices$);
   const invoiceCount = useLiveData(invoicesService.invoices.invoiceCount$);
   const isLoading = useLiveData(invoicesService.invoices.isLoading$);

packages/frontend/core/src/modules/cloud/entities/workspace-invoices.ts

@@ -0,0 +1,83 @@
+import type { InvoicesQuery } from '@affine/graphql';
+import type { WorkspaceService } from '@toeverything/infra';
+import {
+  backoffRetry,
+  catchErrorInto,
+  effect,
+  Entity,
+  exhaustMapSwitchUntilChanged,
+  fromPromise,
+  LiveData,
+  onComplete,
+  onStart,
+} from '@toeverything/infra';
+import { EMPTY, map, mergeMap } from 'rxjs';
+
+import { isBackendError, isNetworkError } from '../error';
+import type { InvoicesStore } from '../stores/invoices';
+
+export type Invoice = NonNullable<
+  InvoicesQuery['currentUser']
+>['invoices'][number];
+
+export class WorkspaceInvoices extends Entity {
+  constructor(
+    private readonly store: InvoicesStore,
+    private readonly workspaceService: WorkspaceService
+  ) {
+    super();
+  }
+
+  pageNum$ = new LiveData(0);
+  invoiceCount$ = new LiveData<number | undefined>(undefined);
+  pageInvoices$ = new LiveData<Invoice[] | undefined>(undefined);
+
+  isLoading$ = new LiveData(false);
+  error$ = new LiveData<any>(null);
+
+  readonly PAGE_SIZE = 8;
+
+  readonly revalidate = effect(
+    map(() => this.pageNum$.value),
+    exhaustMapSwitchUntilChanged(
+      (a, b) => a === b,
+      pageNum => {
+        return fromPromise(async signal => {
+          return this.store.fetchWorkspaceInvoices(
+            pageNum * this.PAGE_SIZE,
+            this.PAGE_SIZE,
+            this.workspaceService.workspace.id,
+            signal
+          );
+        }).pipe(
+          mergeMap(data => {
+            this.invoiceCount$.setValue(data.invoiceCount);
+            this.pageInvoices$.setValue(data.invoices);
+            return EMPTY;
+          }),
+          backoffRetry({
+            when: isNetworkError,
+            count: Infinity,
+          }),
+          backoffRetry({
+            when: isBackendError,
+          }),
+          catchErrorInto(this.error$),
+          onStart(() => {
+            this.pageInvoices$.setValue(undefined);
+            this.isLoading$.setValue(true);
+          }),
+          onComplete(() => this.isLoading$.setValue(false))
+        );
+      }
+    )
+  );
+
+  setPageNum(pageNum: number) {
+    this.pageNum$.setValue(pageNum);
+  }
+
+  override dispose(): void {
+    this.revalidate.unsubscribe();
+  }
+}

packages/frontend/core/src/modules/cloud/entities/workspace-subscription.ts

@@ -0,0 +1,132 @@
+import type { SubscriptionQuery, SubscriptionRecurring } from '@affine/graphql';
+import { SubscriptionPlan } from '@affine/graphql';
+import type { WorkspaceService } from '@toeverything/infra';
+import {
+  backoffRetry,
+  catchErrorInto,
+  effect,
+  Entity,
+  exhaustMapWithTrailing,
+  fromPromise,
+  LiveData,
+  onComplete,
+  onStart,
+} from '@toeverything/infra';
+import { EMPTY, mergeMap } from 'rxjs';
+
+import { isBackendError, isNetworkError } from '../error';
+import type { ServerService } from '../services/server';
+import type { SubscriptionStore } from '../stores/subscription';
+
+export type SubscriptionType = NonNullable<
+  SubscriptionQuery['currentUser']
+>['subscriptions'][number];
+
+export class WorkspaceSubscription extends Entity {
+  subscription$ = new LiveData<SubscriptionType | null | undefined>(null);
+  isRevalidating$ = new LiveData(false);
+  error$ = new LiveData<any | null>(null);
+
+  team$ = this.subscription$.map(
+    subscription => subscription?.plan === SubscriptionPlan.Team
+  );
+
+  constructor(
+    private readonly workspaceService: WorkspaceService,
+    private readonly serverService: ServerService,
+    private readonly store: SubscriptionStore
+  ) {
+    super();
+  }
+
+  async resumeSubscription(idempotencyKey: string, plan?: SubscriptionPlan) {
+    await this.store.mutateResumeSubscription(idempotencyKey, plan);
+    await this.waitForRevalidation();
+  }
+
+  async cancelSubscription(idempotencyKey: string, plan?: SubscriptionPlan) {
+    await this.store.mutateCancelSubscription(idempotencyKey, plan);
+    await this.waitForRevalidation();
+  }
+
+  async setSubscriptionRecurring(
+    idempotencyKey: string,
+    recurring: SubscriptionRecurring,
+    plan?: SubscriptionPlan
+  ) {
+    await this.store.setSubscriptionRecurring(idempotencyKey, recurring, plan);
+    await this.waitForRevalidation();
+  }
+
+  async waitForRevalidation(signal?: AbortSignal) {
+    this.revalidate();
+    await this.isRevalidating$.waitFor(
+      isRevalidating => !isRevalidating,
+      signal
+    );
+  }
+
+  revalidate = effect(
+    exhaustMapWithTrailing(() => {
+      return fromPromise(async signal => {
+        const currentWorkspaceId = this.workspaceService.workspace.id;
+        if (!currentWorkspaceId) {
+          return undefined; // no subscription if no user
+        }
+        const serverConfig =
+          await this.serverService.server.features$.waitForNonNull(signal);
+
+        if (!serverConfig.payment) {
+          // No payment feature, no subscription
+          return {
+            workspaceId: currentWorkspaceId,
+            subscription: null,
+          };
+        }
+        const { workspaceId, subscription } =
+          await this.store.fetchWorkspaceSubscriptions(
+            currentWorkspaceId,
+            signal
+          );
+        return {
+          workspaceId: workspaceId,
+          subscription: subscription,
+        };
+      }).pipe(
+        backoffRetry({
+          when: isNetworkError,
+          count: Infinity,
+        }),
+        backoffRetry({
+          when: isBackendError,
+        }),
+        mergeMap(data => {
+          if (data && data.subscription && data.workspaceId) {
+            this.store.setCachedWorkspaceSubscription(
+              data.workspaceId,
+              data.subscription
+            );
+            this.subscription$.next(data.subscription);
+          } else {
+            this.subscription$.next(undefined);
+          }
+          return EMPTY;
+        }),
+        catchErrorInto(this.error$),
+        onStart(() => this.isRevalidating$.next(true)),
+        onComplete(() => this.isRevalidating$.next(false))
+      );
+    })
+  );
+
+  reset() {
+    this.subscription$.next(null);
+    this.team$.next(false);
+    this.isRevalidating$.next(false);
+    this.error$.next(null);
+  }
+
+  override dispose(): void {
+    this.revalidate.unsubscribe();
+  }
+}

packages/frontend/core/src/modules/cloud/index.ts

@@ -28,7 +28,9 @@ export { UserCopilotQuotaService } from './services/user-copilot-quota';
 export { UserFeatureService } from './services/user-feature';
 export { UserQuotaService } from './services/user-quota';
 export { WebSocketService } from './services/websocket';
+export { WorkspaceInvoicesService } from './services/workspace-invoices';
 export { WorkspaceServerService } from './services/workspace-server';
+export { WorkspaceSubscriptionService } from './services/workspace-subscription';
 export type { ServerConfig } from './types';
 
 import {
@@ -39,6 +41,7 @@ import {
   GlobalState,
   GlobalStateService,
   WorkspaceScope,
+  WorkspaceService,
 } from '@toeverything/infra';
 
 import { UrlService } from '../url';
@@ -51,6 +54,8 @@ import { SubscriptionPrices } from './entities/subscription-prices';
 import { UserCopilotQuota } from './entities/user-copilot-quota';
 import { UserFeature } from './entities/user-feature';
 import { UserQuota } from './entities/user-quota';
+import { WorkspaceInvoices } from './entities/workspace-invoices';
+import { WorkspaceSubscription } from './entities/workspace-subscription';
 import { DefaultRawFetchProvider, RawFetchProvider } from './provider/fetch';
 import { ValidatorProvider } from './provider/validator';
 import { WebSocketAuthProvider } from './provider/websocket-auth';
@@ -70,7 +75,9 @@ import { UserCopilotQuotaService } from './services/user-copilot-quota';
 import { UserFeatureService } from './services/user-feature';
 import { UserQuotaService } from './services/user-quota';
 import { WebSocketService } from './services/websocket';
+import { WorkspaceInvoicesService } from './services/workspace-invoices';
 import { WorkspaceServerService } from './services/workspace-server';
+import { WorkspaceSubscriptionService } from './services/workspace-subscription';
 import { AuthStore } from './stores/auth';
 import { CloudDocMetaStore } from './stores/cloud-doc-meta';
 import { InvoicesStore } from './stores/invoices';
@@ -142,7 +149,16 @@ export function configureCloudModule(framework: Framework) {
     .store(UserFeatureStore, [GraphQLService])
     .service(InvoicesService)
     .store(InvoicesStore, [GraphQLService])
-    .entity(Invoices, [InvoicesStore]);
+    .entity(Invoices, [InvoicesStore])
+    .scope(WorkspaceScope)
+    .service(WorkspaceSubscriptionService, [SubscriptionStore])
+    .entity(WorkspaceSubscription, [
+      WorkspaceService,
+      ServerService,
+      SubscriptionStore,
+    ])
+    .service(WorkspaceInvoicesService)
+    .entity(WorkspaceInvoices, [InvoicesStore, WorkspaceService]);
 
   framework
     .scope(WorkspaceScope)

packages/frontend/core/src/modules/cloud/services/workspace-invoices.ts

@@ -0,0 +1,7 @@
+import { Service } from '@toeverything/infra';
+
+import { WorkspaceInvoices } from '../entities/workspace-invoices';
+
+export class WorkspaceInvoicesService extends Service {
+  invoices = this.framework.createEntity(WorkspaceInvoices);
+}

packages/frontend/core/src/modules/cloud/services/workspace-subscription.ts

@@ -0,0 +1,19 @@
+import { type CreateCheckoutSessionInput } from '@affine/graphql';
+import { Service } from '@toeverything/infra';
+
+import { SubscriptionPrices } from '../entities/subscription-prices';
+import { WorkspaceSubscription } from '../entities/workspace-subscription';
+import type { SubscriptionStore } from '../stores/subscription';
+
+export class WorkspaceSubscriptionService extends Service {
+  subscription = this.framework.createEntity(WorkspaceSubscription);
+  prices = this.framework.createEntity(SubscriptionPrices);
+
+  constructor(private readonly store: SubscriptionStore) {
+    super();
+  }
+
+  async createCheckoutSession(input: CreateCheckoutSessionInput) {
+    return await this.store.createCheckoutSession(input);
+  }
+}

packages/frontend/core/src/modules/cloud/stores/invoices.ts

@@ -1,4 +1,4 @@
-import { invoicesQuery } from '@affine/graphql';
+import { invoicesQuery, workspaceInvoicesQuery } from '@affine/graphql';
 import { Store } from '@toeverything/infra';
 
 import type { GraphQLService } from '../services/graphql';
@@ -21,4 +21,23 @@ export class InvoicesStore extends Store {
 
     return data.currentUser;
   }
+
+  async fetchWorkspaceInvoices(
+    skip: number,
+    take: number,
+    workspaceId: string,
+    signal?: AbortSignal
+  ) {
+    const data = await this.graphqlService.gql({
+      query: workspaceInvoicesQuery,
+      variables: { skip, take, workspaceId },
+      context: { signal },
+    });
+
+    if (!data.workspace) {
+      throw new Error('No workspace');
+    }
+
+    return data.workspace;
+  }
 }

packages/frontend/core/src/modules/cloud/stores/subscription.ts

@@ -5,6 +5,7 @@ import type {
 import {
   cancelSubscriptionMutation,
   createCheckoutSessionMutation,
+  getWorkspaceSubscriptionQuery,
   pricesQuery,
   resumeSubscriptionMutation,
   SubscriptionPlan,
@@ -27,7 +28,11 @@ const getDefaultSubscriptionSuccessCallbackLink = (
   scheme?: string
 ) => {
   const path =
-    plan === SubscriptionPlan.AI ? '/ai-upgrade-success' : '/upgrade-success';
+    plan === SubscriptionPlan.Team
+      ? '/upgrade-success/team'
+      : plan === SubscriptionPlan.AI
+        ? '/ai-upgrade-success'
+        : '/upgrade-success';
   const urlString = baseUrl + path;
   const url = new URL(urlString);
   if (scheme) {
@@ -64,6 +69,30 @@ export class SubscriptionStore extends Store {
     };
   }
 
+  async fetchWorkspaceSubscriptions(
+    workspaceId: string,
+    abortSignal?: AbortSignal
+  ) {
+    const data = await this.gqlService.gql({
+      query: getWorkspaceSubscriptionQuery,
+      variables: {
+        workspaceId,
+      },
+      context: {
+        signal: abortSignal,
+      },
+    });
+
+    if (!data.workspace) {
+      throw new Error('No workspace');
+    }
+
+    return {
+      workspaceId: data.workspace.subscription?.id,
+      subscription: data.workspace.subscription,
+    };
+  }
+
   async mutateResumeSubscription(
     idempotencyKey: string,
     plan?: SubscriptionPlan,
@@ -114,6 +143,22 @@ export class SubscriptionStore extends Store {
     return this.globalCache.set(SUBSCRIPTION_CACHE_KEY + userId, subscriptions);
   }
 
+  getCachedWorkspaceSubscription(workspaceId: string) {
+    return this.globalCache.get<SubscriptionType>(
+      SUBSCRIPTION_CACHE_KEY + workspaceId
+    );
+  }
+
+  setCachedWorkspaceSubscription(
+    workspaceId: string,
+    subscription: SubscriptionType
+  ) {
+    return this.globalCache.set(
+      SUBSCRIPTION_CACHE_KEY + workspaceId,
+      subscription
+    );
+  }
+
   setSubscriptionRecurring(
     idempotencyKey: string,
     recurring: SubscriptionRecurring,

packages/frontend/graphql/src/graphql/get-workspace-subscription.gql

@@ -0,0 +1,15 @@
+query getWorkspaceSubscription($workspaceId: String!) {
+  workspace(id: $workspaceId) {
+    subscription {
+      id
+      status
+      plan
+      recurring
+      start
+      end
+      nextBillAt
+      canceledAt
+      variant
+    }
+  }
+}

packages/frontend/graphql/src/graphql/index.ts

@@ -706,6 +706,29 @@ query getWorkspacePublicPages($workspaceId: String!) {
 }`,
 };
 
+export const getWorkspaceSubscriptionQuery = {
+  id: 'getWorkspaceSubscriptionQuery' as const,
+  operationName: 'getWorkspaceSubscription',
+  definitionName: 'workspace',
+  containsFile: false,
+  query: `
+query getWorkspaceSubscription($workspaceId: String!) {
+  workspace(id: $workspaceId) {
+    subscription {
+      id
+      status
+      plan
+      recurring
+      start
+      end
+      nextBillAt
+      canceledAt
+      variant
+    }
+  }
+}`,
+};
+
 export const getWorkspaceQuery = {
   id: 'getWorkspaceQuery' as const,
   operationName: 'getWorkspace',
@@ -1408,6 +1431,29 @@ mutation revokeInviteLink($workspaceId: String!) {
 }`,
 };
 
+export const workspaceInvoicesQuery = {
+  id: 'workspaceInvoicesQuery' as const,
+  operationName: 'workspaceInvoices',
+  definitionName: 'workspace',
+  containsFile: false,
+  query: `
+query workspaceInvoices($take: Int!, $skip: Int!, $workspaceId: String!) {
+  workspace(id: $workspaceId) {
+    invoiceCount
+    invoices(take: $take, skip: $skip) {
+      id
+      status
+      currency
+      amount
+      reason
+      lastPaymentError
+      link
+      createdAt
+    }
+  }
+}`,
+};
+
 export const workspaceQuotaQuery = {
   id: 'workspaceQuotaQuery' as const,
   operationName: 'workspaceQuota',

packages/frontend/graphql/src/graphql/workspace-invoices.gql

@@ -0,0 +1,15 @@
+query workspaceInvoices($take: Int!, $skip: Int!, $workspaceId: String!) {
+  workspace(id: $workspaceId) {
+    invoiceCount
+    invoices(take: $take, skip: $skip) {
+      id
+      status
+      currency
+      amount
+      reason
+      lastPaymentError
+      link
+      createdAt
+    }
+  }
+}

packages/frontend/graphql/src/schema.ts

@@ -2014,6 +2014,29 @@ export type GetWorkspacePublicPagesQuery = {
   };
 };
 
+export type GetWorkspaceSubscriptionQueryVariables = Exact<{
+  workspaceId: Scalars['String']['input'];
+}>;
+
+export type GetWorkspaceSubscriptionQuery = {
+  __typename?: 'Query';
+  workspace: {
+    __typename?: 'WorkspaceType';
+    subscription: {
+      __typename?: 'SubscriptionType';
+      id: string | null;
+      status: SubscriptionStatus;
+      plan: SubscriptionPlan;
+      recurring: SubscriptionRecurring;
+      start: string;
+      end: string | null;
+      nextBillAt: string | null;
+      canceledAt: string | null;
+      variant: SubscriptionVariant | null;
+    } | null;
+  };
+};
+
 export type GetWorkspaceQueryVariables = Exact<{
   id: Scalars['String']['input'];
 }>;
@@ -2627,6 +2650,31 @@ export type RevokeInviteLinkMutation = {
   revokeInviteLink: boolean;
 };
 
+export type WorkspaceInvoicesQueryVariables = Exact<{
+  take: Scalars['Int']['input'];
+  skip: Scalars['Int']['input'];
+  workspaceId: Scalars['String']['input'];
+}>;
+
+export type WorkspaceInvoicesQuery = {
+  __typename?: 'Query';
+  workspace: {
+    __typename?: 'WorkspaceType';
+    invoiceCount: number;
+    invoices: Array<{
+      __typename?: 'InvoiceType';
+      id: string | null;
+      status: InvoiceStatus;
+      currency: string;
+      amount: number;
+      reason: string;
+      lastPaymentError: string | null;
+      link: string | null;
+      createdAt: string;
+    }>;
+  };
+};
+
 export type WorkspaceQuotaQueryVariables = Exact<{
   id: Scalars['String']['input'];
 }>;
@@ -2813,6 +2861,11 @@ export type Queries =
       variables: GetWorkspacePublicPagesQueryVariables;
       response: GetWorkspacePublicPagesQuery;
     }
+  | {
+      name: 'getWorkspaceSubscriptionQuery';
+      variables: GetWorkspaceSubscriptionQueryVariables;
+      response: GetWorkspaceSubscriptionQuery;
+    }
   | {
       name: 'getWorkspaceQuery';
       variables: GetWorkspaceQueryVariables;
@@ -2883,6 +2936,11 @@ export type Queries =
       variables: ListWorkspaceFeaturesQueryVariables;
       response: ListWorkspaceFeaturesQuery;
     }
+  | {
+      name: 'workspaceInvoicesQuery';
+      variables: WorkspaceInvoicesQueryVariables;
+      response: WorkspaceInvoicesQuery;
+    }
   | {
       name: 'workspaceQuotaQuery';
       variables: WorkspaceQuotaQueryVariables;