From 15b3a3cf9d757fa52a06948b253cb2f7ffe1e28e Mon Sep 17 00:00:00 2001 From: Roman Date: Wed, 24 Dec 2025 14:13:14 +0000 Subject: [PATCH] feat: basic rate limits --- .env.example | 13 ++ messages/en.json | 5 + messages/es.json | 8 +- messages/fr.json | 8 +- messages/ja.json | 7 +- messages/ko.json | 7 +- messages/zh.json | 5 + src/app/api/chat/route.ts | 18 +++ src/app/api/chat/temporary/route.ts | 18 +++ src/components/message.tsx | 48 +++++++ src/lib/ai/rate-limit-message.ts | 49 +++++++ src/lib/ai/rate-limit.test.ts | 101 ++++++++++++++ src/lib/ai/rate-limit.ts | 195 ++++++++++++++++++++++++++++ 13 files changed, 472 insertions(+), 10 deletions(-) create mode 100644 src/lib/ai/rate-limit-message.ts create mode 100644 src/lib/ai/rate-limit.test.ts create mode 100644 src/lib/ai/rate-limit.ts diff --git a/.env.example b/.env.example index e1a6aae00..4d07a43fd 100644 --- a/.env.example +++ b/.env.example @@ -87,6 +87,19 @@ DISABLE_EMAIL_SIGN_UP= # Set this to 1 to disable OAuth sign-ups (Google, GitHub, Microsoft) DISABLE_SIGN_UP= +# (Optional) AI rate limits (per user, by role). Requires REDIS_URL. +# Leave empty to disable for that role. +AI_RATE_LIMIT_USER_PER_HOUR= +AI_RATE_LIMIT_USER_PER_DAY= +AI_RATE_LIMIT_EDITOR_PER_HOUR= +AI_RATE_LIMIT_EDITOR_PER_DAY= +AI_RATE_LIMIT_ADMIN_PER_HOUR= +AI_RATE_LIMIT_ADMIN_PER_DAY= + +# (Optional) Fallback limits for roles not set above +AI_RATE_LIMIT_PER_HOUR= +AI_RATE_LIMIT_PER_DAY= + # (Optional) # Set this to 1 to disallow adding MCP servers. NOT_ALLOW_ADD_MCP_SERVERS= diff --git a/messages/en.json b/messages/en.json index 1ac73a758..99b99d3db 100644 --- a/messages/en.json +++ b/messages/en.json @@ -194,6 +194,11 @@ "Chat": { "Error": "Chat Error", "thisMessageWasNotSavedPleaseTryTheChatAgain": "This message was not saved. Please try the chat again.", + "rateLimitedTitle": "You’ve hit the AI limit", + "rateLimitedHourly": "You can send up to {limit} AI messages per hour.", + "rateLimitedDaily": "You can send up to {limit} AI messages per day.", + "rateLimitedRetryAfter": "Try again in {minutes} min", + "rateLimitedSupport": "Need more? Ask your admin.", "uploadImage": "Upload File", "generateImage": "Generate Image", "imageUploadedSuccessfully": "Image uploaded successfully", diff --git a/messages/es.json b/messages/es.json index 144b3d61a..f60d326cd 100644 --- a/messages/es.json +++ b/messages/es.json @@ -74,6 +74,11 @@ "Chat": { "Error": "Error de Chat", "thisMessageWasNotSavedPleaseTryTheChatAgain": "Este mensaje no se guardó. Por favor, intenta el chat nuevamente.", + "rateLimitedTitle": "Has alcanzado el límite de IA", + "rateLimitedHourly": "Puedes enviar hasta {limit} mensajes de IA por hora.", + "rateLimitedDaily": "Puedes enviar hasta {limit} mensajes de IA por día.", + "rateLimitedRetryAfter": "Inténtalo de nuevo en {minutes} min", + "rateLimitedSupport": "¿Necesitas más? Habla con tu administrador.", "uploadImage": "Subir archivo", "generateImage": "Generar Imagen", "imageUploadedSuccessfully": "Imagen subida exitosamente", @@ -146,7 +151,6 @@ }, "Thread": { "chat": "Chat", - "renameChat": "Renombrar", "deleteChat": "Eliminar Chat", "deleteUnarchivedChats": "Eliminar Todos los Chats No Archivados", @@ -166,7 +170,6 @@ "createLink": "Crear Enlace", "linkCopied": "Enlace copiado" }, - "ChatPreferences": { "title": "Preferencias de Chat", "whatShouldWeCallYou": "¿Cómo deberíamos llamarte?", @@ -233,7 +236,6 @@ "theme": "Tema", "signOut": "Cerrar sesión", "language": "Idioma", - "showAllChats": "Ver Todos los Chats", "showLessChats": "Mostrar menos", "reportAnIssue": "Reportar un problema", diff --git a/messages/fr.json b/messages/fr.json index ea6e7be16..f14b6f765 100644 --- a/messages/fr.json +++ b/messages/fr.json @@ -74,6 +74,11 @@ "Chat": { "Error": "Erreur de Chat", "thisMessageWasNotSavedPleaseTryTheChatAgain": "Ce message n'a pas été enregistré. Veuillez réessayer le chat.", + "rateLimitedTitle": "Vous avez atteint la limite IA", + "rateLimitedHourly": "Vous pouvez envoyer jusqu'à {limit} messages IA par heure.", + "rateLimitedDaily": "Vous pouvez envoyer jusqu'à {limit} messages IA par jour.", + "rateLimitedRetryAfter": "Réessayez dans {minutes} min", + "rateLimitedSupport": "Besoin de plus ? Contactez votre administrateur.", "uploadImage": "Téléverser un fichier", "generateImage": "Générer une Image", "imageUploadedSuccessfully": "Image téléchargée avec succès", @@ -146,7 +151,6 @@ }, "Thread": { "chat": "Chat", - "renameChat": "Renommer", "deleteChat": "Supprimer le Chat", "deleteUnarchivedChats": "Supprimer Tous les Chats Non Archivés", @@ -166,7 +170,6 @@ "createLink": "Créer le Lien", "linkCopied": "Lien copié" }, - "ChatPreferences": { "title": "Préférences de Chat", "whatShouldWeCallYou": "Comment devrions-nous vous appeler ?", @@ -233,7 +236,6 @@ "theme": "Thème", "signOut": "Se déconnecter", "language": "Langue", - "showAllChats": "Voir Tous les Chats", "showLessChats": "Afficher moins", "reportAnIssue": "Signaler un problème", diff --git a/messages/ja.json b/messages/ja.json index e3ab17a40..ba760efb5 100644 --- a/messages/ja.json +++ b/messages/ja.json @@ -74,6 +74,11 @@ "Chat": { "Error": "チャットエラー", "thisMessageWasNotSavedPleaseTryTheChatAgain": "このメッセージは保存されませんでした。もう一度チャットをお試しください。", + "rateLimitedTitle": "AI の制限に達しました", + "rateLimitedHourly": "1 時間に {limit} 件まで AI メッセージを送信できます。", + "rateLimitedDaily": "1 日に {limit} 件まで AI メッセージを送信できます。", + "rateLimitedRetryAfter": "{minutes} 分後に再試行してください", + "rateLimitedSupport": "上限を増やすには管理者に連絡してください。", "uploadImage": "ファイルをアップロード", "generateImage": "画像を生成", "imageUploadedSuccessfully": "画像が正常にアップロードされました", @@ -146,7 +151,6 @@ }, "Thread": { "chat": "チャット", - "renameChat": "名前を変更", "deleteChat": "チャットを削除", "deleteUnarchivedChats": "アーカイブされていないチャットをすべて削除", @@ -166,7 +170,6 @@ "createLink": "リンクを作成", "linkCopied": "リンクがコピーされました" }, - "ChatPreferences": { "title": "チャット設定", "whatShouldWeCallYou": "何とお呼びすればよろしいですか？", diff --git a/messages/ko.json b/messages/ko.json index 07fae0713..e8b8fc190 100644 --- a/messages/ko.json +++ b/messages/ko.json @@ -75,6 +75,11 @@ "Chat": { "Error": "채팅 오류", "thisMessageWasNotSavedPleaseTryTheChatAgain": "이 메시지는 저장되지 않았습니다. 다시 시도해주세요.", + "rateLimitedTitle": "AI 한도에 도달했습니다", + "rateLimitedHourly": "한 시간에 최대 {limit}개의 AI 메시지를 보낼 수 있습니다.", + "rateLimitedDaily": "하루에 최대 {limit}개의 AI 메시지를 보낼 수 있습니다.", + "rateLimitedRetryAfter": "{minutes}분 후에 다시 시도해주세요", + "rateLimitedSupport": "더 필요하면 관리자에게 문의하세요.", "uploadImage": "파일 업로드", "generateImage": "이미지 만들기", "imageUploadedSuccessfully": "이미지가 성공적으로 업로드되었습니다", @@ -147,7 +152,6 @@ }, "Thread": { "chat": "채팅", - "renameChat": "채팅 이름 변경", "deleteChat": "채팅 삭제", "deleteUnarchivedChats": "아카이브되지 않은 채팅 모두 삭제", @@ -167,7 +171,6 @@ "createLink": "링크 만들기", "linkCopied": "링크가 복사되었습니다" }, - "ChatPreferences": { "title": "채팅 환경설정", "whatShouldWeCallYou": "뭐라고 불러드릴까요?", diff --git a/messages/zh.json b/messages/zh.json index d64b07727..7892f1be6 100644 --- a/messages/zh.json +++ b/messages/zh.json @@ -75,6 +75,11 @@ "Chat": { "Error": "聊天错误", "thisMessageWasNotSavedPleaseTryTheChatAgain": "此消息未保存。请重试聊天。", + "rateLimitedTitle": "已达到 AI 限制", + "rateLimitedHourly": "每小时最多可发送 {limit} 条 AI 消息。", + "rateLimitedDaily": "每天最多可发送 {limit} 条 AI 消息。", + "rateLimitedRetryAfter": "请在 {minutes} 分钟后重试", + "rateLimitedSupport": "需要更高额度？请联系管理员。", "uploadImage": "上传文件", "generateImage": "生成图片", "imageUploadedSuccessfully": "图片上传成功", diff --git a/src/app/api/chat/route.ts b/src/app/api/chat/route.ts index 9f6e4e21c..c06100762 100644 --- a/src/app/api/chat/route.ts +++ b/src/app/api/chat/route.ts @@ -51,6 +51,8 @@ import { nanoBananaTool, openaiImageTool } from "lib/ai/tools/image"; import { ImageToolName } from "lib/ai/tools"; import { buildCsvIngestionPreviewParts } from "@/lib/ai/ingest/csv-ingest"; import { serverFileStorage } from "lib/file-storage"; +import { getAiRateLimiter } from "lib/ai/rate-limit"; +import { buildRateLimitMessage } from "lib/ai/rate-limit-message"; const logger = globalLogger.withDefaults({ message: colorize("blackBright", `Chat API: `), @@ -65,6 +67,22 @@ export async function POST(request: Request) { if (!session?.user.id) { return new Response("Unauthorized", { status: 401 }); } + + const rateLimiter = getAiRateLimiter(); + if (rateLimiter) { + const rateLimitResult = await rateLimiter.check( + session.user.id, + // eslint-disable-next-line @typescript-eslint/no-explicit-any + (session.user as any).role, + ); + if (!rateLimitResult.ok) { + const message = buildRateLimitMessage(rateLimitResult); + return new Response(message, { + status: 429, + headers: { "Retry-After": `${rateLimitResult.retryAfterSeconds}` }, + }); + } + } const { id, message, diff --git a/src/app/api/chat/temporary/route.ts b/src/app/api/chat/temporary/route.ts index f17f0fb20..08e0d8903 100644 --- a/src/app/api/chat/temporary/route.ts +++ b/src/app/api/chat/temporary/route.ts @@ -9,6 +9,8 @@ import { customModelProvider } from "lib/ai/models"; import globalLogger from "logger"; import { buildUserSystemPrompt } from "lib/ai/prompts"; import { getUserPreferences } from "lib/user/server"; +import { getAiRateLimiter } from "lib/ai/rate-limit"; +import { buildRateLimitMessage } from "lib/ai/rate-limit-message"; import { colorize } from "consola/utils"; @@ -25,6 +27,22 @@ export async function POST(request: Request) { return new Response("Unauthorized", { status: 401 }); } + const rateLimiter = getAiRateLimiter(); + if (rateLimiter) { + const rateLimitResult = await rateLimiter.check( + session.user.id, + // eslint-disable-next-line @typescript-eslint/no-explicit-any + (session.user as any).role, + ); + if (!rateLimitResult.ok) { + const message = buildRateLimitMessage(rateLimitResult); + return new Response(message, { + status: 429, + headers: { "Retry-After": `${rateLimitResult.retryAfterSeconds}` }, + }); + } + } + const { messages, chatModel, instructions } = json as { messages: UIMessage[]; chatModel?: { diff --git a/src/components/message.tsx b/src/components/message.tsx index 3bcdebc7c..15a9d517f 100644 --- a/src/components/message.tsx +++ b/src/components/message.tsx @@ -18,6 +18,7 @@ import { ChevronDown, ChevronUp, TriangleAlertIcon } from "lucide-react"; import { Button } from "ui/button"; import { useTranslations } from "next-intl"; import { ChatMetadata } from "app-types/chat"; +import { parseRateLimitMessage } from "lib/ai/rate-limit-message"; interface Props { message: UIMessage; @@ -210,6 +211,53 @@ export const ErrorMessage = ({ const [isExpanded, setIsExpanded] = useState(false); const maxLength = 200; const t = useTranslations(); + const parsedRateLimit = parseRateLimitMessage(error.message); + + if (parsedRateLimit) { + const retryMinutes = Math.max( + 1, + Math.ceil(parsedRateLimit.retryAfterSeconds / 60), + ); + + return ( +
+
+
+
+
+
+ +
+
+

