From 200e8d2e38f29a12cd6c7ba4129e8d930f4c4648 Mon Sep 17 00:00:00 2001 From: ramnique <30795890+ramnique@users.noreply.github.com> Date: Tue, 4 Feb 2025 16:35:12 +0530 Subject: [PATCH] add rate-limiting --- apps/rowboat/app/actions.ts | 31 +++++- .../app/api/v1/[projectId]/chat/route.ts | 6 ++ .../widget/v1/chats/[chatId]/turn/route.ts | 6 ++ apps/rowboat/app/lib/client_utils.ts | 6 ++ apps/rowboat/app/lib/rate_limiting.ts | 21 ++++ apps/rowboat/app/lib/redis.ts | 7 ++ apps/rowboat/package-lock.json | 102 ++++++++++++++++++ apps/rowboat/package.json | 2 + docker-compose.yml | 9 ++ 9 files changed, 188 insertions(+), 2 deletions(-) create mode 100644 apps/rowboat/app/lib/client_utils.ts create mode 100644 apps/rowboat/app/lib/rate_limiting.ts create mode 100644 apps/rowboat/app/lib/redis.ts diff --git a/apps/rowboat/app/actions.ts b/apps/rowboat/app/actions.ts index eb0c0d9b..6e6f9998 100644 --- a/apps/rowboat/app/actions.ts +++ b/apps/rowboat/app/actions.ts @@ -3,7 +3,7 @@ import { redirect } from "next/navigation"; import { SimulationData, EmbeddingDoc, GetInformationToolResult, DataSource, PlaygroundChat, AgenticAPIChatRequest, AgenticAPIChatResponse, convertFromAgenticAPIChatMessages, WebpageCrawlResponse, Workflow, WorkflowAgent, CopilotAPIRequest, CopilotAPIResponse, CopilotMessage, CopilotWorkflow, convertToCopilotWorkflow, convertToCopilotApiMessage, convertToCopilotMessage, CopilotAssistantMessage, CopilotChatContext, convertToCopilotApiChatContext, Scenario, ClientToolCallRequestBody, ClientToolCallJwt, ClientToolCallRequest, WithStringId, Project, WorkflowTool, WorkflowPrompt, ApiKey } from "./lib/types"; import { ObjectId, WithId } from "mongodb"; -import { generateObject, generateText, tool, embed } from "ai"; +import { generateObject, generateText, embed } from "ai"; import { dataSourcesCollection, embeddingsCollection, projectsCollection, webpagesCollection, agentWorkflowsCollection, scenariosCollection, projectMembersCollection, apiKeysCollection } from "@/app/lib/mongodb"; import { z } from 'zod'; import { openai } from "@ai-sdk/openai"; @@ -12,12 +12,13 @@ import { embeddingModel } from "./lib/embedding"; import { apiV1 } from "rowboat-shared"; import { zodToJsonSchema } from 'zod-to-json-schema'; import crypto from 'crypto'; -import { SignJWT } from "jose"; import { Claims, getSession } from "@auth0/nextjs-auth0"; import { revalidatePath } from "next/cache"; import { callClientToolWebhook, getAgenticApiResponse } from "./lib/utils"; import { templates } from "./lib/project_templates"; import { assert, error } from "node:console"; +import { check_query_limit } from "./lib/rate_limiting"; +import { QueryLimitError } from "./lib/client_utils"; const crawler = new FirecrawlApp({ apiKey: process.env.FIRECRAWL_API_KEY || '' }); @@ -319,6 +320,19 @@ export async function scrapeWebpage(url: string): Promise 0) { + const count = await projectsCollection.countDocuments({ + createdByUserId: user.sub, + }); + if (count >= projectsLimit) { + throw new Error('You have reached your project limit. Please upgrade your plan.'); + } + } + const name = formData.get('name') as string; const templateKey = formData.get('template') as string; const projectId = crypto.randomUUID(); @@ -492,6 +506,9 @@ export async function getAssistantResponse( rawResponse: unknown, }> { await projectAuthCheck(projectId); + if (!await check_query_limit(projectId)) { + throw new QueryLimitError(); + } const response = await getAgenticApiResponse(request); return { @@ -513,6 +530,9 @@ export async function getCopilotResponse( rawResponse: unknown, }> { await projectAuthCheck(projectId); + if (!await check_query_limit(projectId)) { + throw new QueryLimitError(); + } // prepare request const request: z.infer = { @@ -643,6 +663,9 @@ export async function getCopilotResponse( export async function suggestToolResponse(toolId: string, projectId: string, messages: z.infer[]): Promise { await projectAuthCheck(projectId); + if (!await check_query_limit(projectId)) { + throw new QueryLimitError(); + } const prompt = ` # Your Specific Task: @@ -891,6 +914,10 @@ export async function simulateUserResponse( simulationData: z.infer ): Promise { await projectAuthCheck(projectId); + if (!await check_query_limit(projectId)) { + throw new QueryLimitError(); + } + const articlePrompt = ` # Your Specific Task: diff --git a/apps/rowboat/app/api/v1/[projectId]/chat/route.ts b/apps/rowboat/app/api/v1/[projectId]/chat/route.ts index 0c4f391c..b4a3b9d2 100644 --- a/apps/rowboat/app/api/v1/[projectId]/chat/route.ts +++ b/apps/rowboat/app/api/v1/[projectId]/chat/route.ts @@ -5,6 +5,7 @@ import { ObjectId } from "mongodb"; import { authCheck } from "@/app/api/v1/utils"; import { convertFromApiToAgenticApiMessages, convertFromAgenticApiToApiMessages, AgenticAPIChatRequest, ApiRequest, ApiResponse, convertWorkflowToAgenticAPI } from "@/app/lib/types"; import { getAgenticApiResponse } from "@/app/lib/utils"; +import { check_query_limit } from "@/app/lib/rate_limiting"; // get next turn / agent response export async function POST( @@ -13,6 +14,11 @@ export async function POST( ): Promise { const { projectId } = await params; + // check query limit + if (!await check_query_limit(projectId)) { + return Response.json({ error: "Query limit exceeded" }, { status: 429 }); + } + return await authCheck(projectId, req, async () => { // parse and validate the request body let body; diff --git a/apps/rowboat/app/api/widget/v1/chats/[chatId]/turn/route.ts b/apps/rowboat/app/api/widget/v1/chats/[chatId]/turn/route.ts index d9bb823e..ee6f1633 100644 --- a/apps/rowboat/app/api/widget/v1/chats/[chatId]/turn/route.ts +++ b/apps/rowboat/app/api/widget/v1/chats/[chatId]/turn/route.ts @@ -6,6 +6,7 @@ import { ObjectId, WithId } from "mongodb"; import { authCheck } from "../../../utils"; import { AgenticAPIChatRequest, convertFromAgenticAPIChatMessages, convertToAgenticAPIChatMessages, convertWorkflowToAgenticAPI } from "@/app/lib/types"; import { callClientToolWebhook, getAgenticApiResponse } from "@/app/lib/utils"; +import { check_query_limit } from "@/app/lib/rate_limiting"; const chatsCollection = db.collection>("chats"); const chatMessagesCollection = db.collection>("chatMessages"); @@ -18,6 +19,11 @@ export async function POST( return await authCheck(req, async (session) => { const { chatId } = await params; + // check query limit + if (!await check_query_limit(session.projectId)) { + return Response.json({ error: "Query limit exceeded" }, { status: 429 }); + } + // parse and validate the request body let body; try { diff --git a/apps/rowboat/app/lib/client_utils.ts b/apps/rowboat/app/lib/client_utils.ts new file mode 100644 index 00000000..8688d64e --- /dev/null +++ b/apps/rowboat/app/lib/client_utils.ts @@ -0,0 +1,6 @@ +export class QueryLimitError extends Error { + constructor(message: string = 'Query limit exceeded') { + super(message); + this.name = 'QueryLimitError'; + } +} \ No newline at end of file diff --git a/apps/rowboat/app/lib/rate_limiting.ts b/apps/rowboat/app/lib/rate_limiting.ts new file mode 100644 index 00000000..99ea4663 --- /dev/null +++ b/apps/rowboat/app/lib/rate_limiting.ts @@ -0,0 +1,21 @@ +import { redisClient } from "./redis"; + +const MAX_QUERIES_PER_MINUTE = Number(process.env.MAX_QUERIES_PER_MINUTE) || 0; + +export async function check_query_limit(projectId: string): Promise { + // if the limit is 0, we don't want to check the limit + if (MAX_QUERIES_PER_MINUTE === 0) { + return true; + } + + const minutes_since_epoch = Math.floor(Date.now() / 1000 / 60); // 60 second window + const key = `rate_limit:${projectId}:${minutes_since_epoch}`; + + // increment the counter and return the count + const count = await redisClient.incr(key); + if (count === 1) { + await redisClient.expire(key, 70); // Set TTL to clean up automatically + } + + return count <= MAX_QUERIES_PER_MINUTE; +} \ No newline at end of file diff --git a/apps/rowboat/app/lib/redis.ts b/apps/rowboat/app/lib/redis.ts new file mode 100644 index 00000000..9017fe2f --- /dev/null +++ b/apps/rowboat/app/lib/redis.ts @@ -0,0 +1,7 @@ +import { createClient } from 'redis'; + +export const redisClient = createClient({ + url: process.env.REDIS_URL, +}); + +redisClient.connect(); diff --git a/apps/rowboat/package-lock.json b/apps/rowboat/package-lock.json index 47196c2a..b8d9adac 100644 --- a/apps/rowboat/package-lock.json +++ b/apps/rowboat/package-lock.json @@ -34,6 +34,7 @@ "react-dom": "^18.3.1", "react-markdown": "^9.0.1", "react-resizable-panels": "^2.1.7", + "redis": "^4.7.0", "remark-gfm": "^4.0.0", "rowboat-shared": "github:rowboatlabs/shared", "sharp": "^0.33.4", @@ -50,6 +51,7 @@ "@types/node": "^20", "@types/react": "^18", "@types/react-dom": "^18", + "@types/redis": "^4.0.11", "eslint": "^8", "eslint-config-next": "14.2.5", "postcss": "^8", @@ -6440,6 +6442,64 @@ "react": "^16.8.0 || ^17.0.0-rc.1 || ^18.0.0" } }, + "node_modules/@redis/bloom": { + "version": "1.2.0", + "resolved": "https://registry.npmjs.org/@redis/bloom/-/bloom-1.2.0.tgz", + "integrity": "sha512-HG2DFjYKbpNmVXsa0keLHp/3leGJz1mjh09f2RLGGLQZzSHpkmZWuwJbAvo3QcRY8p80m5+ZdXZdYOSBLlp7Cg==", + "peerDependencies": { + "@redis/client": "^1.0.0" + } + }, + "node_modules/@redis/client": { + "version": "1.6.0", + "resolved": "https://registry.npmjs.org/@redis/client/-/client-1.6.0.tgz", + "integrity": "sha512-aR0uffYI700OEEH4gYnitAnv3vzVGXCFvYfdpu/CJKvk4pHfLPEy/JSZyrpQ+15WhXe1yJRXLtfQ84s4mEXnPg==", + "dependencies": { + "cluster-key-slot": "1.1.2", + "generic-pool": "3.9.0", + "yallist": "4.0.0" + }, + "engines": { + "node": ">=14" + } + }, + "node_modules/@redis/client/node_modules/yallist": { + "version": "4.0.0", + "resolved": "https://registry.npmjs.org/yallist/-/yallist-4.0.0.tgz", + "integrity": "sha512-3wdGidZyq5PB084XLES5TpOSRA3wjXAlIWMhum2kRcv/41Sn2emQ0dycQW4uZXLejwKvg6EsvbdlVL+FYEct7A==" + }, + "node_modules/@redis/graph": { + "version": "1.1.1", + "resolved": "https://registry.npmjs.org/@redis/graph/-/graph-1.1.1.tgz", + "integrity": "sha512-FEMTcTHZozZciLRl6GiiIB4zGm5z5F3F6a6FZCyrfxdKOhFlGkiAqlexWMBzCi4DcRoyiOsuLfW+cjlGWyExOw==", + "peerDependencies": { + "@redis/client": "^1.0.0" + } + }, + "node_modules/@redis/json": { + "version": "1.0.7", + "resolved": "https://registry.npmjs.org/@redis/json/-/json-1.0.7.tgz", + "integrity": "sha512-6UyXfjVaTBTJtKNG4/9Z8PSpKE6XgSyEb8iwaqDcy+uKrd/DGYHTWkUdnQDyzm727V7p21WUMhsqz5oy65kPcQ==", + "peerDependencies": { + "@redis/client": "^1.0.0" + } + }, + "node_modules/@redis/search": { + "version": "1.2.0", + "resolved": "https://registry.npmjs.org/@redis/search/-/search-1.2.0.tgz", + "integrity": "sha512-tYoDBbtqOVigEDMAcTGsRlMycIIjwMCgD8eR2t0NANeQmgK/lvxNAvYyb6bZDD4frHRhIHkJu2TBRvB0ERkOmw==", + "peerDependencies": { + "@redis/client": "^1.0.0" + } + }, + "node_modules/@redis/time-series": { + "version": "1.1.0", + "resolved": "https://registry.npmjs.org/@redis/time-series/-/time-series-1.1.0.tgz", + "integrity": "sha512-c1Q99M5ljsIuc4YdaCwfUEXsofakb9c8+Zse2qxTadu8TalLXuAESzLvFAvNVbkmSlvlzIQOLpBCmWI9wTOt+g==", + "peerDependencies": { + "@redis/client": "^1.0.0" + } + }, "node_modules/@rushstack/eslint-patch": { "version": "1.10.3", "resolved": "https://registry.npmjs.org/@rushstack/eslint-patch/-/eslint-patch-1.10.3.tgz", @@ -6715,6 +6775,16 @@ "@types/react": "*" } }, + "node_modules/@types/redis": { + "version": "4.0.11", + "resolved": "https://registry.npmjs.org/@types/redis/-/redis-4.0.11.tgz", + "integrity": "sha512-bI+gth8La8Wg/QCR1+V1fhrL9+LZUSWfcqpOj2Kc80ZQ4ffbdL173vQd5wovmoV9i071FU9oP2g6etLuEwb6Rg==", + "deprecated": "This is a stub types definition. redis provides its own type definitions, so you do not need this installed.", + "dev": true, + "dependencies": { + "redis": "*" + } + }, "node_modules/@types/retry": { "version": "0.12.0", "resolved": "https://registry.npmjs.org/@types/retry/-/retry-0.12.0.tgz", @@ -7832,6 +7902,14 @@ "node": ">=6" } }, + "node_modules/cluster-key-slot": { + "version": "1.1.2", + "resolved": "https://registry.npmjs.org/cluster-key-slot/-/cluster-key-slot-1.1.2.tgz", + "integrity": "sha512-RMr0FhtfXemyinomL4hrWcYJxmX6deFdCxpJzhDttxgO1+bcCnkk+9drydLVDmAMG7NE6aN/fl4F7ucU/90gAA==", + "engines": { + "node": ">=0.10.0" + } + }, "node_modules/code-red": { "version": "1.0.4", "resolved": "https://registry.npmjs.org/code-red/-/code-red-1.0.4.tgz", @@ -9416,6 +9494,14 @@ "resolved": "https://registry.npmjs.org/fzy.js/-/fzy.js-0.4.1.tgz", "integrity": "sha512-4sPVXf+9oGhzg2tYzgWe4hgAY0wEbkqeuKVEgdnqX8S8VcLosQsDjb0jV+f5uoQlf8INWId1w0IGoufAoik1TA==" }, + "node_modules/generic-pool": { + "version": "3.9.0", + "resolved": "https://registry.npmjs.org/generic-pool/-/generic-pool-3.9.0.tgz", + "integrity": "sha512-hymDOu5B53XvN4QT9dBmZxPX4CWhBPPLguTZ9MMFeFa/Kg0xWVfylOVNlJji/E7yTZWFd/q9GO5TxDLq156D7g==", + "engines": { + "node": ">= 4" + } + }, "node_modules/gensync": { "version": "1.0.0-beta.2", "resolved": "https://registry.npmjs.org/gensync/-/gensync-1.0.0-beta.2.tgz", @@ -12929,6 +13015,22 @@ "node": ">=8.10.0" } }, + "node_modules/redis": { + "version": "4.7.0", + "resolved": "https://registry.npmjs.org/redis/-/redis-4.7.0.tgz", + "integrity": "sha512-zvmkHEAdGMn+hMRXuMBtu4Vo5P6rHQjLoHftu+lBqq8ZTA3RCVC/WzD790bkKKiNFp7d5/9PcSD19fJyyRvOdQ==", + "workspaces": [ + "./packages/*" + ], + "dependencies": { + "@redis/bloom": "1.2.0", + "@redis/client": "1.6.0", + "@redis/graph": "1.1.1", + "@redis/json": "1.0.7", + "@redis/search": "1.2.0", + "@redis/time-series": "1.1.0" + } + }, "node_modules/reflect.getprototypeof": { "version": "1.0.6", "resolved": "https://registry.npmjs.org/reflect.getprototypeof/-/reflect.getprototypeof-1.0.6.tgz", diff --git a/apps/rowboat/package.json b/apps/rowboat/package.json index 86981c68..5c34b892 100644 --- a/apps/rowboat/package.json +++ b/apps/rowboat/package.json @@ -37,6 +37,7 @@ "react-dom": "^18.3.1", "react-markdown": "^9.0.1", "react-resizable-panels": "^2.1.7", + "redis": "^4.7.0", "remark-gfm": "^4.0.0", "rowboat-shared": "github:rowboatlabs/shared", "sharp": "^0.33.4", @@ -53,6 +54,7 @@ "@types/node": "^20", "@types/react": "^18", "@types/react-dom": "^18", + "@types/redis": "^4.0.11", "eslint": "^8", "eslint-config-next": "14.2.5", "postcss": "^8", diff --git a/docker-compose.yml b/docker-compose.yml index 76403396..5be5d381 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -23,6 +23,9 @@ services: - AUTH0_ISSUER_BASE_URL=${AUTH0_ISSUER_BASE_URL} - AUTH0_CLIENT_ID=${AUTH0_CLIENT_ID} - AUTH0_CLIENT_SECRET=${AUTH0_CLIENT_SECRET} + - REDIS_URL=redis://redis:6379 + - MAX_QUERIES_PER_MINUTE=${MAX_QUERIES_PER_MINUTE} + - MAX_PROJECTS_PER_USER=${MAX_PROJECTS_PER_USER} restart: unless-stopped agents: @@ -54,3 +57,9 @@ services: ports: - "8000:8000" restart: unless-stopped + + redis: + image: redis:latest + ports: + - "6379:6379" + restart: unless-stopped