diff --git a/workspaces/lightspeed/.changeset/five-pots-change.md b/workspaces/lightspeed/.changeset/five-pots-change.md new file mode 100644 index 0000000000..434fcc59b8 --- /dev/null +++ b/workspaces/lightspeed/.changeset/five-pots-change.md @@ -0,0 +1,5 @@ +--- +'@red-hat-developer-hub/backstage-plugin-lightspeed-backend': minor +--- + +add rate limiting to lightspeed and notebooks diff --git a/workspaces/lightspeed/plugins/lightspeed-backend/README.md b/workspaces/lightspeed/plugins/lightspeed-backend/README.md index 9f473ab457..cf86b4095a 100644 --- a/workspaces/lightspeed/plugins/lightspeed-backend/README.md +++ b/workspaces/lightspeed/plugins/lightspeed-backend/README.md @@ -36,6 +36,43 @@ lightspeed: mcpServers: # Optional - one or more MCP servers - name: # must match the name configured in LCS token: ${MCP_TOKEN} + rateLimit: # Optional - per-user request rate limits (defaults apply if omitted) + expensive: + max: 25 # Max requests per minute per user for expensive endpoints (default: 25). Set to 0 to disable. + general: + max: 200 # Max requests per minute per user for other authenticated endpoints (default: 200). Set to 0 to disable. +``` + +#### Rate limiting + +The backend applies per-user rate limits to authenticated endpoints as an abuse +prevention measure. Limits are keyed by the authenticated user's entity ref and +use a fixed 1-minute window. + +**Tiers**: + +- **Expensive** (default: 25 requests/minute per user): `POST /v1/query`, and + (when Notebooks is enabled) notebook document uploads and RAG queries. +- **General** (default: 200 requests/minute per user): all other authenticated + endpoints, including conversation listing, MCP server management, feedback, + and notebook session CRUD. +- **Excluded**: `/health` and `/notebooks/health` are not rate limited. + +When a limit is exceeded, the API returns `429 Too Many Requests` with a +`Retry-After` header and a JSON error body (`RateLimitExceeded`). + +Set `max: 0` on a tier to disable rate limiting for that tier. If the entire +`rateLimit` block is omitted, the defaults above apply. + +**Example** — tighter limits for a small deployment: + +```yaml +lightspeed: + rateLimit: + expensive: + max: 10 + general: + max: 100 ``` #### MCP servers settings endpoints diff --git a/workspaces/lightspeed/plugins/lightspeed-backend/config.d.ts b/workspaces/lightspeed/plugins/lightspeed-backend/config.d.ts index 709741b0be..2908be6475 100644 --- a/workspaces/lightspeed/plugins/lightspeed-backend/config.d.ts +++ b/workspaces/lightspeed/plugins/lightspeed-backend/config.d.ts @@ -49,6 +49,38 @@ export interface Config { */ token?: string; }>; + /** + * Per-user rate limiting for Lightspeed API endpoints. + * @visibility backend + */ + rateLimit?: { + /** + * Limits for expensive endpoints (LLM queries, document uploads). + * @visibility backend + */ + expensive?: { + /** + * Maximum requests per minute per user. + * Set to 0 to disable rate limiting for this tier. + * @default 25 + * @visibility backend + */ + max?: number; + }; + /** + * Limits for all other authenticated endpoints. + * @visibility backend + */ + general?: { + /** + * Maximum requests per minute per user. + * Set to 0 to disable rate limiting for this tier. + * @default 200 + * @visibility backend + */ + max?: number; + }; + }; /** * Configuration for AI Notebooks (Developer Preview) */ diff --git a/workspaces/lightspeed/plugins/lightspeed-backend/package.json b/workspaces/lightspeed/plugins/lightspeed-backend/package.json index 2df7504fd2..88a0e2dae3 100644 --- a/workspaces/lightspeed/plugins/lightspeed-backend/package.json +++ b/workspaces/lightspeed/plugins/lightspeed-backend/package.json @@ -52,6 +52,7 @@ "@langchain/openai": "^0.6.0", "@red-hat-developer-hub/backstage-plugin-lightspeed-common": "workspace:^", "express": "^4.21.1", + "express-rate-limit": "^8.2.2", "form-data": "^4.0.5", "htmlparser2": "^9.1.0", "http-proxy-middleware": "^3.0.2", diff --git a/workspaces/lightspeed/plugins/lightspeed-backend/src/service/constant.ts b/workspaces/lightspeed/plugins/lightspeed-backend/src/service/constant.ts index b16b108d92..a78673e282 100644 --- a/workspaces/lightspeed/plugins/lightspeed-backend/src/service/constant.ts +++ b/workspaces/lightspeed/plugins/lightspeed-backend/src/service/constant.ts @@ -26,6 +26,13 @@ export const DEFAULT_LIGHTSPEED_SERVICE_HOST = '127.0.0.1'; // Lightspeed core s export const DEFAULT_LIGHTSPEED_SERVICE_PORT = 8080; // Lightspeed service port export const DEFAULT_MAX_FILE_SIZE_MB = 20 * 1024 * 1024; // 20MB +/** + * Rate limiting defaults (window is fixed at 1 minute) + */ +export const RATE_LIMIT_WINDOW_MS = 60000; +export const DEFAULT_EXPENSIVE_RATE_LIMIT_MAX = 25; +export const DEFAULT_GENERAL_RATE_LIMIT_MAX = 200; + /** * Input validation limits for query endpoints */ diff --git a/workspaces/lightspeed/plugins/lightspeed-backend/src/service/express.d.ts b/workspaces/lightspeed/plugins/lightspeed-backend/src/service/express.d.ts index f512a5d031..f7bf4b4bb3 100644 --- a/workspaces/lightspeed/plugins/lightspeed-backend/src/service/express.d.ts +++ b/workspaces/lightspeed/plugins/lightspeed-backend/src/service/express.d.ts @@ -16,10 +16,13 @@ import type { BackstageCredentials } from '@backstage/backend-plugin-api'; +import type { RateLimitInfo } from 'express-rate-limit'; + // Populated by the identity middleware for use in route handlers. declare module 'express-serve-static-core' { interface Request { credentials?: BackstageCredentials; userEntityRef?: string; + rateLimit?: RateLimitInfo; } } diff --git a/workspaces/lightspeed/plugins/lightspeed-backend/src/service/middleware/createRateLimitMiddleware.test.ts b/workspaces/lightspeed/plugins/lightspeed-backend/src/service/middleware/createRateLimitMiddleware.test.ts new file mode 100644 index 0000000000..e289fa3a2e --- /dev/null +++ b/workspaces/lightspeed/plugins/lightspeed-backend/src/service/middleware/createRateLimitMiddleware.test.ts @@ -0,0 +1,195 @@ +/* + * Copyright Red Hat, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { mockServices } from '@backstage/backend-test-utils'; + +import express from 'express'; +import request from 'supertest'; + +import { + DEFAULT_EXPENSIVE_RATE_LIMIT_MAX, + DEFAULT_GENERAL_RATE_LIMIT_MAX, +} from '../constant'; +import { + createRateLimitMiddleware, + getRateLimitMax, +} from './createRateLimitMiddleware'; + +describe('getRateLimitMax', () => { + it('returns default expensive max when config is omitted', () => { + const config = mockServices.rootConfig({ data: {} }); + expect(getRateLimitMax(config, 'expensive')).toBe( + DEFAULT_EXPENSIVE_RATE_LIMIT_MAX, + ); + }); + + it('returns default general max when config is omitted', () => { + const config = mockServices.rootConfig({ data: {} }); + expect(getRateLimitMax(config, 'general')).toBe( + DEFAULT_GENERAL_RATE_LIMIT_MAX, + ); + }); + + it('returns configured max values when provided', () => { + const config = mockServices.rootConfig({ + data: { + lightspeed: { + rateLimit: { + expensive: { max: 10 }, + general: { max: 50 }, + }, + }, + }, + }); + + expect(getRateLimitMax(config, 'expensive')).toBe(10); + expect(getRateLimitMax(config, 'general')).toBe(50); + }); + + it('treats negative values as disabled (0)', () => { + const config = mockServices.rootConfig({ + data: { + lightspeed: { + rateLimit: { + expensive: { max: -1 }, + general: { max: -5 }, + }, + }, + }, + }); + + expect(getRateLimitMax(config, 'expensive')).toBe(0); + expect(getRateLimitMax(config, 'general')).toBe(0); + }); + + it('floors decimal values to integers', () => { + const config = mockServices.rootConfig({ + data: { + lightspeed: { + rateLimit: { + expensive: { max: 10.7 }, + general: { max: 50.3 }, + }, + }, + }, + }); + + expect(getRateLimitMax(config, 'expensive')).toBe(10); + expect(getRateLimitMax(config, 'general')).toBe(50); + }); +}); + +describe('createRateLimitMiddleware', () => { + function createTestApp( + max: number, + tier: 'expensive' | 'general' = 'general', + ) { + const app = express(); + const config = mockServices.rootConfig({ + data: { + lightspeed: { + rateLimit: { + [tier]: { max }, + }, + }, + }, + }); + + app.use((req, _res, next) => { + req.credentials = { $$type: '@backstage/BackstageCredentials' } as any; + req.userEntityRef = 'user:default/test-user'; + next(); + }); + app.get('/test', createRateLimitMiddleware(config, tier), (_req, res) => { + res.json({ ok: true }); + }); + + return app; + } + + it('allows requests up to the configured max', async () => { + const app = createTestApp(1); + + const first = await request(app).get('/test'); + expect(first.status).toBe(200); + expect(first.body).toEqual({ ok: true }); + }); + + it('returns 429 with Retry-After when limit is exceeded', async () => { + const app = createTestApp(1); + + await request(app).get('/test'); + const second = await request(app).get('/test'); + + expect(second.status).toBe(429); + expect(second.headers['retry-after']).toBeDefined(); + expect(second.body).toEqual({ + error: { + name: 'RateLimitExceeded', + message: 'Too many requests. Please try again later.', + retryAfter: expect.any(Number), + }, + }); + }); + + it('does not rate limit when max is 0', async () => { + const app = createTestApp(0); + + const first = await request(app).get('/test'); + const second = await request(app).get('/test'); + + expect(first.status).toBe(200); + expect(second.status).toBe(200); + }); + + it('tracks limits independently per user', async () => { + const config = mockServices.rootConfig({ + data: { + lightspeed: { + rateLimit: { + general: { max: 1 }, + }, + }, + }, + }); + const app = express(); + + app.use((req, _res, next) => { + req.credentials = { $$type: '@backstage/BackstageCredentials' } as any; + req.userEntityRef = + req.headers['x-test-user']?.toString() ?? 'user:default/user-a'; + next(); + }); + app.get( + '/test', + createRateLimitMiddleware(config, 'general'), + (_req, res) => { + res.json({ ok: true }); + }, + ); + + await request(app).get('/test').set('x-test-user', 'user:default/user-a'); + const blockedForA = await request(app) + .get('/test') + .set('x-test-user', 'user:default/user-a'); + const allowedForB = await request(app) + .get('/test') + .set('x-test-user', 'user:default/user-b'); + + expect(blockedForA.status).toBe(429); + expect(allowedForB.status).toBe(200); + }); +}); diff --git a/workspaces/lightspeed/plugins/lightspeed-backend/src/service/middleware/createRateLimitMiddleware.ts b/workspaces/lightspeed/plugins/lightspeed-backend/src/service/middleware/createRateLimitMiddleware.ts new file mode 100644 index 0000000000..ceec32d974 --- /dev/null +++ b/workspaces/lightspeed/plugins/lightspeed-backend/src/service/middleware/createRateLimitMiddleware.ts @@ -0,0 +1,85 @@ +/* + * Copyright Red Hat, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { Config } from '@backstage/config'; + +import type { RequestHandler } from 'express'; +import rateLimit from 'express-rate-limit'; + +import { + DEFAULT_EXPENSIVE_RATE_LIMIT_MAX, + DEFAULT_GENERAL_RATE_LIMIT_MAX, + RATE_LIMIT_WINDOW_MS, +} from '../constant'; +import { getIdentity } from './getIdentity'; + +export type RateLimitTier = 'expensive' | 'general'; + +export function getRateLimitMax(config: Config, tier: RateLimitTier): number { + const configKey = + tier === 'expensive' + ? 'lightspeed.rateLimit.expensive.max' + : 'lightspeed.rateLimit.general.max'; + const defaultMax = + tier === 'expensive' + ? DEFAULT_EXPENSIVE_RATE_LIMIT_MAX + : DEFAULT_GENERAL_RATE_LIMIT_MAX; + + const configured = config.getOptionalNumber(configKey); + if (configured === undefined) { + return defaultMax; + } + + if (configured <= 0) { + return 0; // rate limit of 0 means disabled + } + + return Math.floor(configured); +} + +export function createRateLimitMiddleware( + config: Config, + tier: RateLimitTier, +): RequestHandler { + const max = getRateLimitMax(config, tier); + + if (max === 0) { + return (_req, _res, next) => next(); + } + + return rateLimit({ + windowMs: RATE_LIMIT_WINDOW_MS, + max, + standardHeaders: false, + legacyHeaders: false, + keyGenerator: req => getIdentity(req).userEntityRef, + handler: (req, res, _next, options) => { + const resetTime = req.rateLimit?.resetTime; + const retryAfter = resetTime + ? Math.max(1, Math.ceil((resetTime.getTime() - Date.now()) / 1000)) + : Math.ceil(RATE_LIMIT_WINDOW_MS / 1000); + + res.setHeader('Retry-After', String(retryAfter)); + res.status(options.statusCode).json({ + error: { + name: 'RateLimitExceeded', + message: 'Too many requests. Please try again later.', + retryAfter, + }, + }); + }, + }); +} diff --git a/workspaces/lightspeed/plugins/lightspeed-backend/src/service/notebooks/notebooksRouter.test.ts b/workspaces/lightspeed/plugins/lightspeed-backend/src/service/notebooks/notebooksRouter.test.ts index 04bf897166..c4542701cd 100644 --- a/workspaces/lightspeed/plugins/lightspeed-backend/src/service/notebooks/notebooksRouter.test.ts +++ b/workspaces/lightspeed/plugins/lightspeed-backend/src/service/notebooks/notebooksRouter.test.ts @@ -469,4 +469,96 @@ describe('Notebooks Router', () => { expect(credentialsSpy).toHaveBeenCalledTimes(1); }); }); + + describe('rate limiting', () => { + let rateLimitedApp: express.Application; + + beforeEach(async () => { + resetMockStorage(); + VectorStoresOperator.resetInstance(); + const logger = mockServices.logger.mock(); + const config = mockServices.rootConfig({ + data: { + lightspeed: { + servicePort: 7007, + rateLimit: { + expensive: { max: 1 }, + general: { max: 1 }, + }, + notebooks: { + enabled: true, + queryDefaults: { + model: 'test-model', + provider_id: 'test-provider', + }, + sessionDefaults: { + provider_id: 'test-notebooks', + embedding_model: 'test-embedding-model', + embedding_dimension: 768, + }, + }, + }, + }, + }); + + const rateLimitedHttpAuth = mockServices.httpAuth(); + const userInfo = mockServices.userInfo.mock({ + getUserInfo: async () => ({ + userEntityRef: mockUserId, + ownershipEntityRefs: [mockUserId], + }), + }); + const permissions = mockServices.permissions.mock({ + authorize: async () => [{ result: AuthorizeResult.ALLOW }], + }); + + const router = await createNotebooksRouter({ + logger, + config, + httpAuth: rateLimitedHttpAuth, + userInfo, + permissions, + }); + + rateLimitedApp = express(); + rateLimitedApp.use(router); + }); + + it('returns 429 on expensive query endpoint when limit exceeded', async () => { + const sessionRes = await request(rateLimitedApp) + .post('/notebooks/v1/sessions') + .send({ name: 'Rate Limit Test' }); + const sessionId = sessionRes.body.session.session_id; + + const first = await request(rateLimitedApp) + .post(`/notebooks/v1/sessions/${sessionId}/query`) + .send({ query: 'What is this about?' }); + const second = await request(rateLimitedApp) + .post(`/notebooks/v1/sessions/${sessionId}/query`) + .send({ query: 'Another question?' }); + + expect(first.status).not.toBe(429); + expect(second.status).toBe(429); + expect(second.headers['retry-after']).toBeDefined(); + expect(second.body.error.name).toBe('RateLimitExceeded'); + }); + + it('returns 429 on general endpoint when limit exceeded', async () => { + const first = await request(rateLimitedApp).get('/notebooks/v1/sessions'); + const second = await request(rateLimitedApp).get( + '/notebooks/v1/sessions', + ); + + expect(first.status).toBe(200); + expect(second.status).toBe(429); + }); + + it('does not rate limit health endpoint', async () => { + const first = await request(rateLimitedApp).get('/notebooks/health'); + const second = await request(rateLimitedApp).get('/notebooks/health'); + + expect(first.status).toBe(200); + expect(second.status).toBe(200); + }); + }); }); diff --git a/workspaces/lightspeed/plugins/lightspeed-backend/src/service/notebooks/notebooksRouters.ts b/workspaces/lightspeed/plugins/lightspeed-backend/src/service/notebooks/notebooksRouters.ts index a9771b5115..2e50e0171e 100644 --- a/workspaces/lightspeed/plugins/lightspeed-backend/src/service/notebooks/notebooksRouters.ts +++ b/workspaces/lightspeed/plugins/lightspeed-backend/src/service/notebooks/notebooksRouters.ts @@ -39,6 +39,7 @@ import { NOTEBOOKS_SYSTEM_PROMPT, upload, } from '../constant'; +import { createRateLimitMiddleware } from '../middleware/createRateLimitMiddleware'; import { createIdentityMiddleware, getIdentity, @@ -100,6 +101,9 @@ export async function createNotebooksRouter( const authorizer = userPermissionAuthorization(permissions); + const expensiveRateLimiter = createRateLimitMiddleware(config, 'expensive'); + const generalRateLimiter = createRateLimitMiddleware(config, 'general'); + const requireNotebooksPermission = async ( req: any, res: any, @@ -266,10 +270,11 @@ export async function createNotebooksRouter( '/v1', createIdentityMiddleware(httpAuth, userInfo, logger), ); - notebooksRouter.use('/v1', requireNotebooksPermission); notebooksRouter.post( '/v1/sessions', + generalRateLimiter, + requireNotebooksPermission, withAuth(async (req, res, userId) => { const { name, description, metadata } = req.body; if (!name) { @@ -288,6 +293,8 @@ export async function createNotebooksRouter( notebooksRouter.get( '/v1/sessions', + generalRateLimiter, + requireNotebooksPermission, withAuth(async (_req, res, userId) => { const sessions = await sessionService.listSessions(userId); res.json(createSessionListResponse(sessions)); @@ -296,6 +303,8 @@ export async function createNotebooksRouter( notebooksRouter.get( '/v1/sessions/:sessionId', + generalRateLimiter, + requireNotebooksPermission, withAuth(async (req, res, userId) => { const { sessionId } = req.params; const session = await sessionService.readSession(sessionId, userId); @@ -307,6 +316,8 @@ export async function createNotebooksRouter( notebooksRouter.put( '/v1/sessions/:sessionId', + generalRateLimiter, + requireNotebooksPermission, withAuth(async (req, res, userId) => { const { sessionId } = req.params; const { name, description, metadata } = req.body; @@ -323,6 +334,8 @@ export async function createNotebooksRouter( notebooksRouter.delete( '/v1/sessions/:sessionId', + generalRateLimiter, + requireNotebooksPermission, withAuth(async (req, res, userId) => { const { sessionId } = req.params; await sessionService.deleteSession(sessionId, userId); @@ -337,6 +350,8 @@ export async function createNotebooksRouter( notebooksRouter.get( '/v1/sessions/:sessionId/documents', + generalRateLimiter, + requireNotebooksPermission, requireSessionOwnership(), withAuth(async (req, res) => { const { sessionId } = req.params; @@ -351,6 +366,8 @@ export async function createNotebooksRouter( notebooksRouter.put( '/v1/sessions/:sessionId/documents', + expensiveRateLimiter, + requireNotebooksPermission, upload.single('file') as any, withAuth(async (req, res, userId) => { const { sessionId } = req.params; @@ -401,6 +418,8 @@ export async function createNotebooksRouter( notebooksRouter.get( '/v1/sessions/:sessionId/documents/:documentId/status', + generalRateLimiter, + requireNotebooksPermission, requireSessionOwnership(), withAuth(async (req, res) => { const { sessionId, documentId } = req.params; @@ -419,6 +438,8 @@ export async function createNotebooksRouter( notebooksRouter.delete( '/v1/sessions/:sessionId/documents/:documentId', + generalRateLimiter, + requireNotebooksPermission, requireSessionOwnership(), withAuth(async (req, res) => { const { sessionId, documentId } = req.params; @@ -443,6 +464,8 @@ export async function createNotebooksRouter( notebooksRouter.post( '/v1/sessions/:sessionId/query', + expensiveRateLimiter, + requireNotebooksPermission, express.json({ limit: EXPRESS_JSON_BODY_LIMIT }), withAuth(async (req, res, userId) => { const { sessionId } = req.params; diff --git a/workspaces/lightspeed/plugins/lightspeed-backend/src/service/router.test.ts b/workspaces/lightspeed/plugins/lightspeed-backend/src/service/router.test.ts index 5f4744abd0..9084d9e39f 100644 --- a/workspaces/lightspeed/plugins/lightspeed-backend/src/service/router.test.ts +++ b/workspaces/lightspeed/plugins/lightspeed-backend/src/service/router.test.ts @@ -1140,4 +1140,102 @@ describe('lightspeed router tests', () => { expect(response.statusCode).toEqual(403); }); }); + + describe('rate limiting', () => { + const RATE_LIMIT_CONFIG = { + lightspeed: { + ...BASE_CONFIG.lightspeed, + rateLimit: { + expensive: { max: 1 }, + general: { max: 1 }, + }, + }, + }; + + it('returns 429 on expensive endpoint when limit exceeded', async () => { + const backendServer = await startBackendServer(RATE_LIMIT_CONFIG); + const body = { + model: mockModel, + query: 'Hello', + provider: 'test-server', + }; + + const first = await request(backendServer) + .post('/api/lightspeed/v1/query') + .send(body); + const second = await request(backendServer) + .post('/api/lightspeed/v1/query') + .send(body); + + expect(first.statusCode).toBe(200); + expect(second.statusCode).toBe(429); + expect(second.headers['retry-after']).toBeDefined(); + expect(second.body.error).toEqual({ + name: 'RateLimitExceeded', + message: 'Too many requests. Please try again later.', + retryAfter: expect.any(Number), + }); + }); + + it('returns 429 on repeated invalid /v1/query requests after limit exceeded', async () => { + const backendServer = await startBackendServer(RATE_LIMIT_CONFIG); + const invalidBody = { + model: mockModel, + query: 'Hello', + }; + + const first = await request(backendServer) + .post('/api/lightspeed/v1/query') + .send(invalidBody); + const second = await request(backendServer) + .post('/api/lightspeed/v1/query') + .send(invalidBody); + + expect(first.statusCode).toBe(400); + expect(first.body.error).toBe( + 'provider is required and must be a non-empty string', + ); + expect(second.statusCode).toBe(429); + }); + + it('returns 429 on general endpoint when limit exceeded', async () => { + const backendServer = await startBackendServer(RATE_LIMIT_CONFIG); + + const first = await request(backendServer).get( + '/api/lightspeed/v1/models', + ); + const second = await request(backendServer).get( + '/api/lightspeed/v1/models', + ); + + expect(first.statusCode).toBe(200); + expect(second.statusCode).toBe(429); + }); + + it('does not rate limit health endpoint', async () => { + const backendServer = await startBackendServer(RATE_LIMIT_CONFIG); + + const first = await request(backendServer).get('/api/lightspeed/health'); + const second = await request(backendServer).get('/api/lightspeed/health'); + + expect(first.statusCode).toBe(200); + expect(second.statusCode).toBe(200); + }); + + it('uses separate counters for expensive and general tiers', async () => { + const backendServer = await startBackendServer(RATE_LIMIT_CONFIG); + const body = { + model: mockModel, + query: 'Hello', + provider: 'test-server', + }; + + await request(backendServer).post('/api/lightspeed/v1/query').send(body); + const modelsResponse = await request(backendServer).get( + '/api/lightspeed/v1/models', + ); + + expect(modelsResponse.statusCode).toBe(200); + }); + }); }); diff --git a/workspaces/lightspeed/plugins/lightspeed-backend/src/service/router.ts b/workspaces/lightspeed/plugins/lightspeed-backend/src/service/router.ts index 79a4aa787d..b081fadddd 100644 --- a/workspaces/lightspeed/plugins/lightspeed-backend/src/service/router.ts +++ b/workspaces/lightspeed/plugins/lightspeed-backend/src/service/router.ts @@ -46,6 +46,7 @@ import { } from './mcp-server-types'; import { McpServerValidator } from './mcp-server-validator'; import { createPermissionMiddleware } from './middleware/createPermissionMiddleware'; +import { createRateLimitMiddleware } from './middleware/createRateLimitMiddleware'; import { createIdentityMiddleware, getIdentity, @@ -220,12 +221,16 @@ export async function createRouter( return createPermissionMiddleware(authorizer, permission, logger); } + const expensiveRateLimiter = createRateLimitMiddleware(config, 'expensive'); + const generalRateLimiter = createRateLimitMiddleware(config, 'general'); + // ─── MCP Server Management Endpoints ──────────────────────────────── // All MCP servers are admin-configured (static). Users can view the // list, toggle servers on/off, and provide personal access tokens. router.get( '/mcp-servers', + generalRateLimiter, requirePermission(lightspeedMcpReadPermission), async (req, res) => { try { @@ -262,6 +267,7 @@ export async function createRouter( router.post( '/mcp-servers/validate', + generalRateLimiter, requirePermission(lightspeedMcpReadPermission), async (req, res) => { try { @@ -296,6 +302,7 @@ export async function createRouter( router.post( '/mcp-servers/:name/validate', + generalRateLimiter, requirePermission(lightspeedMcpManagePermission), async (req, res) => { try { @@ -361,6 +368,7 @@ export async function createRouter( router.patch( '/mcp-servers/:name', + generalRateLimiter, requirePermission(lightspeedMcpManagePermission), async (req, res) => { try { @@ -442,75 +450,88 @@ export async function createRouter( ); // Returns conversation IDs associated with notebook sessions for filtering - router.get('/notebook-conversation-ids', async (req, res) => { - try { - const { userEntityRef } = getIdentity(req); + router.get( + '/notebook-conversation-ids', + generalRateLimiter, + async (req, res) => { + try { + const { userEntityRef } = getIdentity(req); - const vectorStoresPage = await vectorStoresOperator.vectorStores.list(); - const vectorStores = vectorStoresPage.data || []; + const vectorStoresPage = await vectorStoresOperator.vectorStores.list(); + const vectorStores = vectorStoresPage.data || []; - const conversationIds: string[] = []; + const conversationIds: string[] = []; - for (const store of vectorStores) { - const sessionUserId = store.metadata?.user_id as string; - const conversationId = store.metadata?.conversation_id as string | null; + for (const store of vectorStores) { + const sessionUserId = store.metadata?.user_id as string; + const conversationId = store.metadata?.conversation_id as + | string + | null; - // Only include this user's sessions with a conversation_id - if (sessionUserId === userEntityRef && conversationId) { - conversationIds.push(conversationId); + // Only include this user's sessions with a conversation_id + if (sessionUserId === userEntityRef && conversationId) { + conversationIds.push(conversationId); + } } - } - res.json({ - conversation_ids: conversationIds, - }); - } catch (error) { - const errormsg = `Error fetching notebook conversation IDs: ${error}`; - logger.error(errormsg); + res.json({ + conversation_ids: conversationIds, + }); + } catch (error) { + const errormsg = `Error fetching notebook conversation IDs: ${error}`; + logger.error(errormsg); - if (error instanceof NotAllowedError) { - res.status(403).json({ error: error.message }); - } else { - res.status(500).json({ error: errormsg }); + if (error instanceof NotAllowedError) { + res.status(403).json({ error: error.message }); + } else { + res.status(500).json({ error: errormsg }); + } } - } - }); + }, + ); // ─── Proxy Routes ─────────────────────────────────────────────────── router.get( '/v1/models', + generalRateLimiter, requirePermission(lightspeedChatReadPermission), apiProxy, ); router.get( '/v1/shields', + generalRateLimiter, requirePermission(lightspeedChatReadPermission), apiProxy, ); router.get( '/v2/conversations', + generalRateLimiter, requirePermission(lightspeedChatReadPermission), apiProxy, ); router.get( '/v2/conversations/:conversation_id', + generalRateLimiter, requirePermission(lightspeedChatReadPermission), apiProxy, ); router.delete( '/v2/conversations/:conversation_id', + generalRateLimiter, requirePermission(lightspeedChatDeletePermission), apiProxy, ); router.get( '/v1/feedback/status', + generalRateLimiter, requirePermission(lightspeedChatReadPermission), apiProxy, ); router.post( '/v1/feedback', + generalRateLimiter, requirePermission(lightspeedChatCreatePermission), async (request, response) => { try { @@ -553,6 +574,7 @@ export async function createRouter( router.post( '/v1/query/interrupt', + generalRateLimiter, requirePermission(lightspeedChatCreatePermission), async (request, response) => { try { @@ -590,6 +612,7 @@ export async function createRouter( router.post( '/v1/query', express.json({ limit: EXPRESS_JSON_BODY_LIMIT }), + expensiveRateLimiter, validateCompletionsRequest, requirePermission(lightspeedChatCreatePermission), async (request, response) => { @@ -690,6 +713,7 @@ export async function createRouter( router.put( '/v2/conversations/:conversation_id', + generalRateLimiter, requirePermission(lightspeedChatCreatePermission), async (request, response) => { try { diff --git a/workspaces/lightspeed/yarn.lock b/workspaces/lightspeed/yarn.lock index 13e74a9b47..d48f4b0fbd 100644 --- a/workspaces/lightspeed/yarn.lock +++ b/workspaces/lightspeed/yarn.lock @@ -11195,6 +11195,7 @@ __metadata: "@types/multer": "npm:^1.4.12" "@types/supertest": "npm:2.0.16" express: "npm:^4.21.1" + express-rate-limit: "npm:^8.2.2" form-data: "npm:^4.0.5" htmlparser2: "npm:^9.1.0" http-proxy-middleware: "npm:^3.0.2"