fix(core): fix error when server not support ai (#6796)

packages/frontend/core/src/components/affine/setting-modal/account-setting/ai-usage-panel.tsx

@@ -4,7 +4,7 @@ import { openSettingModalAtom } from '@affine/core/atoms';
 import {
   ServerConfigService,
   SubscriptionService,
-  UserQuotaService,
+  UserCopilotQuotaService,
 } from '@affine/core/modules/cloud';
 import { useAFFiNEI18N } from '@affine/i18n/hooks';
 import { useLiveData, useService } from '@toeverything/infra';
@@ -28,14 +28,18 @@ export const AIUsagePanel = () => {
     // revalidate latest subscription status
     subscriptionService.subscription.revalidate();
   }, [subscriptionService]);
-  const quotaService = useService(UserQuotaService);
+  const copilotQuotaService = useService(UserCopilotQuotaService);
   useEffect(() => {
-    quotaService.quota.revalidate();
+    copilotQuotaService.copilotQuota.revalidate();
-  }, [quotaService]);
+  }, [copilotQuotaService]);
-  const aiActionLimit = useLiveData(quotaService.quota.aiActionLimit$);
+  const copilotActionLimit = useLiveData(
-  const aiActionUsed = useLiveData(quotaService.quota.aiActionUsed$);
+    copilotQuotaService.copilotQuota.copilotActionLimit$
-  const loading = aiActionLimit === null || aiActionUsed === null;
+  );
-  const loadError = useLiveData(quotaService.quota.error$);
+  const copilotActionUsed = useLiveData(
+    copilotQuotaService.copilotQuota.copilotActionUsed$
+  );
+  const loading = copilotActionLimit === null || copilotActionUsed === null;
+  const loadError = useLiveData(copilotQuotaService.copilotQuota.error$);
 
   const openBilling = useCallback(() => {
     setOpenSettingModal({
@@ -69,13 +73,13 @@ export const AIUsagePanel = () => {
   }
 
   const percent =
-    aiActionLimit === 'unlimited'
+    copilotActionLimit === 'unlimited'
       ? 0
       : Math.min(
           100,
           Math.max(
             0.5,
-            Number(((aiActionUsed / aiActionLimit) * 100).toFixed(4))
+            Number(((copilotActionUsed / copilotActionLimit) * 100).toFixed(4))
           )
         );
 
@@ -91,7 +95,7 @@ export const AIUsagePanel = () => {
       }
       name={t['com.affine.payment.ai.usage-title']()}
     >
-      {aiActionLimit === 'unlimited' ? (
+      {copilotActionLimit === 'unlimited' ? (
         hasPaymentFeature && aiSubscription?.canceledAt ? (
           <AIResume />
         ) : (
@@ -106,8 +110,8 @@ export const AIUsagePanel = () => {
               <span>{t['com.affine.payment.ai.usage.used-caption']()}</span>
               <span>
                 {t['com.affine.payment.ai.usage.used-detail']({
-                  used: aiActionUsed.toString(),
+                  used: copilotActionUsed.toString(),
-                  limit: aiActionLimit.toString(),
+                  limit: copilotActionLimit.toString(),
                 })}
               </span>
             </div>

packages/frontend/core/src/components/affine/setting-modal/account-setting/index.tsx

@@ -8,7 +8,12 @@ import { Button } from '@affine/component/ui/button';
 import { useAsyncCallback } from '@affine/core/hooks/affine-async-hooks';
 import { useAFFiNEI18N } from '@affine/i18n/hooks';
 import { ArrowRightSmallIcon, CameraIcon } from '@blocksuite/icons';
-import { useEnsureLiveData, useService } from '@toeverything/infra';
+import {
+  useEnsureLiveData,
+  useLiveData,
+  useService,
+  useServices,
+} from '@toeverything/infra';
 import { useSetAtom } from 'jotai';
 import type { FC, MouseEvent } from 'react';
 import { useCallback, useEffect, useState } from 'react';
@@ -18,7 +23,7 @@ import {
   openSettingModalAtom,
   openSignOutModalAtom,
 } from '../../../../atoms';
-import { AuthService } from '../../../../modules/cloud';
+import { AuthService, ServerConfigService } from '../../../../modules/cloud';
 import { mixpanel } from '../../../../utils';
 import { Upload } from '../../../pure/file-upload';
 import { AIUsagePanel } from './ai-usage-panel';
@@ -178,8 +183,15 @@ const StoragePanel = () => {
 };
 
 export const AccountSetting: FC = () => {
+  const { authService, serverConfigService } = useServices({
+    AuthService,
+    ServerConfigService,
+  });
+  const serverFeatures = useLiveData(
+    serverConfigService.serverConfig.features$
+  );
   const t = useAFFiNEI18N();
-  const session = useService(AuthService).session;
+  const session = authService.session;
   useEffect(() => {
     session.revalidate();
   }, [session]);
@@ -235,7 +247,7 @@ export const AccountSetting: FC = () => {
         </Button>
       </SettingRow>
       <StoragePanel />
-      <AIUsagePanel />
+      {serverFeatures?.copilot && <AIUsagePanel />}
       <SettingRow
         name={t[`Sign out`]()}
         desc={t['com.affine.setting.sign.out.message']()}

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

@@ -86,9 +86,6 @@ export class Subscription extends Entity {
             return undefined; // no subscription if no user
           }
 
-          // ensure server config is loaded
-          this.serverConfigService.serverConfig.revalidateIfNeeded();
-
           const serverConfig =
             await this.serverConfigService.serverConfig.features$.waitForNonNull(
               signal

packages/frontend/core/src/modules/cloud/entities/user-copilot-quota.ts

@@ -0,0 +1,100 @@
+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 { AuthService } from '../services/auth';
+import type { ServerConfigService } from '../services/server-config';
+import type { UserCopilotQuotaStore } from '../stores/user-copilot-quota';
+
+export class UserCopilotQuota extends Entity {
+  copilotActionLimit$ = new LiveData<number | 'unlimited' | null>(null);
+  copilotActionUsed$ = new LiveData<number | null>(null);
+
+  isRevalidating$ = new LiveData(false);
+  error$ = new LiveData<any | null>(null);
+
+  constructor(
+    private readonly authService: AuthService,
+    private readonly store: UserCopilotQuotaStore,
+    private readonly serverConfigService: ServerConfigService
+  ) {
+    super();
+  }
+
+  revalidate = effect(
+    map(() => ({
+      accountId: this.authService.session.account$.value?.id,
+    })),
+    exhaustMapSwitchUntilChanged(
+      (a, b) => a.accountId === b.accountId,
+      ({ accountId }) =>
+        fromPromise(async signal => {
+          if (!accountId) {
+            return; // no quota if no user
+          }
+
+          const serverConfig =
+            await this.serverConfigService.serverConfig.features$.waitForNonNull(
+              signal
+            );
+
+          let aiQuota = null;
+
+          if (serverConfig.copilot) {
+            aiQuota = await this.store.fetchUserCopilotQuota(signal);
+          }
+
+          return aiQuota;
+        }).pipe(
+          backoffRetry({
+            when: isNetworkError,
+            count: Infinity,
+          }),
+          backoffRetry({
+            when: isBackendError,
+          }),
+          mergeMap(data => {
+            if (data) {
+              const { limit, used } = data;
+              this.copilotActionUsed$.next(used);
+              this.copilotActionLimit$.next(
+                limit === null ? 'unlimited' : limit
+              ); // fix me: unlimited status
+            } else {
+              this.copilotActionUsed$.next(null);
+              this.copilotActionLimit$.next(null);
+            }
+            return EMPTY;
+          }),
+          catchErrorInto(this.error$),
+          onStart(() => this.isRevalidating$.next(true)),
+          onComplete(() => this.isRevalidating$.next(false))
+        ),
+      () => {
+        // Reset the state when the user is changed
+        this.reset();
+      }
+    )
+  );
+
+  reset() {
+    this.copilotActionUsed$.next(null);
+    this.copilotActionLimit$.next(null);
+    this.error$.next(null);
+    this.isRevalidating$.next(false);
+  }
+
+  override dispose(): void {
+    this.revalidate.unsubscribe();
+  }
+}

packages/frontend/core/src/modules/cloud/entities/user-quota.ts

@@ -31,9 +31,6 @@ export class UserQuota extends Entity {
   /** Maximum storage limit formatted */
   maxFormatted$ = this.max$.map(max => (max ? bytes.format(max) : null));
 
-  aiActionLimit$ = new LiveData<number | 'unlimited' | null>(null);
-  aiActionUsed$ = new LiveData<number | null>(null);
-
   /** Percentage of storage used */
   percent$ = LiveData.computed(get => {
     const max = get(this.max$);
@@ -76,10 +73,9 @@ export class UserQuota extends Entity {
           if (!accountId) {
             return; // no quota if no user
           }
-          const { quota, aiQuota, used } =
+          const { quota, used } = await this.store.fetchUserQuota(signal);
-            await this.store.fetchUserQuota(signal);
 
-          return { quota, aiQuota, used };
+          return { quota, used };
         }).pipe(
           backoffRetry({
             when: isNetworkError,
@@ -90,18 +86,12 @@ export class UserQuota extends Entity {
           }),
           mergeMap(data => {
             if (data) {
-              const { aiQuota, quota, used } = data;
+              const { quota, used } = data;
               this.quota$.next(quota);
               this.used$.next(used);
-              this.aiActionUsed$.next(aiQuota.used);
-              this.aiActionLimit$.next(
-                aiQuota.limit === null ? 'unlimited' : aiQuota.limit
-              ); // fix me: unlimited status
             } else {
               this.quota$.next(null);
               this.used$.next(null);
-              this.aiActionUsed$.next(null);
-              this.aiActionLimit$.next(null);
             }
             return EMPTY;
           }),
@@ -119,8 +109,6 @@ export class UserQuota extends Entity {
   reset() {
     this.quota$.next(null);
     this.used$.next(null);
-    this.aiActionUsed$.next(null);
-    this.aiActionLimit$.next(null);
     this.error$.next(null);
     this.isRevalidating$.next(false);
   }

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

@@ -10,6 +10,7 @@ export { FetchService } from './services/fetch';
 export { GraphQLService } from './services/graphql';
 export { ServerConfigService } from './services/server-config';
 export { SubscriptionService } from './services/subscription';
+export { UserCopilotQuotaService } from './services/user-copilot-quota';
 export { UserFeatureService } from './services/user-feature';
 export { UserQuotaService } from './services/user-quota';
 export { WebSocketService } from './services/websocket';
@@ -24,6 +25,7 @@ import { ServerConfig } from './entities/server-config';
 import { AuthSession } from './entities/session';
 import { Subscription } from './entities/subscription';
 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 { AuthService } from './services/auth';
@@ -31,12 +33,14 @@ import { FetchService } from './services/fetch';
 import { GraphQLService } from './services/graphql';
 import { ServerConfigService } from './services/server-config';
 import { SubscriptionService } from './services/subscription';
+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 { AuthStore } from './stores/auth';
 import { ServerConfigStore } from './stores/server-config';
 import { SubscriptionStore } from './stores/subscription';
+import { UserCopilotQuotaStore } from './stores/user-copilot-quota';
 import { UserFeatureStore } from './stores/user-feature';
 import { UserQuotaStore } from './stores/user-quota';
 
@@ -58,6 +62,13 @@ export function configureCloudModule(framework: Framework) {
     .service(UserQuotaService)
     .store(UserQuotaStore, [GraphQLService])
     .entity(UserQuota, [AuthService, UserQuotaStore])
+    .service(UserCopilotQuotaService)
+    .store(UserCopilotQuotaStore, [GraphQLService])
+    .entity(UserCopilotQuota, [
+      AuthService,
+      UserCopilotQuotaStore,
+      ServerConfigService,
+    ])
     .service(UserFeatureService)
     .entity(UserFeature, [AuthService, UserFeatureStore])
     .store(UserFeatureStore, [GraphQLService]);

packages/frontend/core/src/modules/cloud/services/user-copilot-quota.ts

@@ -0,0 +1,13 @@
+import { OnEvent, Service } from '@toeverything/infra';
+
+import { UserCopilotQuota } from '../entities/user-copilot-quota';
+import { AccountChanged } from './auth';
+
+@OnEvent(AccountChanged, e => e.onAccountChanged)
+export class UserCopilotQuotaService extends Service {
+  copilotQuota = this.framework.createEntity(UserCopilotQuota);
+
+  private onAccountChanged() {
+    this.copilotQuota.revalidate();
+  }
+}

packages/frontend/core/src/modules/cloud/stores/user-copilot-quota.ts

@@ -0,0 +1,25 @@
+import { copilotQuotaQuery } from '@affine/graphql';
+import { Store } from '@toeverything/infra';
+
+import type { GraphQLService } from '../services/graphql';
+
+export class UserCopilotQuotaStore extends Store {
+  constructor(private readonly graphqlService: GraphQLService) {
+    super();
+  }
+
+  async fetchUserCopilotQuota(abortSignal?: AbortSignal) {
+    const data = await this.graphqlService.gql({
+      query: copilotQuotaQuery,
+      context: {
+        signal: abortSignal,
+      },
+    });
+
+    if (!data.currentUser) {
+      throw new Error('No logged in');
+    }
+
+    return data.currentUser.copilot.quota;
+  }
+}

packages/frontend/core/src/modules/cloud/stores/user-quota.ts

@@ -22,7 +22,6 @@ export class UserQuotaStore extends Store {
 
     return {
       userId: data.currentUser.id,
-      aiQuota: data.currentUser.copilot.quota,
       quota: data.currentUser.quota,
       used: data.collectAllBlobSizes.size,
     };

packages/frontend/graphql/src/graphql/copilot-quota.gql

@@ -1,4 +1,4 @@
-query getCopilotQuota {
+query copilotQuota {
   currentUser {
     copilot {
       quota {

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

@@ -94,6 +94,24 @@ mutation changePassword($token: String!, $newPassword: String!) {
 }`,
 };
 
+export const copilotQuotaQuery = {
+  id: 'copilotQuotaQuery' as const,
+  operationName: 'copilotQuota',
+  definitionName: 'currentUser',
+  containsFile: false,
+  query: `
+query copilotQuota {
+  currentUser {
+    copilot {
+      quota {
+        limit
+        used
+      }
+    }
+  }
+}`,
+};
+
 export const createCheckoutSessionMutation = {
   id: 'createCheckoutSessionMutation' as const,
   operationName: 'createCheckoutSession',
@@ -238,24 +256,6 @@ query getCopilotHistories($workspaceId: String!, $docId: String, $options: Query
 }`,
 };
 
-export const getCopilotQuotaQuery = {
-  id: 'getCopilotQuotaQuery' as const,
-  operationName: 'getCopilotQuota',
-  definitionName: 'currentUser',
-  containsFile: false,
-  query: `
-query getCopilotQuota {
-  currentUser {
-    copilot {
-      quota {
-        limit
-        used
-      }
-    }
-  }
-}`,
-};
-
 export const getCopilotSessionsQuery = {
   id: 'getCopilotSessionsQuery' as const,
   operationName: 'getCopilotSessions',
@@ -607,12 +607,6 @@ export const quotaQuery = {
 query quota {
   currentUser {
     id
-    copilot {
-      quota {
-        limit
-        used
-      }
-    }
     quota {
       name
       blobLimit

packages/frontend/graphql/src/graphql/quota.gql

@@ -1,12 +1,6 @@
 query quota {
   currentUser {
     id
-    copilot {
-      quota {
-        limit
-        used
-      }
-    }
     quota {
       name
       blobLimit

packages/frontend/graphql/src/schema.ts

@@ -213,6 +213,23 @@ export type ChangePasswordMutation = {
   changePassword: { __typename?: 'UserType'; id: string };
 };
 
+export type CopilotQuotaQueryVariables = Exact<{ [key: string]: never }>;
+
+export type CopilotQuotaQuery = {
+  __typename?: 'Query';
+  currentUser: {
+    __typename?: 'UserType';
+    copilot: {
+      __typename?: 'Copilot';
+      quota: {
+        __typename?: 'CopilotQuota';
+        limit: number | null;
+        used: number;
+      };
+    };
+  } | null;
+};
+
 export type CreateCheckoutSessionMutationVariables = Exact<{
   input: CreateCheckoutSessionInput;
 }>;
@@ -353,23 +370,6 @@ export type GetCopilotHistoriesQuery = {
   } | null;
 };
 
-export type GetCopilotQuotaQueryVariables = Exact<{ [key: string]: never }>;
-
-export type GetCopilotQuotaQuery = {
-  __typename?: 'Query';
-  currentUser: {
-    __typename?: 'UserType';
-    copilot: {
-      __typename?: 'Copilot';
-      quota: {
-        __typename?: 'CopilotQuota';
-        limit: number | null;
-        used: number;
-      };
-    };
-  } | null;
-};
-
 export type GetCopilotSessionsQueryVariables = Exact<{
   workspaceId: Scalars['String']['input'];
 }>;
@@ -677,14 +677,6 @@ export type QuotaQuery = {
   currentUser: {
     __typename?: 'UserType';
     id: string;
-    copilot: {
-      __typename?: 'Copilot';
-      quota: {
-        __typename?: 'CopilotQuota';
-        limit: number | null;
-        used: number;
-      };
-    };
     quota: {
       __typename?: 'UserQuota';
       name: string;
@@ -1038,6 +1030,11 @@ export type Queries =
       variables: ListBlobsQueryVariables;
       response: ListBlobsQuery;
     }
+  | {
+      name: 'copilotQuotaQuery';
+      variables: CopilotQuotaQueryVariables;
+      response: CopilotQuotaQuery;
+    }
   | {
       name: 'earlyAccessUsersQuery';
       variables: EarlyAccessUsersQueryVariables;
@@ -1048,11 +1045,6 @@ export type Queries =
       variables: GetCopilotHistoriesQueryVariables;
       response: GetCopilotHistoriesQuery;
     }
-  | {
-      name: 'getCopilotQuotaQuery';
-      variables: GetCopilotQuotaQueryVariables;
-      response: GetCopilotQuotaQuery;
-    }
   | {
       name: 'getCopilotSessionsQuery';
       variables: GetCopilotSessionsQueryVariables;