Skip to content

Commit

Permalink
fix(core): fix error when server not support ai (#6796)
Browse files Browse the repository at this point in the history
  • Loading branch information
EYHN committed May 7, 2024
1 parent a0e0b6b commit 35ce4ad
Show file tree
Hide file tree
Showing 13 changed files with 226 additions and 97 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand All @@ -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({
Expand Down Expand Up @@ -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))
)
);

Expand All @@ -91,7 +95,7 @@ export const AIUsagePanel = () => {
}
name={t['com.affine.payment.ai.usage-title']()}
>
{aiActionLimit === 'unlimited' ? (
{copilotActionLimit === 'unlimited' ? (
hasPaymentFeature && aiSubscription?.canceledAt ? (
<AIResume />
) : (
Expand All @@ -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(),
limit: aiActionLimit.toString(),
used: copilotActionUsed.toString(),
limit: copilotActionLimit.toString(),
})}
</span>
</div>
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand All @@ -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';
Expand Down Expand Up @@ -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]);
Expand Down Expand Up @@ -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']()}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
@@ -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();
}
}
18 changes: 3 additions & 15 deletions packages/frontend/core/src/modules/cloud/entities/user-quota.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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$);
Expand Down Expand Up @@ -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,
Expand All @@ -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;
}),
Expand All @@ -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);
}
Expand Down
11 changes: 11 additions & 0 deletions packages/frontend/core/src/modules/cloud/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand All @@ -24,19 +25,22 @@ 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';
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';

Expand All @@ -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]);
Expand Down
Original file line number Diff line number Diff line change
@@ -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();
}
}
Original file line number Diff line number Diff line change
@@ -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;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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,
};
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
query getCopilotQuota {
query copilotQuota {
currentUser {
copilot {
quota {
Expand Down
Loading

0 comments on commit 35ce4ad

Please sign in to comment.