From 9fe76ea4eb7d4187ce2a53c76da9e492d6ca5200 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fatih=20Kadir=20Ak=C4=B1n?= Date: Wed, 25 Mar 2026 10:50:07 +0300 Subject: [PATCH] feat: add tiered rate limiting for MCP tool call endpoints --- prompts.config.ts | 2 +- src/lib/rate-limit.ts | 93 +++++++++++++++++++++++++++++++++++++++++++ src/pages/api/mcp.ts | 55 +++++++++++++++++++++++++ 3 files changed, 149 insertions(+), 1 deletion(-) create mode 100644 src/lib/rate-limit.ts diff --git a/prompts.config.ts b/prompts.config.ts index 9da02fe2..84dc3c68 100644 --- a/prompts.config.ts +++ b/prompts.config.ts @@ -61,7 +61,7 @@ export default defineConfig({ // Enable AI-powered generation features (requires OPENAI_API_KEY) aiGeneration: false, // Enable MCP (Model Context Protocol) features including API key generation - mcp: false, + mcp: true, // Enable comments on prompts comments: true, }, diff --git a/src/lib/rate-limit.ts b/src/lib/rate-limit.ts new file mode 100644 index 00000000..7dfe6764 --- /dev/null +++ b/src/lib/rate-limit.ts @@ -0,0 +1,93 @@ +/** + * In-memory sliding window rate limiter. + * + * Each limiter instance tracks request timestamps per identifier (IP or API key) + * and rejects requests that exceed the configured window/max. + * + * NOTE: This is per-process. In a multi-instance deployment, consider a + * Redis-backed implementation instead. + */ + +interface RateLimitEntry { + timestamps: number[]; +} + +interface RateLimiterOptions { + /** Maximum number of requests allowed within the window. */ + max: number; + /** Time window in seconds. */ + windowSeconds: number; +} + +export class RateLimiter { + private store = new Map(); + private readonly max: number; + private readonly windowMs: number; + // Cleanup stale entries every 60 s to prevent memory leaks + private cleanupInterval: ReturnType; + + constructor(opts: RateLimiterOptions) { + this.max = opts.max; + this.windowMs = opts.windowSeconds * 1000; + this.cleanupInterval = setInterval(() => this.cleanup(), 60_000); + // Allow the process to exit without waiting for the interval + if (this.cleanupInterval.unref) { + this.cleanupInterval.unref(); + } + } + + /** + * Check whether the given identifier is allowed to make a request. + * Returns `{ allowed: true, remaining }` or `{ allowed: false, retryAfterSeconds }`. + */ + check(identifier: string): { allowed: true; remaining: number } | { allowed: false; retryAfterSeconds: number } { + const now = Date.now(); + const windowStart = now - this.windowMs; + + let entry = this.store.get(identifier); + if (!entry) { + entry = { timestamps: [] }; + this.store.set(identifier, entry); + } + + // Drop timestamps outside the current window + entry.timestamps = entry.timestamps.filter((t) => t > windowStart); + + if (entry.timestamps.length >= this.max) { + // Earliest timestamp that will leave the window + const oldest = entry.timestamps[0]; + const retryAfterMs = oldest + this.windowMs - now; + return { allowed: false, retryAfterSeconds: Math.ceil(retryAfterMs / 1000) }; + } + + entry.timestamps.push(now); + return { allowed: true, remaining: this.max - entry.timestamps.length }; + } + + private cleanup() { + const now = Date.now(); + const windowStart = now - this.windowMs; + for (const [key, entry] of this.store) { + entry.timestamps = entry.timestamps.filter((t) => t > windowStart); + if (entry.timestamps.length === 0) { + this.store.delete(key); + } + } + } +} + +// --------------------------------------------------------------------------- +// Pre-configured limiters for MCP tool calls +// --------------------------------------------------------------------------- + +/** General MCP POST requests – 20 req / min per identifier */ +export const mcpGeneralLimiter = new RateLimiter({ max: 20, windowSeconds: 60 }); + +/** tool calls (tools/call) – 10 req / min per identifier */ +export const mcpToolCallLimiter = new RateLimiter({ max: 10, windowSeconds: 60 }); + +/** Write-mutation tools – 5 req / min per identifier */ +export const mcpWriteToolLimiter = new RateLimiter({ max: 5, windowSeconds: 60 }); + +/** AI-powered tools (improve_prompt) – 2 req / min per identifier */ +export const mcpAiToolLimiter = new RateLimiter({ max: 2, windowSeconds: 60 }); diff --git a/src/pages/api/mcp.ts b/src/pages/api/mcp.ts index 77e0d3b7..ce9f9a1b 100644 --- a/src/pages/api/mcp.ts +++ b/src/pages/api/mcp.ts @@ -12,6 +12,12 @@ import { db } from "@/lib/db"; import { isValidApiKeyFormat } from "@/lib/api-key"; import { parseSkillFiles, serializeSkillFiles, DEFAULT_SKILL_FILE } from "@/lib/skill-files"; import appConfig from "@/../prompts.config"; +import { + mcpGeneralLimiter, + mcpToolCallLimiter, + mcpWriteToolLimiter, + mcpAiToolLimiter, +} from "@/lib/rate-limit"; interface AuthenticatedUser { id: string; @@ -1463,6 +1469,55 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse) const body = await parseBody(req); + // --- Rate limiting --- + const rateLimitId = apiKey || req.headers["x-forwarded-for"]?.toString().split(",")[0]?.trim() || req.socket.remoteAddress || "unknown"; + + const generalCheck = mcpGeneralLimiter.check(rateLimitId); + if (!generalCheck.allowed) { + return res.status(429).json({ + jsonrpc: "2.0", + error: { code: -32000, message: `Rate limit exceeded. Try again in ${generalCheck.retryAfterSeconds}s.` }, + id: null, + }); + } + + // Apply stricter limits for tool calls based on tool name + const WRITE_TOOLS = new Set(["save_prompt", "save_skill", "add_file_to_skill", "update_skill_file", "remove_file_from_skill"]); + const AI_TOOLS = new Set(["improve_prompt"]); + + const rpcBody = body as { method?: string; params?: { name?: string } }; + if (rpcBody?.method === "tools/call") { + const toolCallCheck = mcpToolCallLimiter.check(rateLimitId); + if (!toolCallCheck.allowed) { + return res.status(429).json({ + jsonrpc: "2.0", + error: { code: -32000, message: `Tool call rate limit exceeded. Try again in ${toolCallCheck.retryAfterSeconds}s.` }, + id: null, + }); + } + + const toolName = rpcBody.params?.name; + if (toolName && AI_TOOLS.has(toolName)) { + const aiCheck = mcpAiToolLimiter.check(rateLimitId); + if (!aiCheck.allowed) { + return res.status(429).json({ + jsonrpc: "2.0", + error: { code: -32000, message: `AI tool rate limit exceeded (${toolName}). Try again in ${aiCheck.retryAfterSeconds}s.` }, + id: null, + }); + } + } else if (toolName && WRITE_TOOLS.has(toolName)) { + const writeCheck = mcpWriteToolLimiter.check(rateLimitId); + if (!writeCheck.allowed) { + return res.status(429).json({ + jsonrpc: "2.0", + error: { code: -32000, message: `Write tool rate limit exceeded (${toolName}). Try again in ${writeCheck.retryAfterSeconds}s.` }, + id: null, + }); + } + } + } + await transport.handleRequest(req, res, body); res.on("close", () => {