Просмотр исходного кода

add rate limits (#46)

* add rate limits

* fix lint
Chris Tate 5 месяцев назад
Родитель
Сommit
a48f6d665d

+ 1 - 0
.gitignore

@@ -39,3 +39,4 @@ yarn-error.log*
 
 # opensrc - source code for packages
 opensrc/
+.env*.local

+ 5 - 0
apps/web/.env.example

@@ -7,3 +7,8 @@ AI_GATEWAY_API_KEY=
 # Override the default model used for UI generation
 # Default: anthropic/claude-haiku-4.5
 AI_GATEWAY_MODEL=anthropic/claude-haiku-4.5
+
+# Vercel KV (Rate Limiting)
+# Automatically populated when you add Vercel KV to your project
+KV_REST_API_URL=
+KV_REST_API_TOKEN=

+ 1 - 0
apps/web/.gitignore

@@ -35,3 +35,4 @@ yarn-error.log*
 # typescript
 *.tsbuildinfo
 next-env.d.ts
+.env*.local

+ 28 - 0
apps/web/app/api/generate/route.ts

@@ -1,4 +1,6 @@
 import { streamText } from "ai";
+import { headers } from "next/headers";
+import { minuteRateLimit, dailyRateLimit } from "@/lib/rate-limit";
 
 export const maxDuration = 30;
 
@@ -81,6 +83,32 @@ const MAX_PROMPT_LENGTH = 500;
 const DEFAULT_MODEL = "anthropic/claude-haiku-4.5";
 
 export async function POST(req: Request) {
+  // Get client IP for rate limiting
+  const headersList = await headers();
+  const ip = headersList.get("x-forwarded-for")?.split(",")[0] ?? "anonymous";
+
+  // Check rate limits (minute and daily)
+  const [minuteResult, dailyResult] = await Promise.all([
+    minuteRateLimit.limit(ip),
+    dailyRateLimit.limit(ip),
+  ]);
+
+  if (!minuteResult.success || !dailyResult.success) {
+    const isMinuteLimit = !minuteResult.success;
+    return new Response(
+      JSON.stringify({
+        error: "Rate limit exceeded",
+        message: isMinuteLimit
+          ? "Too many requests. Please wait a moment before trying again."
+          : "Daily limit reached. Please try again tomorrow.",
+      }),
+      {
+        status: 429,
+        headers: { "Content-Type": "application/json" },
+      },
+    );
+  }
+
   const { prompt, context } = await req.json();
   const previousTree = context?.previousTree;
 

+ 4 - 1
apps/web/components/demo.tsx

@@ -219,7 +219,10 @@ export function Demo({
     clear,
   } = useUIStream({
     api: "/api/generate",
-    onError: (err: Error) => console.error("Generation error:", err),
+    onError: (err: Error) => {
+      console.error("Generation error:", err);
+      toast.error(err.message || "Generation failed. Please try again.");
+    },
   } as Parameters<typeof useUIStream>[0]);
 
   // Initialize interactive state for Select components

+ 18 - 2
apps/web/components/playground.tsx

@@ -28,7 +28,7 @@ interface Version {
   id: string;
   prompt: string;
   tree: UITree | null;
-  status: "generating" | "complete";
+  status: "generating" | "complete" | "error";
 }
 
 const EXAMPLE_PROMPTS = [
@@ -64,7 +64,20 @@ export function Playground() {
     clear,
   } = useUIStream({
     api: "/api/generate",
-    onError: (err: Error) => console.error("Generation error:", err),
+    onError: (err: Error) => {
+      console.error("Generation error:", err);
+      toast.error(err.message || "Generation failed. Please try again.");
+      // Mark the version as errored
+      if (generatingVersionIdRef.current) {
+        const erroredVersionId = generatingVersionIdRef.current;
+        setVersions((prev) =>
+          prev.map((v) =>
+            v.id === erroredVersionId ? { ...v, status: "error" as const } : v,
+          ),
+        );
+        generatingVersionIdRef.current = null;
+      }
+    },
   } as Parameters<typeof useUIStream>[0]);
 
   useInteractiveState();
@@ -307,6 +320,9 @@ ${jsx}
                     ...
                   </span>
                 )}
+                {version.status === "error" && (
+                  <span className="text-xs text-red-500 shrink-0">failed</span>
+                )}
               </div>
             </button>
           ))

+ 21 - 0
apps/web/lib/rate-limit.ts

@@ -0,0 +1,21 @@
+import { Ratelimit } from "@upstash/ratelimit";
+import { Redis } from "@upstash/redis";
+
+const redis = new Redis({
+  url: process.env.KV_REST_API_URL!,
+  token: process.env.KV_REST_API_TOKEN!,
+});
+
+// 10 requests per minute (sliding window)
+export const minuteRateLimit = new Ratelimit({
+  redis,
+  limiter: Ratelimit.slidingWindow(10, "1 m"),
+  prefix: "ratelimit:minute",
+});
+
+// 100 requests per day (fixed window)
+export const dailyRateLimit = new Ratelimit({
+  redis,
+  limiter: Ratelimit.fixedWindow(100, "1 d"),
+  prefix: "ratelimit:daily",
+});

+ 2 - 0
apps/web/package.json

@@ -19,6 +19,8 @@
     "@radix-ui/react-dialog": "^1.1.15",
     "@radix-ui/react-slot": "^1.2.4",
     "@radix-ui/react-tabs": "^1.1.13",
+    "@upstash/ratelimit": "^2.0.8",
+    "@upstash/redis": "^1.36.1",
     "@vercel/analytics": "^1.6.1",
     "@vercel/speed-insights": "^1.3.1",
     "ai": "^6.0.33",

+ 13 - 1
packages/react/src/hooks.ts

