add rate-limiting for non-admin users

Author: HugoRCD

apps/app/server/api/chats.post.ts

@@ -2,6 +2,7 @@ import type { UIMessage } from 'ai'
 import { db, schema } from '@nuxthub/db'
 import { z } from 'zod'
 import type { CreateChatBody, CreateChatResponse } from '#shared/types/chat'
+import { checkRateLimit, incrementRateLimit } from '../utils/rate-limit'
 
 const bodySchema = z.object({
   id: z.string(),
@@ -22,6 +23,13 @@ export default defineEventHandler(async (event) => {
     throw createError({ statusCode: 403, statusMessage: 'Admin access required', data: { why: 'Admin chat mode requires the admin role', fix: 'Contact an administrator to be granted access' } })
   }
 
+  if (user.role !== 'admin') {
+    const rateLimit = await checkRateLimit(user.id)
+    if (!rateLimit.allowed) {
+      throw createError({ statusCode: 429, statusMessage: 'Rate limit exceeded', data: { why: `You have reached the daily limit of ${rateLimit.limit} messages`, fix: 'Wait until tomorrow or contact an administrator' } })
+    }
+  }
+
   const [chat] = await db.insert(schema.chats).values({
     id,
     title: '',
@@ -40,5 +48,9 @@ export default defineEventHandler(async (event) => {
     parts: message.parts
   })
 
+  if (user.role !== 'admin') {
+    await incrementRateLimit(user.id)
+  }
+
   return chat satisfies CreateChatResponse
 })

apps/app/server/api/chats/[id].post.ts

@@ -11,6 +11,7 @@ import { generateTitle } from '../../utils/chat/generate-title'
 import { getAgentConfig } from '../../utils/agent-config'
 import { KV_KEYS } from '../../utils/sandbox/types'
 import { adminTools } from '../../utils/chat/admin-tools'
+import { checkRateLimit, incrementRateLimit } from '../../utils/rate-limit'
 
 defineRouteMeta({
   openAPI: {
@@ -64,6 +65,14 @@ export default defineEventHandler(async (event) => {
       throw createError({ statusCode: 403, statusMessage: 'Admin access required', data: { why: 'This chat is in admin mode and requires the admin role', fix: 'Contact an administrator to be granted access' } })
     }
 
+    let rateLimitInfo: { allowed: boolean, remaining: number, limit: number } | undefined
+    if (user.role !== 'admin') {
+      rateLimitInfo = await checkRateLimit(user.id)
+      if (!rateLimitInfo.allowed) {
+        throw createError({ statusCode: 429, statusMessage: 'Rate limit exceeded', data: { why: `You have reached the daily limit of ${rateLimitInfo.limit} messages`, fix: 'Wait until tomorrow or contact an administrator' } })
+      }
+    }
+
     const lastMessage = messages[messages.length - 1]
     if (lastMessage?.role === 'user' && messages.length > 1) {
       await db.insert(schema.messages).values({
@@ -221,6 +230,10 @@ export default defineEventHandler(async (event) => {
           }
         }
 
+        if (user.role !== 'admin') {
+          await incrementRateLimit(user.id)
+        }
+
         requestLog.set({
           outcome: 'success',
           responseMessageCount: responseMessages.length,
@@ -230,6 +243,11 @@ export default defineEventHandler(async (event) => {
       },
     })
 
+    if (rateLimitInfo) {
+      setHeader(event, 'X-RateLimit-Limit', String(rateLimitInfo.limit))
+      setHeader(event, 'X-RateLimit-Remaining', String(rateLimitInfo.remaining - 1))
+    }
+
     return createUIMessageStreamResponse({ stream })
   } catch (error) {
     requestLog.error(error instanceof Error ? error : new Error(String(error)))

apps/app/server/utils/rate-limit.ts

@@ -0,0 +1,26 @@
+import { kv } from '@nuxthub/kv'
+import { KV_KEYS } from './sandbox/types'
+
+const DAILY_MESSAGE_LIMIT = 15
+
+function todayKey(): string {
+  return new Date().toISOString().slice(0, 10)
+}
+
+export async function checkRateLimit(userId: string): Promise<{ allowed: boolean, remaining: number, limit: number }> {
+  const date = todayKey()
+  const count = await kv.get<number>(KV_KEYS.rateLimit(userId, date)) ?? 0
+
+  if (count >= DAILY_MESSAGE_LIMIT) {
+    return { allowed: false, remaining: 0, limit: DAILY_MESSAGE_LIMIT }
+  }
+
+  return { allowed: true, remaining: DAILY_MESSAGE_LIMIT - count, limit: DAILY_MESSAGE_LIMIT }
+}
+
+export async function incrementRateLimit(userId: string): Promise<void> {
+  const date = todayKey()
+  const key = KV_KEYS.rateLimit(userId, date)
+  const count = await kv.get<number>(key) ?? 0
+  await kv.set(key, count + 1, { ttl: 86400 })
+}

apps/app/server/utils/sandbox/types.ts

@@ -36,4 +36,5 @@ export const KV_KEYS = {
   AGENT_CONFIG_CACHE: 'agent:config-cache',
   session: (sessionId: string) => `session:${sessionId}`,
   ACTIVE_SANDBOX_SESSION: 'sandbox:active-session',
+  rateLimit: (userId: string, date: string) => `ratelimit:${userId}:${date}`,
 } as const