add rate limits

app/api/ai/generate/route.ts

@@ -1,6 +1,7 @@
 import { streamText } from "ai";
 import { NextResponse } from "next/server";
 import { auth } from "@/lib/auth";
+import { checkAIRateLimit, getRateLimitHeaders } from "@/lib/rate-limit";
 import { generateAIActionPrompts } from "@/plugins";
 
 // Simple type for operations
@@ -256,6 +257,15 @@ export async function POST(request: Request) {
       return NextResponse.json({ error: "Unauthorized" }, { status: 401 });
     }
 
+    // Check rate limit
+    const rateLimit = await checkAIRateLimit(session.user.id);
+    if (!rateLimit.allowed) {
+      return NextResponse.json(
+        { error: "Rate limit exceeded. Please try again later." },
+        { status: 429, headers: getRateLimitHeaders(rateLimit) }
+      );
+    }
+
     const body = await request.json();
     const { prompt, existingWorkflow } = body;
 

app/api/workflows/[workflowId]/webhook/route.ts

@@ -5,6 +5,7 @@ import { start } from "workflow/api";
 import { db } from "@/lib/db";
 import { validateWorkflowIntegrations } from "@/lib/db/integrations";
 import { apiKeys, workflowExecutions, workflows } from "@/lib/db/schema";
+import { checkExecutionRateLimit, getRateLimitHeaders } from "@/lib/rate-limit";
 import { executeWorkflow } from "@/lib/workflow-executor.workflow";
 import type { WorkflowEdge, WorkflowNode } from "@/lib/workflow-store";
 
@@ -149,6 +150,18 @@ export async function POST(
       );
     }
 
+    // Check rate limit
+    const rateLimit = await checkExecutionRateLimit(workflow.userId);
+    if (!rateLimit.allowed) {
+      return NextResponse.json(
+        { error: "Rate limit exceeded. Please try again later." },
+        {
+          status: 429,
+          headers: { ...corsHeaders, ...getRateLimitHeaders(rateLimit) },
+        }
+      );
+    }
+
     // Verify this is a webhook-triggered workflow
     const triggerNode = (workflow.nodes as WorkflowNode[]).find(
       (node) => node.data.type === "trigger"

lib/rate-limit.ts

@@ -0,0 +1,68 @@
+import { Ratelimit } from "@upstash/ratelimit";
+import { Redis } from "@upstash/redis";
+
+type RateLimitResult = {
+  allowed: boolean;
+  remaining: number;
+  resetAt: Date;
+};
+
+// Create Redis client with KV_ prefix
+const redis = new Redis({
+  url: process.env.KV_REST_API_URL ?? "",
+  token: process.env.KV_REST_API_TOKEN ?? "",
+});
+
+// AI generation: 50 requests per hour
+const aiRatelimit = new Ratelimit({
+  redis,
+  limiter: Ratelimit.slidingWindow(50, "1 h"),
+  prefix: "ratelimit:ai",
+});
+
+// Webhook execution: 1000 requests per hour
+const webhookRatelimit = new Ratelimit({
+  redis,
+  limiter: Ratelimit.slidingWindow(1000, "1 h"),
+  prefix: "ratelimit:webhook",
+});
+
+/**
+ * Check rate limit for AI generation requests
+ */
+export async function checkAIRateLimit(
+  userId: string
+): Promise<RateLimitResult> {
+  const { success, remaining, reset } = await aiRatelimit.limit(userId);
+
+  return {
+    allowed: success,
+    remaining,
+    resetAt: new Date(reset),
+  };
+}
+
+/**
+ * Check rate limit for workflow executions
+ */
+export async function checkExecutionRateLimit(
+  userId: string
+): Promise<RateLimitResult> {
+  const { success, remaining, reset } = await webhookRatelimit.limit(userId);
+
+  return {
+    allowed: success,
+    remaining,
+    resetAt: new Date(reset),
+  };
+}
+
+/**
+ * Get rate limit headers for response
+ */
+export function getRateLimitHeaders(result: RateLimitResult): HeadersInit {
+  return {
+    "X-RateLimit-Remaining": result.remaining.toString(),
+    "X-RateLimit-Reset": result.resetAt.toISOString(),
+  };
+}

package.json

@@ -26,6 +26,8 @@
     "@radix-ui/react-context-menu": "^2.2.16",
     "@radix-ui/react-dialog": "^1.1.15",
     "@slack/web-api": "^7.12.0",
+    "@upstash/ratelimit": "^2.0.7",
+    "@upstash/redis": "^1.35.7",
     "@vercel/analytics": "^1.5.0",
     "@vercel/og": "^0.8.5",
     "@vercel/sdk": "^1.17.1",