diff --git a/src/rate-limit.ts b/src/rate-limit.ts index 2104e30..12bff6f 100644 --- a/src/rate-limit.ts +++ b/src/rate-limit.ts @@ -36,6 +36,7 @@ const redis = createRedisClient({ export const rateLimitPlugin: FastifyPluginAsync = fp.default(async (app) => { await app.register(rateLimit, { global: true, + enableDraftSpec: true, max: (request: FastifyRequest) => { const authType = determineAuthType(request); switch (authType.type) { diff --git a/src/runs/runs.service.ts b/src/runs/runs.service.ts index 79f1b64..98a7a49 100644 --- a/src/runs/runs.service.ts +++ b/src/runs/runs.service.ts @@ -69,12 +69,14 @@ import { ensureRequestContextData } from '@/context.js'; import { getProjectPrincipal } from '@/administration/helpers.js'; import { RUNS_QUOTA_DAILY } from '@/config.js'; import { dayjs, getLatestDailyFixedTime } from '@/utils/datetime.js'; +import { updateRateLimitHeadersWithDailyQuota } from '@/utils/rate-limit.js'; export async function assertRunsQuota(newRuns = 1) { const count = await ORM.em.getRepository(Run).count({ createdBy: getProjectPrincipal(), createdAt: { $gte: getLatestDailyFixedTime().toDate() } }); + updateRateLimitHeadersWithDailyQuota({ quota: RUNS_QUOTA_DAILY, used: count }); if (count + newRuns > RUNS_QUOTA_DAILY) { throw new APIError({ message: 'Your daily runs quota has been exceeded', diff --git a/src/streaming/sse.ts b/src/streaming/sse.ts index decd85f..8f2e244 100644 --- a/src/streaming/sse.ts +++ b/src/streaming/sse.ts @@ -15,12 +15,17 @@ */ import { FastifyReply } from 'fastify'; +import { entries } from 'remeda'; import { Event } from './dtos/event.js'; export const init = (res: FastifyReply) => { res.hijack(); if (!res.raw.headersSent) { + const headers = res.getHeaders(); + entries(headers).forEach(([key, value]) => { + if (value) res.raw.setHeader(key, value); + }); res.raw.setHeader('Content-Type', 'text/event-stream'); res.raw.setHeader('Connection', 'keep-alive'); res.raw.setHeader('Cache-Control', 'no-cache,no-transform'); diff --git a/src/utils/rate-limit.ts b/src/utils/rate-limit.ts new file mode 100644 index 0000000..ca6c0fe --- /dev/null +++ b/src/utils/rate-limit.ts @@ -0,0 +1,45 @@ +import { FastifyReply } from 'fastify'; + +import { dayjs, getLatestDailyFixedTime } from './datetime'; + +import { ensureRequestContextData } from '@/context'; + +function getNumericHeader(res: FastifyReply, header: string, fallback: number) { + const value = res.getHeader(header); + if (value === undefined) return fallback; + if (typeof value !== 'number') throw new Error('Invalid header type'); + return value; +} + +const RateLimitHeaders = { + LIMIT: 'ratelimit-limit', + REMAINING: 'ratelimit-remaining', + RESET: 'ratelimit-reset', + RETRY: 'retry-after' +} as const; + +export function updateRateLimitHeadersWithDailyQuota({ + quota, + used +}: { + quota: number; + used: number; +}) { + const res = ensureRequestContextData('res'); + res.header( + RateLimitHeaders.LIMIT, + Math.min(getNumericHeader(res, RateLimitHeaders.LIMIT, Infinity), quota) + ); + res.header( + RateLimitHeaders.REMAINING, + Math.min(getNumericHeader(res, RateLimitHeaders.REMAINING, Infinity), quota - used) + ); + if (quota === used) { + const reset = Math.max( + getNumericHeader(res, RateLimitHeaders.RESET, 0), + getLatestDailyFixedTime().add(1, 'day').unix() - dayjs().unix() + ); + res.header(RateLimitHeaders.RESET, reset); + res.header(RateLimitHeaders.RETRY, reset); + } +} diff --git a/src/vector-store-files/vector-store-files.service.ts b/src/vector-store-files/vector-store-files.service.ts index 121cd99..a7e3625 100644 --- a/src/vector-store-files/vector-store-files.service.ts +++ b/src/vector-store-files/vector-store-files.service.ts @@ -53,6 +53,7 @@ import { QueueName } from '@/jobs/constants.js'; import { getProjectPrincipal } from '@/administration/helpers.js'; import { VECTOR_STORE_FILE_QUOTA_DAILY } from '@/config.js'; import { dayjs, getLatestDailyFixedTime } from '@/utils/datetime.js'; +import { updateRateLimitHeadersWithDailyQuota } from '@/utils/rate-limit.js'; const getFileLogger = (vectorStoreFileIds?: string[]) => getServiceLogger('vector-store-files').child({ vectorStoreFileIds }); @@ -62,6 +63,7 @@ export async function assertVectorStoreFilesQuota(newFilesCount = 1) { createdBy: getProjectPrincipal(), createdAt: { $gte: getLatestDailyFixedTime().toDate() } }); + updateRateLimitHeadersWithDailyQuota({ quota: VECTOR_STORE_FILE_QUOTA_DAILY, used: count }); if (count + newFilesCount > VECTOR_STORE_FILE_QUOTA_DAILY) { throw new APIError({ message: 'Your daily vector store file quota has been exceeded',