+ {t("Chat.rateLimitedTitle")} +

+

+ {parsedRateLimit.window === "hour" + ? t("Chat.rateLimitedHourly", { + limit: parsedRateLimit.limit, + }) + : t("Chat.rateLimitedDaily", { + limit: parsedRateLimit.limit, + })} +

+
+
+ +
+ + + {t("Chat.rateLimitedRetryAfter", { minutes: retryMinutes })} + + + {t("Chat.rateLimitedSupport")} + +
+
+
+
+ ); + } return (
diff --git a/src/lib/ai/rate-limit-message.ts b/src/lib/ai/rate-limit-message.ts new file mode 100644 index 000000000..a8c2dfa76 --- /dev/null +++ b/src/lib/ai/rate-limit-message.ts @@ -0,0 +1,49 @@ +export const RATE_LIMIT_ERROR_PREFIX = "AI_RATE_LIMIT"; + +export type RateLimitWindow = "hour" | "day"; + +export interface RateLimitMessagePayload { + window: RateLimitWindow; + retryAfterSeconds: number; + limit: number; +} + +export const buildRateLimitMessage = ( + payload: RateLimitMessagePayload, +): string => { + return [ + RATE_LIMIT_ERROR_PREFIX, + payload.window, + Math.max(0, Math.ceil(payload.retryAfterSeconds)), + payload.limit, + ].join("|"); +}; + +export const parseRateLimitMessage = ( + message?: string, +): RateLimitMessagePayload | null => { + if (!message) return null; + const parts = message.split("|"); + if (parts[0] !== RATE_LIMIT_ERROR_PREFIX) return null; + if (parts.length < 4) return null; + + const window = parts[1] as RateLimitWindow; + const retryAfterSeconds = Number(parts[2]); + const limit = Number(parts[3]); + + if (!Number.isFinite(retryAfterSeconds) || retryAfterSeconds < 0) { + return null; + } + if (!Number.isFinite(limit) || limit <= 0) { + return null; + } + if (window !== "hour" && window !== "day") { + return null; + } + + return { + window, + retryAfterSeconds, + limit, + }; +}; diff --git a/src/lib/ai/rate-limit.test.ts b/src/lib/ai/rate-limit.test.ts new file mode 100644 index 000000000..30b640a2d --- /dev/null +++ b/src/lib/ai/rate-limit.test.ts @@ -0,0 +1,101 @@ +import { describe, expect, it, vi } from "vitest"; +import { AiRateLimiter, MemoryRateLimitStore } from "./rate-limit"; + +const createLimiter = (options: { + perHour?: number; + perDay?: number; + role?: string; + defaultPerDay?: number; +}) => { + return new AiRateLimiter({ + redisUrl: "redis://example.com", + limitsByRole: options.role + ? { [options.role]: { perHour: options.perHour, perDay: options.perDay } } + : {}, + defaultLimits: options.defaultPerDay + ? { perDay: options.defaultPerDay, perHour: options.perHour } + : { perHour: options.perHour, perDay: options.perDay }, + store: new MemoryRateLimitStore(), + }); +}; + +describe("AiRateLimiter", () => { + it("allows requests under the limit", async () => { + const limiter = createLimiter({ perHour: 2, role: "user" }); + const userId = "user-1"; + + const first = await limiter.check(userId, "user"); + const second = await limiter.check(userId, "user"); + + expect(first.ok).toBe(true); + expect(second.ok).toBe(true); + }); + + it("blocks when hourly limit exceeded", async () => { + const limiter = createLimiter({ perHour: 1, role: "user" }); + const userId = "user-2"; + + await limiter.check(userId, "user"); + const blocked = await limiter.check(userId, "user"); + + expect(blocked.ok).toBe(false); + if (!blocked.ok) { + expect(blocked.window).toBe("hour"); + expect(blocked.limit).toBe(1); + expect(blocked.retryAfterSeconds).toBeGreaterThan(0); + } + }); + + it("resets after window expires", async () => { + const realNow = Date.now; + let now = realNow(); + // Override Date.now to simulate time passing without timers API + Date.now = () => now; + + try { + const limiter = createLimiter({ perHour: 1, role: "user" }); + const userId = "user-3"; + + await limiter.check(userId, "user"); + const blocked = await limiter.check(userId, "user"); + expect(blocked.ok).toBe(false); + + // Advance just over one hour + now = now + 60 * 60 * 1000 + 10; + + const allowedAgain = await limiter.check(userId, "user"); + expect(allowedAgain.ok).toBe(true); + } finally { + Date.now = realNow; + } + }); + + it("applies daily limit when configured", async () => { + const limiter = createLimiter({ perHour: 0, perDay: 2, role: "editor" }); + const userId = "user-4"; + + await limiter.check(userId, "editor"); + await limiter.check(userId, "editor"); + const blocked = await limiter.check(userId, "editor"); + + expect(blocked.ok).toBe(false); + if (!blocked.ok) { + expect(blocked.window).toBe("day"); + expect(blocked.limit).toBe(2); + } + }); + + it("uses default limits when role not configured", async () => { + const limiter = createLimiter({ perHour: 0, defaultPerDay: 1 }); + const userId = "user-5"; + + await limiter.check(userId, "unknown"); + const blocked = await limiter.check(userId, "unknown"); + + expect(blocked.ok).toBe(false); + if (!blocked.ok) { + expect(blocked.window).toBe("day"); + expect(blocked.limit).toBe(1); + } + }); +}); diff --git a/src/lib/ai/rate-limit.ts b/src/lib/ai/rate-limit.ts new file mode 100644 index 000000000..095f46067 --- /dev/null +++ b/src/lib/ai/rate-limit.ts @@ -0,0 +1,195 @@ +import Redis from "ioredis"; +import logger from "logger"; +import { RateLimitMessagePayload } from "./rate-limit-message"; + +export type RateLimitCheckResult = + | { ok: true } + | ({ ok: false } & RateLimitMessagePayload); + +export interface RateLimitConfig { + perHour?: number; + perDay?: number; +} + +export interface RateLimitOptions { + redisUrl: string; + limitsByRole?: Record; + defaultLimits?: RateLimitConfig; +} + +interface RateLimitStore { + increment( + key: string, + windowMs: number, + ): Promise<{ count: number; ttlMs: number }>; +} + +class RedisRateLimitStore implements RateLimitStore { + private redis: Redis; + + constructor(redisUrl: string) { + this.redis = new Redis(redisUrl, { + // Allow brief disconnects without throwing "stream isn't writeable" + enableOfflineQueue: true, + maxRetriesPerRequest: 2, + connectTimeout: 5000, + commandTimeout: 5000, + }); + } + + async increment(key: string, windowMs: number) { + const count = await this.redis.incr(key); + if (count === 1) { + await this.redis.pexpire(key, windowMs); + } + const ttlMs = await this.redis.pttl(key); + return { count, ttlMs }; + } +} + +export class AiRateLimiter { + private store: RateLimitStore; + private limitsByRole: Record; + private defaultLimits?: RateLimitConfig; + + constructor(options: RateLimitOptions & { store?: RateLimitStore }) { + if (!options.redisUrl) { + throw new Error("Redis URL is required for AI rate limiting"); + } + this.limitsByRole = this.sanitizeLimits(options.limitsByRole ?? {}); + this.defaultLimits = this.sanitizeLimits({ + default: options.defaultLimits ?? {}, + }).default; + this.store = options.store ?? new RedisRateLimitStore(options.redisUrl); + } + + private safeLimit(limit?: number) { + if (!limit || limit <= 0 || !Number.isFinite(limit)) return undefined; + return Math.floor(limit); + } + + private sanitizeLimits(limits: Record) { + return Object.entries(limits).reduce>( + (acc, [role, config]) => { + const perHour = this.safeLimit(config.perHour); + const perDay = this.safeLimit(config.perDay); + if (perHour || perDay) { + acc[role.toLowerCase()] = { perHour, perDay }; + } + return acc; + }, + {}, + ); + } + + private async checkWindow( + userId: string, + window: "hour" | "day", + limit: number, + ): Promise { + const windowMs = window === "hour" ? 60 * 60 * 1000 : 24 * 60 * 60 * 1000; + const key = `ai:rate:${window}:${userId}`; + const { count, ttlMs } = await this.store.increment(key, windowMs); + + if (count > limit) { + const retryAfterSeconds = Math.max(1, Math.ceil(ttlMs / 1000)); + logger.warn( + `AI rate limit exceeded: user=${userId} window=${window} limit=${limit} count=${count}`, + ); + return { + ok: false, + window, + retryAfterSeconds, + limit, + }; + } + + return { ok: true } as const; + } + + async check(userId: string, role?: string): Promise { + const roleKey = (role ?? "").toLowerCase(); + const limits = this.limitsByRole[roleKey] || this.defaultLimits; + + if (!limits || (!limits.perHour && !limits.perDay)) { + return { ok: true } as const; + } + + if (limits.perHour) { + const result = await this.checkWindow(userId, "hour", limits.perHour); + if (!result.ok) return result; + } + + if (limits.perDay) { + const result = await this.checkWindow(userId, "day", limits.perDay); + if (!result.ok) return result; + } + + return { ok: true } as const; + } +} + +let limiter: AiRateLimiter | null = null; + +export const getAiRateLimiter = () => { + if (limiter) return limiter; + + const redisUrl = process.env.REDIS_URL; + const roles = ["USER", "EDITOR", "ADMIN"] as const; + + const readLimit = (role: string, window: "HOUR" | "DAY") => { + const value = process.env[`AI_RATE_LIMIT_${role}_PER_${window}`]; + const num = Number(value); + return Number.isFinite(num) && num > 0 ? num : undefined; + }; + + const limitsByRole = roles.reduce>( + (acc, role) => { + const perHour = readLimit(role, "HOUR"); + const perDay = readLimit(role, "DAY"); + if (perHour || perDay) { + acc[role.toLowerCase()] = { perHour, perDay }; + } + return acc; + }, + {}, + ); + + const defaultPerHour = Number(process.env.AI_RATE_LIMIT_PER_HOUR ?? "0"); + const defaultPerDay = Number(process.env.AI_RATE_LIMIT_PER_DAY ?? "0"); + + const isEnabled = Boolean( + redisUrl && + (Object.keys(limitsByRole).length > 0 || + defaultPerHour > 0 || + defaultPerDay > 0), + ); + if (!isEnabled) return null; + + limiter = new AiRateLimiter({ + redisUrl: redisUrl!, + limitsByRole, + defaultLimits: { + perHour: defaultPerHour, + perDay: defaultPerDay, + }, + }); + return limiter; +}; + +// Exported for testing +export class MemoryRateLimitStore implements RateLimitStore { + private store = new Map(); + + async increment(key: string, windowMs: number) { + const now = Date.now(); + const existing = this.store.get(key); + if (!existing || existing.expiresAt <= now) { + this.store.set(key, { count: 1, expiresAt: now + windowMs }); + return { count: 1, ttlMs: windowMs }; + } + const next = { count: existing.count + 1, expiresAt: existing.expiresAt }; + this.store.set(key, next); + return { count: next.count, ttlMs: next.expiresAt - now }; + } +}