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;