fix(limits): merge window and rate limits, fix streaming headers (#134)

Signed-off-by: Tomas Pilar <[email protected]>

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) {

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',

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');

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);
+  }
+}

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',