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

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 } =
-            await this.store.fetchUserQuota(signal);
+          const { quota, used } = 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,
     };