diff --git a/.github/workflows/release-desktop.yml b/.github/workflows/release-desktop.yml index 237de8ced3..a44e2a252a 100644 --- a/.github/workflows/release-desktop.yml +++ b/.github/workflows/release-desktop.yml @@ -325,7 +325,7 @@ jobs: mkdir -p builds mv packages/frontend/electron/out/*/make/zip/win32/x64/AFFiNE*-win32-x64-*.zip ./builds/affine-${{ needs.before-make.outputs.RELEASE_VERSION }}-${{ env.BUILD_TYPE }}-windows-x64.zip mv packages/frontend/electron/out/*/make/squirrel.windows/x64/*.exe ./builds/affine-${{ needs.before-make.outputs.RELEASE_VERSION }}-${{ env.BUILD_TYPE }}-windows-x64.exe - mv packages/frontend/electron/out/*/make/nsis.windows/x64/*.exe ./builds/affine-${{ needs.before-make.outputs.RELEASE_VERSION }}-${{ env.BUILD_TYPE }}-windows-x64.snis.exe + mv packages/frontend/electron/out/*/make/nsis.windows/x64/*.exe ./builds/affine-${{ needs.before-make.outputs.RELEASE_VERSION }}-${{ env.BUILD_TYPE }}-windows-x64.nsis.exe - name: Upload Artifact uses: actions/upload-artifact@v4 diff --git a/packages/backend/server/src/core/auth/resolver.ts b/packages/backend/server/src/core/auth/resolver.ts index b347ca1f28..baa32a7dec 100644 --- a/packages/backend/server/src/core/auth/resolver.ts +++ b/packages/backend/server/src/core/auth/resolver.ts @@ -12,7 +12,7 @@ import { } from '@nestjs/graphql'; import type { Request, Response } from 'express'; -import { Config, Throttle } from '../../fundamentals'; +import { Config, SkipThrottle, Throttle } from '../../fundamentals'; import { UserService } from '../user'; import { UserType } from '../user/types'; import { validators } from '../utils/validators'; @@ -33,12 +33,6 @@ export class ClientTokenType { sessionToken?: string; } -/** - * Auth resolver - * Token rate limit: 20 req/m - * Sign up/in rate limit: 10 req/m - * Other rate limit: 5 req/m - */ @Throttle('strict') @Resolver(() => UserType) export class AuthResolver { @@ -49,6 +43,7 @@ export class AuthResolver { private readonly token: TokenService ) {} + @SkipThrottle() @Public() @Query(() => UserType, { name: 'currentUser', diff --git a/packages/backend/server/src/fundamentals/index.ts b/packages/backend/server/src/fundamentals/index.ts index 5abf98febb..5060d35432 100644 --- a/packages/backend/server/src/fundamentals/index.ts +++ b/packages/backend/server/src/fundamentals/index.ts @@ -27,7 +27,7 @@ export { export type { PrismaTransaction } from './prisma'; export * from './storage'; export { type StorageProvider, StorageProviderFactory } from './storage'; -export { CloudThrottlerGuard, Throttle } from './throttler'; +export { CloudThrottlerGuard, SkipThrottle, Throttle } from './throttler'; export { getRequestFromHost, getRequestResponseFromContext, diff --git a/packages/backend/server/src/fundamentals/throttler/decorators.ts b/packages/backend/server/src/fundamentals/throttler/decorators.ts index 1baa7d9dd0..742a32d729 100644 --- a/packages/backend/server/src/fundamentals/throttler/decorators.ts +++ b/packages/backend/server/src/fundamentals/throttler/decorators.ts @@ -1,7 +1,7 @@ import { applyDecorators, SetMetadata } from '@nestjs/common'; import { SkipThrottle, Throttle as RawThrottle } from '@nestjs/throttler'; -export type Throttlers = 'default' | 'strict'; +export type Throttlers = 'default' | 'strict' | 'authenticated'; export const THROTTLER_PROTECTED = 'affine_throttler:protected'; /** @@ -10,8 +10,9 @@ export const THROTTLER_PROTECTED = 'affine_throttler:protected'; * If a Controller or Query do not protected behind a Throttler, * it will never be rate limited. * - * - Ease: 120 calls within 60 seconds - * - Strict: 10 calls within 60 seconds + * - default: 120 calls within 60 seconds + * - strict: 10 calls within 60 seconds + * - authenticated: no rate limit for authenticated users, apply [default] throttler for unauthenticated users * * @example * diff --git a/packages/backend/server/src/fundamentals/throttler/index.ts b/packages/backend/server/src/fundamentals/throttler/index.ts index a75490f093..f15f408c12 100644 --- a/packages/backend/server/src/fundamentals/throttler/index.ts +++ b/packages/backend/server/src/fundamentals/throttler/index.ts @@ -166,10 +166,12 @@ export class CloudThrottlerGuard extends ThrottlerGuard { } getSpecifiedThrottler(context: ExecutionContext) { - return this.reflector.getAllAndOverride( + const throttler = this.reflector.getAllAndOverride( THROTTLER_PROTECTED, [context.getHandler(), context.getClass()] ); + + return throttler === 'authenticated' ? undefined : throttler; } } diff --git a/packages/backend/server/src/plugins/copilot/prompt.ts b/packages/backend/server/src/plugins/copilot/prompt.ts index c71c4fc2fd..74c51127e8 100644 --- a/packages/backend/server/src/plugins/copilot/prompt.ts +++ b/packages/backend/server/src/plugins/copilot/prompt.ts @@ -1,4 +1,4 @@ -import { Injectable } from '@nestjs/common'; +import { Injectable, Logger } from '@nestjs/common'; import { AiPrompt, PrismaClient } from '@prisma/client'; import Mustache from 'mustache'; import { Tiktoken } from 'tiktoken'; @@ -26,6 +26,7 @@ function extractMustacheParams(template: string) { } export class ChatPrompt { + private readonly logger = new Logger(ChatPrompt.name); public readonly encoder?: Tiktoken; private readonly promptTokenSize: number; private readonly templateParamKeys: string[] = []; @@ -88,7 +89,7 @@ export class ChatPrompt { return this.encoder?.encode_ordinary(message).length || 0; } - private checkParams(params: PromptParams) { + private checkParams(params: PromptParams, sessionId?: string) { const selfParams = this.templateParams; for (const key of Object.keys(selfParams)) { const options = selfParams[key]; @@ -97,7 +98,20 @@ export class ChatPrompt { typeof income !== 'string' || (Array.isArray(options) && !options.includes(income)) ) { - throw new Error(`Invalid param: ${key}`); + if (sessionId) { + const prefix = income + ? `Invalid param value: ${key}=${income}` + : `Missing param value: ${key}`; + this.logger.warn( + `${prefix} in session ${sessionId}, use default options: ${options[0]}` + ); + } + if (Array.isArray(options)) { + // use the first option if income is not in options + params[key] = options[0]; + } else { + params[key] = options; + } } } } @@ -107,8 +121,8 @@ export class ChatPrompt { * @param params record of params, e.g. { name: 'Alice' } * @returns e.g. [{ role: 'system', content: 'Hello, {{name}}' }] => [{ role: 'system', content: 'Hello, Alice' }] */ - finish(params: PromptParams): PromptMessage[] { - this.checkParams(params); + finish(params: PromptParams, sessionId?: string): PromptMessage[] { + this.checkParams(params, sessionId); return this.messages.map(({ content, params: _, ...rest }) => ({ ...rest, params, diff --git a/packages/backend/server/src/plugins/copilot/session.ts b/packages/backend/server/src/plugins/copilot/session.ts index aa5b033a3c..c7e28c22c1 100644 --- a/packages/backend/server/src/plugins/copilot/session.ts +++ b/packages/backend/server/src/plugins/copilot/session.ts @@ -112,7 +112,8 @@ export class ChatSession implements AsyncDisposable { const messages = this.takeMessages(); return [ ...this.state.prompt.finish( - Object.keys(params).length ? params : messages[0]?.params || {} + Object.keys(params).length ? params : messages[0]?.params || {}, + this.config.sessionId ), ...messages.filter(m => m.content || m.attachments?.length), ]; @@ -354,7 +355,7 @@ export class ChatSessionService { // render system prompt const preload = withPrompt ? prompt - .finish(ret.data[0]?.params || {}) + .finish(ret.data[0]?.params || {}, id) .filter(({ role }) => role !== 'system') : []; diff --git a/packages/backend/server/src/plugins/payment/service.ts b/packages/backend/server/src/plugins/payment/service.ts index 758b8bd677..68dd218a00 100644 --- a/packages/backend/server/src/plugins/payment/service.ts +++ b/packages/backend/server/src/plugins/payment/service.ts @@ -95,11 +95,8 @@ export class SubscriptionService { }); oldSubscriptions.data.forEach(sub => { - if ( - (sub.status === 'past_due' || sub.status === 'canceled') && - sub.items.data[0].price.lookup_key - ) { - const [oldPlan] = decodeLookupKey(sub.items.data[0].price.lookup_key); + if (sub.status === 'past_due' || sub.status === 'canceled') { + const [oldPlan] = this.decodePlanFromSubscription(sub); if (oldPlan === SubscriptionPlan.Pro) { canHaveEarlyAccessDiscount = false; } @@ -418,9 +415,13 @@ export class SubscriptionService { @OnEvent('customer.subscription.created') @OnEvent('customer.subscription.updated') async onSubscriptionChanges(subscription: Stripe.Subscription) { + // webhook call may not in sequential order, get the latest status + subscription = await this.stripe.subscriptions.retrieve(subscription.id); if (subscription.status === 'active') { const user = await this.retrieveUserFromCustomer( - subscription.customer as string + typeof subscription.customer === 'string' + ? subscription.customer + : subscription.customer.id ); await this.saveSubscription(user, subscription); @@ -431,6 +432,19 @@ export class SubscriptionService { @OnEvent('customer.subscription.deleted') async onSubscriptionDeleted(subscription: Stripe.Subscription) { + subscription = await this.stripe.subscriptions.retrieve(subscription.id); + const user = await this.retrieveUserFromCustomer( + typeof subscription.customer === 'string' + ? subscription.customer + : subscription.customer.id + ); + + const [plan] = this.decodePlanFromSubscription(subscription); + this.event.emit('user.subscription.canceled', { + userId: user.id, + plan, + }); + await this.db.userSubscription.deleteMany({ where: { stripeSubscriptionId: subscription.id, @@ -440,6 +454,7 @@ export class SubscriptionService { @OnEvent('invoice.paid') async onInvoicePaid(stripeInvoice: Stripe.Invoice) { + stripeInvoice = await this.stripe.invoices.retrieve(stripeInvoice.id); await this.saveInvoice(stripeInvoice); const line = stripeInvoice.lines.data[0]; @@ -453,6 +468,7 @@ export class SubscriptionService { @OnEvent('invoice.finalization_failed') @OnEvent('invoice.payment_failed') async saveInvoice(stripeInvoice: Stripe.Invoice) { + stripeInvoice = await this.stripe.invoices.retrieve(stripeInvoice.id); if (!stripeInvoice.customer) { throw new Error('Unexpected invoice with no customer'); } @@ -552,26 +568,14 @@ export class SubscriptionService { throw new Error('Unexpected subscription with no key'); } - const [plan, recurring] = decodeLookupKey(price.lookup_key); + const [plan, recurring] = this.decodePlanFromSubscription(subscription); const planActivated = SubscriptionActivated.includes(subscription.status); let nextBillAt: Date | null = null; - if (planActivated) { - this.event.emit('user.subscription.activated', { - userId: user.id, - plan, - }); - + if (planActivated && !subscription.canceled_at) { // get next bill date from upcoming invoice // see https://stripe.com/docs/api/invoices/upcoming - if (!subscription.canceled_at) { - nextBillAt = new Date(subscription.current_period_end * 1000); - } - } else { - this.event.emit('user.subscription.canceled', { - userId: user.id, - plan, - }); + nextBillAt = new Date(subscription.current_period_end * 1000); } const commonData = { @@ -620,6 +624,11 @@ export class SubscriptionService { data: update, }); } else { + this.event.emit('user.subscription.activated', { + userId: user.id, + plan, + }); + return await this.db.userSubscription.create({ data: { userId: user.id, @@ -749,14 +758,11 @@ export class SubscriptionService { }); const subscribed = oldSubscriptions.data.some(sub => { - if (sub.items.data[0].price.lookup_key) { - const [oldPlan] = decodeLookupKey(sub.items.data[0].price.lookup_key); - return ( - oldPlan === plan && - (sub.status === 'past_due' || sub.status === 'canceled') - ); - } - return false; + const [oldPlan] = this.decodePlanFromSubscription(sub); + return ( + oldPlan === plan && + (sub.status === 'past_due' || sub.status === 'canceled') + ); }); if (plan === SubscriptionPlan.Pro) { @@ -830,4 +836,14 @@ export class SubscriptionService { return available ? code.id : null; } + + private decodePlanFromSubscription(sub: Stripe.Subscription) { + const price = sub.items.data[0].price; + + if (!price.lookup_key) { + throw new Error('Unexpected subscription with no key'); + } + + return decodeLookupKey(price.lookup_key); + } } diff --git a/packages/backend/server/tests/copilot.spec.ts b/packages/backend/server/tests/copilot.spec.ts index 141e6c477b..75b023ec3d 100644 --- a/packages/backend/server/tests/copilot.spec.ts +++ b/packages/backend/server/tests/copilot.spec.ts @@ -105,9 +105,14 @@ test('should be able to render prompt', async t => { 'should have param keys' ); t.deepEqual(testPrompt?.params, msg.params, 'should have params'); - t.throws(() => testPrompt?.finish({ src_language: 'abc' }), { - instanceOf: Error, - }); + // will use first option if a params not provided + t.deepEqual(testPrompt?.finish({ src_language: 'abc' }), [ + { + content: 'translate eng to chs: ', + params: { dest_language: 'chs', src_language: 'eng' }, + role: 'system', + }, + ]); }); test('should be able to render listed prompt', async t => { diff --git a/packages/backend/server/tests/nestjs/throttler.spec.ts b/packages/backend/server/tests/nestjs/throttler.spec.ts index 77921961f5..2c113b1fcc 100644 --- a/packages/backend/server/tests/nestjs/throttler.spec.ts +++ b/packages/backend/server/tests/nestjs/throttler.spec.ts @@ -48,6 +48,13 @@ class ThrottledController { return 'default3'; } + @Public() + @Get('/authenticated') + @Throttle('authenticated') + none() { + return 'none'; + } + @Throttle('strict') @Get('/strict') strict() { @@ -156,7 +163,6 @@ test('should use default throttler for unauthenticated user when not specified', t.is(headers.limit, '120'); t.is(headers.remaining, '119'); - t.regex(headers.reset, /59|60/); }); test('should skip throttler for unauthenticated user when specified', async t => { @@ -192,7 +198,6 @@ test('should use specified throttler for unauthenticated user', async t => { t.is(headers.limit, '20'); t.is(headers.remaining, '19'); - t.regex(headers.reset, /59|60/); }); // ==== authenticated user visits ==== @@ -223,7 +228,6 @@ test('should use default throttler for authenticated user when not specified', a t.is(headers.limit, '120'); t.is(headers.remaining, '119'); - t.regex(headers.reset, /59|60/); }); test('should use same throttler for multiple routes', async t => { @@ -238,7 +242,6 @@ test('should use same throttler for multiple routes', async t => { t.is(headers.limit, '120'); t.is(headers.remaining, '119'); - t.regex(headers.reset, /59|60/); res = await request(app.getHttpServer()) .get('/throttled/default2') @@ -263,7 +266,6 @@ test('should use different throttler if specified', async t => { t.is(headers.limit, '120'); t.is(headers.remaining, '119'); - t.regex(headers.reset, /59|60/); res = await request(app.getHttpServer()) .get('/throttled/default3') @@ -274,7 +276,34 @@ test('should use different throttler if specified', async t => { t.is(headers.limit, '10'); t.is(headers.remaining, '9'); - t.regex(headers.reset, /59|60/); +}); + +test('should skip throttler for authenticated if `authenticated` throttler used', async t => { + const { app, cookie } = t.context; + + const res = await request(app.getHttpServer()) + .get('/throttled/authenticated') + .set('Cookie', cookie) + .expect(200); + + const headers = rateLimitHeaders(res); + + t.is(headers.limit, undefined!); + t.is(headers.remaining, undefined!); + t.is(headers.reset, undefined!); +}); + +test('should apply `default` throttler for authenticated user if `authenticated` throttler used', async t => { + const { app } = t.context; + + const res = await request(app.getHttpServer()) + .get('/throttled/authenticated') + .expect(200); + + const headers = rateLimitHeaders(res); + + t.is(headers.limit, '120'); + t.is(headers.remaining, '119'); }); test('should skip throttler for authenticated user when specified', async t => { @@ -304,7 +333,6 @@ test('should use specified throttler for authenticated user', async t => { t.is(headers.limit, '20'); t.is(headers.remaining, '19'); - t.regex(headers.reset, /59|60/); }); test('should separate anonymous and authenticated user throttlers', async t => { @@ -323,9 +351,7 @@ test('should separate anonymous and authenticated user throttlers', async t => { t.is(authenticatedResHeaders.limit, '120'); t.is(authenticatedResHeaders.remaining, '119'); - t.regex(authenticatedResHeaders.reset, /59|60/); t.is(unauthenticatedResHeaders.limit, '120'); t.is(unauthenticatedResHeaders.remaining, '119'); - t.regex(unauthenticatedResHeaders.reset, /59|60/); });