From c0a525c9e40381740866e4aeae3c97b8e226213f Mon Sep 17 00:00:00 2001 From: Chris Tate Date: Mon, 9 Mar 2026 15:43:25 -0500 Subject: [PATCH] rate limits for demo (#695) --- examples/demo/.env.example | 9 +++++ examples/demo/app/actions/browse.ts | 42 +++++++++++++++++++++ examples/demo/app/api/browse/route.ts | 34 ++++++++++++----- examples/demo/app/page.tsx | 22 ++++++----- examples/demo/lib/constants.ts | 8 ++++ examples/demo/lib/rate-limit.ts | 53 +++++++++++++++++++++++++++ examples/demo/package.json | 2 + examples/demo/pnpm-lock.yaml | 36 ++++++++++++++++++ 8 files changed, 188 insertions(+), 18 deletions(-) create mode 100644 examples/demo/lib/constants.ts create mode 100644 examples/demo/lib/rate-limit.ts diff --git a/examples/demo/.env.example b/examples/demo/.env.example index 5127dd3..b413934 100644 --- a/examples/demo/.env.example +++ b/examples/demo/.env.example @@ -7,3 +7,12 @@ # Snapshot ID with agent-browser + Chromium pre-installed (optional, speeds up startup) # Create one with: npx tsx scripts/create-snapshot.ts # AGENT_BROWSER_SNAPSHOT_ID=snap_xxxxxxxxxxxx + +# --- Rate Limiting (Upstash / Vercel KV) --- +# Automatically populated when you add Vercel KV to your project +KV_REST_API_URL= +KV_REST_API_TOKEN= + +# Optional: override default limits +# RATE_LIMIT_PER_MINUTE=10 +# RATE_LIMIT_PER_DAY=100 diff --git a/examples/demo/app/actions/browse.ts b/examples/demo/app/actions/browse.ts index 6008c49..8837b5a 100644 --- a/examples/demo/app/actions/browse.ts +++ b/examples/demo/app/actions/browse.ts @@ -1,7 +1,31 @@ "use server"; +import { headers } from "next/headers"; import * as serverless from "@/lib/agent-browser"; import * as sandbox from "@/lib/agent-browser-sandbox"; +import { ALLOWED_URLS } from "@/lib/constants"; +import { minuteRateLimit, dailyRateLimit } from "@/lib/rate-limit"; + +async function checkRateLimit(): Promise { + const h = await headers(); + const ip = h.get("x-forwarded-for")?.split(",")[0] ?? "anonymous"; + + const minute = await minuteRateLimit.limit(ip); + if (!minute.success) { + return "Too many requests. Please wait a moment before trying again."; + } + + const daily = await dailyRateLimit.limit(ip); + if (!daily.success) { + return "Daily limit reached. Please try again tomorrow."; + } + + return null; +} + +function isAllowedUrl(url: string): boolean { + return (ALLOWED_URLS as readonly string[]).includes(url); +} export type EnvStatus = { serverless: { @@ -52,6 +76,15 @@ export async function takeScreenshot( url: string, mode: Mode = "serverless", ): Promise { + if (!isAllowedUrl(url)) { + return { ok: false, error: "URL not allowed" }; + } + + const rateLimitError = await checkRateLimit(); + if (rateLimitError) { + return { ok: false, error: rateLimitError }; + } + try { if (mode === "sandbox") { const { screenshot, title } = await sandbox.screenshotUrl(url); @@ -78,6 +111,15 @@ export async function takeSnapshot( url: string, mode: Mode = "serverless", ): Promise { + if (!isAllowedUrl(url)) { + return { ok: false, error: "URL not allowed" }; + } + + const rateLimitError = await checkRateLimit(); + if (rateLimitError) { + return { ok: false, error: rateLimitError }; + } + try { if (mode === "sandbox") { const { snapshot, title } = await sandbox.snapshotUrl(url, { diff --git a/examples/demo/app/api/browse/route.ts b/examples/demo/app/api/browse/route.ts index 583ae7a..c87b507 100644 --- a/examples/demo/app/api/browse/route.ts +++ b/examples/demo/app/api/browse/route.ts @@ -1,17 +1,29 @@ import { NextRequest, NextResponse } from "next/server"; import * as ab from "@/lib/agent-browser"; +import { ALLOWED_URLS } from "@/lib/constants"; +import { minuteRateLimit, dailyRateLimit } from "@/lib/rate-limit"; -/** - * POST /api/browse - * - * Programmatic API route for browser automation. - * Uses @sparticuz/chromium + puppeteer-core directly in the function. - * - * Body: { "action": "screenshot", "url": "https://example.com" } - * Or: { "action": "snapshot", "url": "https://example.com" } - */ export async function POST(req: NextRequest) { try { + const ip = + req.headers.get("x-forwarded-for")?.split(",")[0] ?? "anonymous"; + + const minute = await minuteRateLimit.limit(ip); + if (!minute.success) { + return NextResponse.json( + { error: "Too many requests. Please wait a moment before trying again." }, + { status: 429 }, + ); + } + + const daily = await dailyRateLimit.limit(ip); + if (!daily.success) { + return NextResponse.json( + { error: "Daily limit reached. Please try again tomorrow." }, + { status: 429 }, + ); + } + const body = await req.json(); const url = body.url; @@ -19,6 +31,10 @@ export async function POST(req: NextRequest) { return NextResponse.json({ error: "Provide a 'url'" }, { status: 400 }); } + if (!(ALLOWED_URLS as readonly string[]).includes(url)) { + return NextResponse.json({ error: "URL not allowed" }, { status: 400 }); + } + if (body.action === "screenshot") { const result = await ab.screenshotUrl(url, { fullPage: body.fullPage, diff --git a/examples/demo/app/page.tsx b/examples/demo/app/page.tsx index 3d7af7b..39c2d7a 100644 --- a/examples/demo/app/page.tsx +++ b/examples/demo/app/page.tsx @@ -13,8 +13,8 @@ import { ResizablePanel, ResizableHandle, } from "@/components/ui/resizable"; +import { ALLOWED_URLS } from "@/lib/constants"; import { Button } from "@/components/ui/button"; -import { Input } from "@/components/ui/input"; import { Label } from "@/components/ui/label"; import { Badge } from "@/components/ui/badge"; import { Alert, AlertTitle, AlertDescription } from "@/components/ui/alert"; @@ -179,7 +179,7 @@ function ModeCard({ export default function Home() { const isMobile = useIsMobile(); - const [url, setUrl] = useState("https://example.com"); + const [url, setUrl] = useState(ALLOWED_URLS[0]); const [loading, setLoading] = useState(false); const [action, setAction] = useState("screenshot"); const [mode, setMode] = useState("serverless"); @@ -238,22 +238,26 @@ export default function Home() {
- { setUrl(e.target.value); clearResults(); }} - placeholder="https://example.com" - required - /> + className="flex h-9 w-full rounded-md border border-input bg-background px-3 py-1 text-sm shadow-xs transition-colors focus-visible:outline-none focus-visible:ring-1 focus-visible:ring-ring" + > + {ALLOWED_URLS.map((u) => ( + + ))} +
diff --git a/examples/demo/lib/constants.ts b/examples/demo/lib/constants.ts new file mode 100644 index 0000000..7e49aa0 --- /dev/null +++ b/examples/demo/lib/constants.ts @@ -0,0 +1,8 @@ +export const ALLOWED_URLS = [ + "https://example.com", + "https://ai-sdk.dev", + "https://useworkflow.dev", + "https://vercel.com", +] as const; + +export type AllowedUrl = (typeof ALLOWED_URLS)[number]; diff --git a/examples/demo/lib/rate-limit.ts b/examples/demo/lib/rate-limit.ts new file mode 100644 index 0000000..7baf1cf --- /dev/null +++ b/examples/demo/lib/rate-limit.ts @@ -0,0 +1,53 @@ +import { Ratelimit } from "@upstash/ratelimit"; +import { Redis } from "@upstash/redis"; + +let _minuteRateLimit: Ratelimit | null = null; +let _dailyRateLimit: Ratelimit | null = null; + +function getRedis(): Redis | null { + const url = process.env.KV_REST_API_URL; + const token = process.env.KV_REST_API_TOKEN; + + if (!url || !token) { + return null; + } + + return new Redis({ url, token }); +} + +const noopRateLimiter = { + limit: async () => ({ success: true, limit: 0, remaining: 0, reset: 0 }), +}; + +const MINUTE_LIMIT = Number(process.env.RATE_LIMIT_PER_MINUTE) || 10; +const DAILY_LIMIT = Number(process.env.RATE_LIMIT_PER_DAY) || 100; + +export const minuteRateLimit = { + limit: async (identifier: string) => { + if (!_minuteRateLimit) { + const redis = getRedis(); + if (!redis) return noopRateLimiter.limit(); + _minuteRateLimit = new Ratelimit({ + redis, + limiter: Ratelimit.slidingWindow(MINUTE_LIMIT, "1 m"), + prefix: "ratelimit:minute", + }); + } + return _minuteRateLimit.limit(identifier); + }, +}; + +export const dailyRateLimit = { + limit: async (identifier: string) => { + if (!_dailyRateLimit) { + const redis = getRedis(); + if (!redis) return noopRateLimiter.limit(); + _dailyRateLimit = new Ratelimit({ + redis, + limiter: Ratelimit.fixedWindow(DAILY_LIMIT, "1 d"), + prefix: "ratelimit:daily", + }); + } + return _dailyRateLimit.limit(identifier); + }, +}; diff --git a/examples/demo/package.json b/examples/demo/package.json index 3db66df..f4c012e 100644 --- a/examples/demo/package.json +++ b/examples/demo/package.json @@ -12,6 +12,8 @@ "@base-ui/react": "^1.2.0", "@sparticuz/chromium": "^143.0.4", "@tailwindcss/postcss": "^4.2.1", + "@upstash/ratelimit": "^2.0.8", + "@upstash/redis": "^1.36.4", "@vercel/sandbox": "^1.0.0", "class-variance-authority": "^0.7.1", "clsx": "^2.1.1", diff --git a/examples/demo/pnpm-lock.yaml b/examples/demo/pnpm-lock.yaml index ea06d29..11b21c6 100644 --- a/examples/demo/pnpm-lock.yaml +++ b/examples/demo/pnpm-lock.yaml @@ -17,6 +17,12 @@ importers: '@tailwindcss/postcss': specifier: ^4.2.1 version: 4.2.1 + '@upstash/ratelimit': + specifier: ^2.0.8 + version: 2.0.8(@upstash/redis@1.36.4) + '@upstash/redis': + specifier: ^1.36.4 + version: 1.36.4 '@vercel/sandbox': specifier: ^1.0.0 version: 1.8.0 @@ -693,6 +699,18 @@ packages: '@types/yauzl@2.10.3': resolution: {integrity: sha512-oJoftv0LSuaDZE3Le4DbKX+KS9G36NzOeSap90UIK0yMA/NhKJhqlSGtNDORNRaIbQfzjXDrQa0ytJ6mNRGz/Q==} + '@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.4': + resolution: {integrity: sha512-w4s/msmyMqxOxaVhC8TQ2whJ77+Zd8YaSFokXL4mULQopaYb4xNJcm/PedtFQyLJn65nneySw9IwYnlMBBmFHg==} + '@vercel/oidc@3.2.0': resolution: {integrity: sha512-UycprH3T6n3jH0k44NHMa7pnFHGu/N05MjojYr+Mc6I7obkoLIJujSWwin1pCvdy/eOxrI/l3uDLQsmcrOb4ug==} engines: {node: '>= 20'} @@ -2122,6 +2140,9 @@ packages: engines: {node: '>=14.17'} hasBin: true + uncrypto@0.1.3: + resolution: {integrity: sha512-Ql87qFHB3s/De2ClA9e0gsnS6zXG27SkTiSJwjCc9MebbfapQfuPzumMIUMi38ezPZVNFcHI9sUIepeQfw8J8Q==} + undici-types@6.21.0: resolution: {integrity: sha512-iwDZqg0QAGrg9Rav5H4n0M64c3mkR59cJ6wQp+7C4nI0gsmExaedaYLNO44eT4AtBBwjbTiGPMlt2Md0T9H9JQ==} @@ -2878,6 +2899,19 @@ snapshots: '@types/node': 22.19.15 optional: true + '@upstash/core-analytics@0.0.10': + dependencies: + '@upstash/redis': 1.36.4 + + '@upstash/ratelimit@2.0.8(@upstash/redis@1.36.4)': + dependencies: + '@upstash/core-analytics': 0.0.10 + '@upstash/redis': 1.36.4 + + '@upstash/redis@1.36.4': + dependencies: + uncrypto: 0.1.3 + '@vercel/oidc@3.2.0': {} '@vercel/sandbox@1.8.0': @@ -4310,6 +4344,8 @@ snapshots: typescript@5.9.3: {} + uncrypto@0.1.3: {} + undici-types@6.21.0: {} undici@7.22.0: {}