feat: add tiered rate limiting for MCP tool call endpoints

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

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<string, RateLimitEntry>();
+  private readonly max: number;
+  private readonly windowMs: number;
+  // Cleanup stale entries every 60 s to prevent memory leaks
+  private cleanupInterval: ReturnType<typeof setInterval>;
+
+  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 });

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", () => {