Merge commit '2c228a35f8405e18c2886c7503f7027491efe7a3' into beta

This commit is contained in:
DarkSky
2024-04-25 20:40:12 +08:00
10 changed files with 121 additions and 61 deletions
+1 -1
View File
@@ -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
@@ -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',
@@ -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,
@@ -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
*
@@ -166,10 +166,12 @@ export class CloudThrottlerGuard extends ThrottlerGuard {
}
getSpecifiedThrottler(context: ExecutionContext) {
return this.reflector.getAllAndOverride<Throttlers | undefined>(
const throttler = this.reflector.getAllAndOverride<Throttlers | undefined>(
THROTTLER_PROTECTED,
[context.getHandler(), context.getClass()]
);
return throttler === 'authenticated' ? undefined : throttler;
}
}
@@ -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,
@@ -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')
: [];
@@ -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);
}
}
@@ -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 => {
@@ -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/);
});