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

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,