diff --git a/packages/frontend/core/src/components/affine/setting-modal/account-setting/ai-usage-panel.tsx b/packages/frontend/core/src/components/affine/setting-modal/account-setting/ai-usage-panel.tsx
index 7f434b5a5..3b1e12161 100644
--- a/packages/frontend/core/src/components/affine/setting-modal/account-setting/ai-usage-panel.tsx
+++ b/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();
- }, [quotaService]);
- const aiActionLimit = useLiveData(quotaService.quota.aiActionLimit$);
- const aiActionUsed = useLiveData(quotaService.quota.aiActionUsed$);
- const loading = aiActionLimit === null || aiActionUsed === null;
- const loadError = useLiveData(quotaService.quota.error$);
+ copilotQuotaService.copilotQuota.revalidate();
+ }, [copilotQuotaService]);
+ const copilotActionLimit = useLiveData(
+ copilotQuotaService.copilotQuota.copilotActionLimit$
+ );
+ 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 ? (
) : (
@@ -106,8 +110,8 @@ export const AIUsagePanel = () => {
{t['com.affine.payment.ai.usage.used-caption']()}
{t['com.affine.payment.ai.usage.used-detail']({
- used: aiActionUsed.toString(),
- limit: aiActionLimit.toString(),
+ used: copilotActionUsed.toString(),
+ limit: copilotActionLimit.toString(),
})}
diff --git a/packages/frontend/core/src/components/affine/setting-modal/account-setting/index.tsx b/packages/frontend/core/src/components/affine/setting-modal/account-setting/index.tsx
index 0aad1bbfd..1800b407b 100644
--- a/packages/frontend/core/src/components/affine/setting-modal/account-setting/index.tsx
+++ b/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 = () => {
-
+ {serverFeatures?.copilot && }
(null);
+ copilotActionUsed$ = new LiveData(null);
+
+ isRevalidating$ = new LiveData(false);
+ error$ = new LiveData(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();
+ }
+}
diff --git a/packages/frontend/core/src/modules/cloud/entities/user-quota.ts b/packages/frontend/core/src/modules/cloud/entities/user-quota.ts
index 3294d9abf..01110aa06 100644
--- a/packages/frontend/core/src/modules/cloud/entities/user-quota.ts
+++ b/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(null);
- aiActionUsed$ = new LiveData(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);
}
diff --git a/packages/frontend/core/src/modules/cloud/index.ts b/packages/frontend/core/src/modules/cloud/index.ts
index 0bb6e2766..499840d6a 100644
--- a/packages/frontend/core/src/modules/cloud/index.ts
+++ b/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]);
diff --git a/packages/frontend/core/src/modules/cloud/services/user-copilot-quota.ts b/packages/frontend/core/src/modules/cloud/services/user-copilot-quota.ts
new file mode 100644
index 000000000..740c38046
--- /dev/null
+++ b/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();
+ }
+}
diff --git a/packages/frontend/core/src/modules/cloud/stores/user-copilot-quota.ts b/packages/frontend/core/src/modules/cloud/stores/user-copilot-quota.ts
new file mode 100644
index 000000000..ebe5b5737
--- /dev/null
+++ b/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;
+ }
+}
diff --git a/packages/frontend/core/src/modules/cloud/stores/user-quota.ts b/packages/frontend/core/src/modules/cloud/stores/user-quota.ts
index 84819127b..0413d46dc 100644
--- a/packages/frontend/core/src/modules/cloud/stores/user-quota.ts
+++ b/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,
};
diff --git a/packages/frontend/graphql/src/graphql/get-copilot-quota.gql b/packages/frontend/graphql/src/graphql/copilot-quota.gql
similarity index 79%
rename from packages/frontend/graphql/src/graphql/get-copilot-quota.gql
rename to packages/frontend/graphql/src/graphql/copilot-quota.gql
index 5ec67af2c..0f3570e5a 100644
--- a/packages/frontend/graphql/src/graphql/get-copilot-quota.gql
+++ b/packages/frontend/graphql/src/graphql/copilot-quota.gql
@@ -1,4 +1,4 @@
-query getCopilotQuota {
+query copilotQuota {
currentUser {
copilot {
quota {
diff --git a/packages/frontend/graphql/src/graphql/index.ts b/packages/frontend/graphql/src/graphql/index.ts
index d381e9ec2..3ab0d2e1a 100644
--- a/packages/frontend/graphql/src/graphql/index.ts
+++ b/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
diff --git a/packages/frontend/graphql/src/graphql/quota.gql b/packages/frontend/graphql/src/graphql/quota.gql
index 176828f00..c0268b644 100644
--- a/packages/frontend/graphql/src/graphql/quota.gql
+++ b/packages/frontend/graphql/src/graphql/quota.gql
@@ -1,12 +1,6 @@
query quota {
currentUser {
id
- copilot {
- quota {
- limit
- used
- }
- }
quota {
name
blobLimit
diff --git a/packages/frontend/graphql/src/schema.ts b/packages/frontend/graphql/src/schema.ts
index 62a37bc87..b8a12bae5 100644
--- a/packages/frontend/graphql/src/schema.ts
+++ b/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;