feat(prediction): opt-in default rate limit keyed by user or IP

packages/server/src/controllers/predictions/index.ts

@@ -104,7 +104,7 @@ const createPrediction = async (req: Request, res: Response, next: NextFunction)
 
 const getRateLimiterMiddleware = async (req: Request, res: Response, next: NextFunction) => {
     try {
-        return RateLimiterManager.getInstance().getRateLimiter()(req, res, next)
+        return RateLimiterManager.getInstance().getRateLimiter({ withDefault: true })(req, res, next)
     } catch (error) {
         next(error)
     }

packages/server/src/routes/predictions/index.ts

@@ -8,6 +8,7 @@ const router = express.Router()
 
 // /prediction is whitelisted, so the global auth never reads the token. Identify Auth0 callers
 // (browser extension) here; anonymous callers pass through. The data source exists only at request time.
+// It runs before the rate limiter so the default limit can count per user, not per office IP.
 const optionalAAIAuth = (req: Request, res: Response, next: NextFunction) =>
     verifyAAIToken(getRunningExpressApp().AppDataSource, { optional: true })(req, res, next)
 
@@ -17,8 +18,8 @@ const optionalAAIAuth = (req: Request, res: Response, next: NextFunction) =>
 router.post(
     ['/', '/:id'],
     getMulterStorage().array('files'),
-    predictionsController.getRateLimiterMiddleware,
     optionalAAIAuth,
+    predictionsController.getRateLimiterMiddleware,
     predictionsController.createPrediction
 )
 

packages/server/src/utils/rateLimit.ts

@@ -1,5 +1,5 @@
 import { NextFunction, Request, Response } from 'express'
-import { rateLimit, RateLimitRequestHandler } from 'express-rate-limit'
+import { Options, rateLimit, RateLimitRequestHandler } from 'express-rate-limit'
 import { IChatFlow, MODE } from '../Interface'
 import { Mutex } from 'async-mutex'
 import { RedisStore } from 'rate-limit-redis'
@@ -12,9 +12,17 @@ interface CustomListener extends QueueEventsListener {
 
 const QUEUE_NAME = 'ratelimit'
 const QUEUE_EVENT_NAME = 'updateRateLimiter'
+const DEFAULT_LIMIT_MESSAGE = 'Too many requests, please try again later.'
+
+// Default limiter key: per chatflow, per identified user (so people behind one office IP don't share a bucket), else per IP.
+export const rateLimitKey = (req: Request): string => {
+    const id = req.params.id ?? ''
+    return req.user?.id ? `${id}:user:${req.user.id}` : `${id}:ip:${req.ip}`
+}
 
 export class RateLimiterManager {
     private rateLimiters: Record<string, RateLimitRequestHandler> = {}
+    private defaultLimiter?: RateLimitRequestHandler
     private rateLimiterMutex: Mutex = new Mutex()
     private redisClient: Redis
     private static instance: RateLimiterManager
@@ -53,6 +61,36 @@ export class RateLimiterManager {
             this.queueEventsProducer = new QueueEventsProducer(QUEUE_NAME, { connection: this.getConnection() })
             this.queueEvents = new QueueEvents(QUEUE_NAME, { connection: this.getConnection() })
         }
+
+        // Opt-in: off unless DEFAULT_RATE_LIMIT_MAX is set, so a release changes nothing for deployments that don't set it.
+        const max = Number(process.env.DEFAULT_RATE_LIMIT_MAX)
+        if (Number.isInteger(max) && max > 0) {
+            const windowSeconds = Number(process.env.DEFAULT_RATE_LIMIT_WINDOW_SECONDS) || 60
+            this.defaultLimiter = this.createLimiter('rl:default:', windowSeconds * 1000, max, DEFAULT_LIMIT_MESSAGE, {
+                keyGenerator: rateLimitKey,
+                standardHeaders: true,
+                legacyHeaders: false
+            })
+        }
+    }
+
+    private createLimiter(prefix: string, windowMs: number, max: number, message: string, extra: Partial<Options> = {}) {
+        if (process.env.MODE === MODE.QUEUE) {
+            return rateLimit({
+                windowMs,
+                max,
+                standardHeaders: true,
+                legacyHeaders: false,
+                message,
+                store: new RedisStore({
+                    prefix,
+                    // @ts-expect-error - Known issue: the `call` function is not present in @types/ioredis
+                    sendCommand: (...args: string[]) => this.redisClient.call(...args)
+                }),
+                ...extra
+            })
+        }
+        return rateLimit({ windowMs, max, message, ...extra })
     }
 
     getConnection() {
@@ -94,26 +132,7 @@ export class RateLimiterManager {
     public async addRateLimiter(id: string, duration: number, limit: number, message: string): Promise<void> {
         const release = await this.rateLimiterMutex.acquire()
         try {
-            if (process.env.MODE === MODE.QUEUE) {
-                this.rateLimiters[id] = rateLimit({
-                    windowMs: duration * 1000,
-                    max: limit,
-                    standardHeaders: true,
-                    legacyHeaders: false,
-                    message,
-                    store: new RedisStore({
-                        prefix: `rl:${id}`,
-                        // @ts-expect-error - Known issue: the `call` function is not present in @types/ioredis
-                        sendCommand: (...args: string[]) => this.redisClient.call(...args)
-                    })
-                })
-            } else {
-                this.rateLimiters[id] = rateLimit({
-                    windowMs: duration * 1000,
-                    max: limit,
-                    message
-                })
-            }
+            this.rateLimiters[id] = this.createLimiter(`rl:${id}`, duration * 1000, limit, message)
         } finally {
             release()
         }
@@ -125,12 +144,11 @@ export class RateLimiterManager {
         }
     }
 
-    public getRateLimiter(): (req: Request, res: Response, next: NextFunction) => void {
+    public getRateLimiter(opts: { withDefault?: boolean } = {}): (req: Request, res: Response, next: NextFunction) => void {
         return (req: Request, res: Response, next: NextFunction) => {
-            const id = req.params.id
-            if (!this.rateLimiters[id]) return next()
-            const idRateLimiter = this.rateLimiters[id]
-            return idRateLimiter(req, res, next)
+            const limiter = this.rateLimiters[req.params.id] ?? (opts.withDefault ? this.defaultLimiter : undefined)
+            if (!limiter) return next()
+            return limiter(req, res, next)
         }
     }
 

packages/server/test/integration/auth0Fixture.ts

@@ -0,0 +1,63 @@
+/**
+ * Shared setup for local integration tests: a local Auth0-like issuer (RS256 + JWKS) and a Postgres
+ * DataSource migrated with the app's own migrations. Call setupAuth0AndDb before importing anything
+ * that reads AUTH0_* at module load (jwtCheck).
+ */
+import crypto from 'node:crypto'
+import http from 'node:http'
+import { AddressInfo } from 'node:net'
+import { DataSource } from 'typeorm'
+
+export const AUDIENCE = 'https://api.test.local'
+export const ORG = 'org_ias'
+const b64 = (o: object) => Buffer.from(JSON.stringify(o)).toString('base64url')
+
+export const setupAuth0AndDb = async (pgUrl: string) => {
+    const keys = crypto.generateKeyPairSync('rsa', { modulusLength: 2048 })
+    const otherKeys = crypto.generateKeyPairSync('rsa', { modulusLength: 2048 })
+    let issuer = ''
+    const jwk = { ...keys.publicKey.export({ format: 'jwk' }), kid: 'k1', alg: 'RS256', use: 'sig' }
+    const server = http.createServer((req, res) => {
+        res.setHeader('Content-Type', 'application/json')
+        if (req.url === '/.well-known/openid-configuration') return res.end(JSON.stringify({ issuer, jwks_uri: `${issuer}jwks` }))
+        if (req.url === '/jwks') return res.end(JSON.stringify({ keys: [jwk] }))
+        res.statusCode = 404
+        res.end('{}')
+    })
+    await new Promise<void>((resolve) => server.listen(0, '127.0.0.1', resolve))
+    issuer = `http://127.0.0.1:${(server.address() as AddressInfo).port}/`
+
+    const sign = (claims: Record<string, any>, key = keys.privateKey) => {
+        const now = Math.floor(Date.now() / 1000)
+        const head = b64({ alg: 'RS256', typ: 'JWT', kid: 'k1' })
+        const body = b64({ iss: issuer, aud: AUDIENCE, iat: now, exp: now + 600, ...claims })
+        const sig = crypto.sign('sha256', Buffer.from(`${head}.${body}`), key).toString('base64url')
+        return `${head}.${body}.${sig}`
+    }
+
+    const url = new URL(pgUrl)
+    Object.assign(process.env, {
+        AUTH0_ISSUER_BASE_URL: issuer,
+        AUTH0_AUDIENCE: AUDIENCE,
+        AUTH0_ORGANIZATION_ID: ORG,
+        DATABASE_TYPE: 'postgres',
+        DATABASE_HOST: url.hostname,
+        DATABASE_PORT: url.port,
+        DATABASE_USER: decodeURIComponent(url.username),
+        DATABASE_PASSWORD: decodeURIComponent(url.password),
+        DATABASE_NAME: url.pathname.slice(1)
+    })
+
+    const dsModule = await import('../../src/DataSource')
+    await dsModule.init()
+    const ds: DataSource = dsModule.getDataSource()
+    await ds.initialize()
+    await ds.query('DROP SCHEMA public CASCADE; CREATE SCHEMA public; CREATE EXTENSION IF NOT EXISTS "uuid-ossp"')
+    await ds.runMigrations({ transaction: 'each' })
+
+    const close = async () => {
+        await ds.destroy()
+        await new Promise((resolve) => server.close(resolve))
+    }
+    return { ds, sign, otherKey: otherKeys.privateKey, close }
+}

packages/server/test/integration/prediction/index.test.ts

@@ -12,11 +12,10 @@
  * NODE_OPTIONS is needed because langfuse-core (loaded through flowise-components) runs a dynamic import() at load.
  */
 import crypto from 'node:crypto'
-import http from 'node:http'
-import { AddressInfo } from 'node:net'
 import express from 'express'
 import supertest from 'supertest'
 import { DataSource } from 'typeorm'
+import { ORG, setupAuth0AndDb } from '../auth0Fixture'
 
 const mockApp: { AppDataSource?: DataSource } = {}
 jest.mock('../../../src/utils/getRunningExpressApp', () => ({
@@ -31,30 +30,17 @@ jest.mock('../../../src/services/predictions', () => ({
 const PG = process.env.TEST_PG_URL
 const d = PG ? describe : describe.skip
 
-const AUDIENCE = 'https://api.test.local'
-const ORG = 'org_ias'
-const b64 = (o: object) => Buffer.from(JSON.stringify(o)).toString('base64url')
-
 d('POST /api/v1/prediction (integration)', () => {
     jest.setTimeout(180000)
 
-    const keys = crypto.generateKeyPairSync('rsa', { modulusLength: 2048 })
-    const otherKeys = crypto.generateKeyPairSync('rsa', { modulusLength: 2048 })
-    let issuer = ''
-    let issuerServer: http.Server
     let ds: DataSource
     let app: express.Express
     let entities: { User: any; Organization: any; Workspace: any; WorkspaceUser: any; ChatFlow: any }
     let userId = ''
     let chatflowId = ''
+    let fx: Awaited<ReturnType<typeof setupAuth0AndDb>>
+    const sign = (claims: Record<string, any>, key?: crypto.KeyObject) => fx.sign(claims, key)
 
-    const sign = (claims: Record<string, any>, key = keys.privateKey) => {
-        const now = Math.floor(Date.now() / 1000)
-        const head = b64({ alg: 'RS256', typ: 'JWT', kid: 'k1' })
-        const body = b64({ iss: issuer, aud: AUDIENCE, iat: now, exp: now + 600, ...claims })
-        const sig = crypto.sign('sha256', Buffer.from(`${head}.${body}`), key).toString('base64url')
-        return `${head}.${body}.${sig}`
-    }
     const alice = (extra: Record<string, any> = {}) =>
         sign({ sub: 'auth0|alice', email: 'alice@test.local', name: 'Alice', org_id: ORG, ...extra })
 
@@ -73,37 +59,9 @@ d('POST /api/v1/prediction (integration)', () => {
     })
 
     beforeAll(async () => {
-        const jwk = { ...keys.publicKey.export({ format: 'jwk' }), kid: 'k1', alg: 'RS256', use: 'sig' }
-        issuerServer = http.createServer((req, res) => {
-            res.setHeader('Content-Type', 'application/json')
-            if (req.url === '/.well-known/openid-configuration') return res.end(JSON.stringify({ issuer, jwks_uri: `${issuer}jwks` }))
-            if (req.url === '/jwks') return res.end(JSON.stringify({ keys: [jwk] }))
-            res.statusCode = 404
-            res.end('{}')
-        })
-        await new Promise<void>((resolve) => issuerServer.listen(0, '127.0.0.1', resolve))
-        issuer = `http://127.0.0.1:${(issuerServer.address() as AddressInfo).port}/`
-
-        const url = new URL(PG!)
-        Object.assign(process.env, {
-            AUTH0_ISSUER_BASE_URL: issuer,
-            AUTH0_AUDIENCE: AUDIENCE,
-            AUTH0_ORGANIZATION_ID: ORG,
-            DATABASE_TYPE: 'postgres',
-            DATABASE_HOST: url.hostname,
-            DATABASE_PORT: url.port,
-            DATABASE_USER: decodeURIComponent(url.username),
-            DATABASE_PASSWORD: decodeURIComponent(url.password),
-            DATABASE_NAME: url.pathname.slice(1)
-        })
-
-        // Imported after env is set: jwtCheck reads AUTH0_* at module load.
-        const dsModule = await import('../../../src/DataSource')
-        await dsModule.init()
-        ds = dsModule.getDataSource()
-        await ds.initialize()
-        await ds.query('DROP SCHEMA public CASCADE; CREATE SCHEMA public; CREATE EXTENSION IF NOT EXISTS "uuid-ossp"')
-        await ds.runMigrations({ transaction: 'each' })
+        // Before any import that reads AUTH0_* at module load (jwtCheck).
+        fx = await setupAuth0AndDb(PG!)
+        ds = fx.ds
         mockApp.AppDataSource = ds
 
         entities = {
@@ -150,8 +108,7 @@ d('POST /api/v1/prediction (integration)', () => {
     })
 
     afterAll(async () => {
-        await ds?.destroy()
-        await new Promise((resolve) => issuerServer?.close(resolve))
+        await fx?.close()
     })
 
     beforeEach(() => {
@@ -192,7 +149,7 @@ d('POST /api/v1/prediction (integration)', () => {
     })
 
     it('forged token stays anonymous', async () => {
-        const forged = sign({ sub: 'auth0|alice', email: 'alice@test.local', org_id: ORG }, otherKeys.privateKey)
+        const forged = sign({ sub: 'auth0|alice', email: 'alice@test.local', org_id: ORG }, fx.otherKey)
         await post(chatflowId, `Bearer ${forged}`).then(ok)
         expect(seenUser).toBeUndefined()
     })