@@ -153,7 +153,19 @@ export function useUIStream({
         });
 
         if (!response.ok) {
-          throw new Error(`HTTP error: ${response.status}`);
+          // Try to parse JSON error response for better error messages
+          let errorMessage = `HTTP error: ${response.status}`;
+          try {
+            const errorData = await response.json();
+            if (errorData.message) {
+              errorMessage = errorData.message;
+            } else if (errorData.error) {
+              errorMessage = errorData.error;
+            }
+          } catch {
+            // Ignore JSON parsing errors, use default message
+          }
+          throw new Error(errorMessage);
         }
 
         const reader = response.body?.getReader();

+ 36 - 0
pnpm-lock.yaml

@@ -59,6 +59,12 @@ importers:
       '@radix-ui/react-tabs':
         specifier: ^1.1.13
         version: 1.1.13(@types/react-dom@19.2.3(@types/react@19.2.3))(@types/react@19.2.3)(react-dom@19.2.3(react@19.2.3))(react@19.2.3)
+      '@upstash/ratelimit':
+        specifier: ^2.0.8
+        version: 2.0.8(@upstash/redis@1.36.1)
+      '@upstash/redis':
+        specifier: ^1.36.1
+        version: 1.36.1
       '@vercel/analytics':
         specifier: ^1.6.1
         version: 1.6.1(next@16.1.1(@opentelemetry/api@1.9.0)(react-dom@19.2.3(react@19.2.3))(react@19.2.3))(react@19.2.3)
@@ -1421,6 +1427,18 @@ packages:
   '@ungap/structured-clone@1.3.0':
     resolution: {integrity: sha512-WmoN8qaIAo7WTYWbAZuG8PYEhn5fkz7dZrqTBZ7dtt//lL2Gwms1IcnQ5yHqjDfX8Ft5j4YzDM23f87zBfDe9g==}
 
+  '@upstash/core-analytics@0.0.10':
+    resolution: {integrity: sha512-7qJHGxpQgQr9/vmeS1PktEwvNAF7TI4iJDi8Pu2CFZ9YUGHZH4fOP5TfYlZ4aVxfopnELiE4BS4FBjyK7V1/xQ==}
+    engines: {node: '>=16.0.0'}
+
+  '@upstash/ratelimit@2.0.8':
+    resolution: {integrity: sha512-YSTMBJ1YIxsoPkUMX/P4DDks/xV5YYCswWMamU8ZIfK9ly6ppjRnVOyBhMDXBmzjODm4UQKcxsJPvaeFAijp5w==}
+    peerDependencies:
+      '@upstash/redis': ^1.34.3
+
+  '@upstash/redis@1.36.1':
+    resolution: {integrity: sha512-N6SjDcgXdOcTAF+7uNoY69o7hCspe9BcA7YjQdxVu5d25avljTwyLaHBW3krWjrP0FfocgMk94qyVtQbeDp39A==}
+
   '@vercel/analytics@1.6.1':
     resolution: {integrity: sha512-oH9He/bEM+6oKlv3chWuOOcp8Y6fo6/PSro8hEkgCW3pu9/OiCXiUpRUogDh3Fs3LH2sosDrx8CxeOLBEE+afg==}
     peerDependencies:
@@ -3127,6 +3145,9 @@ packages:
     resolution: {integrity: sha512-nWJ91DjeOkej/TA8pXQ3myruKpKEYgqvpw9lz4OPHj/NWFNluYrjbz9j01CJ8yKQd2g4jFoOkINCTW2I5LEEyw==}
     engines: {node: '>= 0.4'}
 
+  uncrypto@0.1.3:
+    resolution: {integrity: sha512-Ql87qFHB3s/De2ClA9e0gsnS6zXG27SkTiSJwjCc9MebbfapQfuPzumMIUMi38ezPZVNFcHI9sUIepeQfw8J8Q==}
+
   undici-types@6.21.0:
     resolution: {integrity: sha512-iwDZqg0QAGrg9Rav5H4n0M64c3mkR59cJ6wQp+7C4nI0gsmExaedaYLNO44eT4AtBBwjbTiGPMlt2Md0T9H9JQ==}
 
@@ -4241,6 +4262,19 @@ snapshots:
 
   '@ungap/structured-clone@1.3.0': {}
 
+  '@upstash/core-analytics@0.0.10':
+    dependencies:
+      '@upstash/redis': 1.36.1
+
+  '@upstash/ratelimit@2.0.8(@upstash/redis@1.36.1)':
+    dependencies:
+      '@upstash/core-analytics': 0.0.10
+      '@upstash/redis': 1.36.1
+
+  '@upstash/redis@1.36.1':
+    dependencies:
+      uncrypto: 0.1.3
+
   '@vercel/analytics@1.6.1(next@16.1.1(@opentelemetry/api@1.9.0)(react-dom@19.2.3(react@19.2.3))(react@19.2.3))(react@19.2.3)':
     optionalDependencies:
       next: 16.1.1(@opentelemetry/api@1.9.0)(react-dom@19.2.3(react@19.2.3))(react@19.2.3)
@@ -6154,6 +6188,8 @@ snapshots:
       has-symbols: 1.1.0
       which-boxed-primitive: 1.1.1
 
+  uncrypto@0.1.3: {}
+
   undici-types@6.21.0: {}
 
   unist-util-is@6.0.1:

+ 1 - 1
turbo.json

@@ -1,7 +1,7 @@
 {
   "$schema": "https://turborepo.dev/schema.json",
   "ui": "tui",
-  "globalEnv": ["AI_GATEWAY_MODEL"],
+  "globalEnv": ["AI_GATEWAY_MODEL", "KV_REST_API_URL", "KV_REST_API_TOKEN"],
   "tasks": {
     "build": {
       "dependsOn": ["^build"],