mirror of
https://github.com/DayuanJiang/next-ai-draw-io.git
synced 2026-01-02 14:22:28 +08:00
feat: add server-side quota tracking with DynamoDB (#379)
- Add dynamo-quota-manager.ts for atomic quota checks using ConditionExpression - Enforce daily request limit, daily token limit, and TPM limit - Return 429 with quota details (type, used, limit) when exceeded - Quota is opt-in: only enabled when DYNAMODB_QUOTA_TABLE env var is set - Remove client-side quota enforcement (server is now source of truth) - Simplify use-quota-manager.tsx to only display toasts - Add @aws-sdk/client-dynamodb dependency
This commit is contained in:
Dayuan Jiang
238
lib/dynamo-quota-manager.ts
Normal file
238
lib/dynamo-quota-manager.ts
Normal file
@@ -0,0 +1,238 @@
|
||||
import {
|
||||
ConditionalCheckFailedException,
|
||||
DynamoDBClient,
|
||||
GetItemCommand,
|
||||
UpdateItemCommand,
|
||||
} from "@aws-sdk/client-dynamodb"
|
||||
|
||||
// Quota tracking is OPT-IN: only enabled if DYNAMODB_QUOTA_TABLE is explicitly set
|
||||
// OSS users who don't need quota tracking can simply not set this env var
|
||||
const TABLE = process.env.DYNAMODB_QUOTA_TABLE
|
||||
const DYNAMODB_REGION = process.env.DYNAMODB_REGION || "ap-northeast-1"
|
||||
|
||||
// Only create client if quota is enabled
|
||||
const client = TABLE ? new DynamoDBClient({ region: DYNAMODB_REGION }) : null
|
||||
|
||||
/**
|
||||
* Check if server-side quota tracking is enabled.
|
||||
* Quota is opt-in: only enabled when DYNAMODB_QUOTA_TABLE env var is set.
|
||||
*/
|
||||
export function isQuotaEnabled(): boolean {
|
||||
return !!TABLE
|
||||
}
|
||||
|
||||
interface QuotaLimits {
|
||||
requests: number // Daily request limit
|
||||
tokens: number // Daily token limit
|
||||
tpm: number // Tokens per minute
|
||||
}
|
||||
|
||||
interface QuotaCheckResult {
|
||||
allowed: boolean
|
||||
error?: string
|
||||
type?: "request" | "token" | "tpm"
|
||||
used?: number
|
||||
limit?: number
|
||||
}
|
||||
|
||||
/**
|
||||
* Check all quotas and increment request count atomically.
|
||||
* Uses ConditionExpression to prevent race conditions.
|
||||
* Returns which limit was exceeded if any.
|
||||
*/
|
||||
export async function checkAndIncrementRequest(
|
||||
ip: string,
|
||||
limits: QuotaLimits,
|
||||
): Promise<QuotaCheckResult> {
|
||||
// Skip if quota tracking not enabled
|
||||
if (!client || !TABLE) {
|
||||
return { allowed: true }
|
||||
}
|
||||
|
||||
const today = new Date().toISOString().split("T")[0]
|
||||
const currentMinute = Math.floor(Date.now() / 60000).toString()
|
||||
const ttl = Math.floor(Date.now() / 1000) + 7 * 24 * 60 * 60
|
||||
|
||||
try {
|
||||
// Atomic check-and-increment with ConditionExpression
|
||||
// This prevents race conditions by failing if limits are exceeded
|
||||
await client.send(
|
||||
new UpdateItemCommand({
|
||||
TableName: TABLE,
|
||||
Key: { PK: { S: `IP#${ip}` } },
|
||||
// Reset counts if new day/minute, then increment request count
|
||||
UpdateExpression: `
|
||||
SET lastResetDate = :today,
|
||||
dailyReqCount = if_not_exists(dailyReqCount, :zero) + :one,
|
||||
dailyTokenCount = if_not_exists(dailyTokenCount, :zero),
|
||||
lastMinute = :minute,
|
||||
tpmCount = if_not_exists(tpmCount, :zero),
|
||||
#ttl = :ttl
|
||||
`,
|
||||
// Atomic condition: only succeed if ALL limits pass
|
||||
// Uses attribute_not_exists for new items, then checks limits for existing items
|
||||
ConditionExpression: `
|
||||
(attribute_not_exists(lastResetDate) OR lastResetDate < :today OR
|
||||
((attribute_not_exists(dailyReqCount) OR dailyReqCount < :reqLimit) AND
|
||||
(attribute_not_exists(dailyTokenCount) OR dailyTokenCount < :tokenLimit))) AND
|
||||
(attribute_not_exists(lastMinute) OR lastMinute <> :minute OR
|
||||
attribute_not_exists(tpmCount) OR tpmCount < :tpmLimit)
|
||||
`,
|
||||
ExpressionAttributeNames: { "#ttl": "ttl" },
|
||||
ExpressionAttributeValues: {
|
||||
":today": { S: today },
|
||||
":zero": { N: "0" },
|
||||
":one": { N: "1" },
|
||||
":minute": { S: currentMinute },
|
||||
":ttl": { N: String(ttl) },
|
||||
":reqLimit": { N: String(limits.requests || 999999) },
|
||||
":tokenLimit": { N: String(limits.tokens || 999999) },
|
||||
":tpmLimit": { N: String(limits.tpm || 999999) },
|
||||
},
|
||||
}),
|
||||
)
|
||||
|
||||
return { allowed: true }
|
||||
} catch (e: any) {
|
||||
// Condition failed - need to determine which limit was exceeded
|
||||
if (e instanceof ConditionalCheckFailedException) {
|
||||
// Get current counts to determine which limit was hit
|
||||
try {
|
||||
const getResult = await client.send(
|
||||
new GetItemCommand({
|
||||
TableName: TABLE,
|
||||
Key: { PK: { S: `IP#${ip}` } },
|
||||
}),
|
||||
)
|
||||
|
||||
const item = getResult.Item
|
||||
const storedDate = item?.lastResetDate?.S
|
||||
const storedMinute = item?.lastMinute?.S
|
||||
const isNewDay = !storedDate || storedDate < today
|
||||
|
||||
const dailyReqCount = isNewDay
|
||||
? 0
|
||||
: Number(item?.dailyReqCount?.N || 0)
|
||||
const dailyTokenCount = isNewDay
|
||||
? 0
|
||||
: Number(item?.dailyTokenCount?.N || 0)
|
||||
const tpmCount =
|
||||
storedMinute !== currentMinute
|
||||
? 0
|
||||
: Number(item?.tpmCount?.N || 0)
|
||||
|
||||
// Determine which limit was exceeded
|
||||
if (limits.requests > 0 && dailyReqCount >= limits.requests) {
|
||||
return {
|
||||
allowed: false,
|
||||
type: "request",
|
||||
error: "Daily request limit exceeded",
|
||||
used: dailyReqCount,
|
||||
limit: limits.requests,
|
||||
}
|
||||
}
|
||||
if (limits.tokens > 0 && dailyTokenCount >= limits.tokens) {
|
||||
return {
|
||||
allowed: false,
|
||||
type: "token",
|
||||
error: "Daily token limit exceeded",
|
||||
used: dailyTokenCount,
|
||||
limit: limits.tokens,
|
||||
}
|
||||
}
|
||||
if (limits.tpm > 0 && tpmCount >= limits.tpm) {
|
||||
return {
|
||||
allowed: false,
|
||||
type: "tpm",
|
||||
error: "Rate limit exceeded (tokens per minute)",
|
||||
used: tpmCount,
|
||||
limit: limits.tpm,
|
||||
}
|
||||
}
|
||||
|
||||
// Condition failed but no limit clearly exceeded - race condition edge case
|
||||
// Fail safe by allowing (could be a reset race)
|
||||
console.warn(
|
||||
`[quota] Condition failed but no limit exceeded for IP prefix: ${ip.slice(0, 8)}...`,
|
||||
)
|
||||
return { allowed: true }
|
||||
} catch (getError: any) {
|
||||
console.error(
|
||||
`[quota] Failed to get quota details after condition failure, IP prefix: ${ip.slice(0, 8)}..., error: ${getError.message}`,
|
||||
)
|
||||
return { allowed: true } // Fail open
|
||||
}
|
||||
}
|
||||
|
||||
// Other DynamoDB errors - fail open
|
||||
console.error(
|
||||
`[quota] DynamoDB error (fail-open), IP prefix: ${ip.slice(0, 8)}..., error: ${e.message}`,
|
||||
)
|
||||
return { allowed: true }
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Record token usage after response completes.
|
||||
* Uses atomic operations to update both daily token count and TPM count.
|
||||
* Handles minute boundaries atomically to prevent race conditions.
|
||||
*/
|
||||
export async function recordTokenUsage(
|
||||
ip: string,
|
||||
tokens: number,
|
||||
): Promise<void> {
|
||||
// Skip if quota tracking not enabled
|
||||
if (!client || !TABLE) return
|
||||
if (!Number.isFinite(tokens) || tokens <= 0) return
|
||||
|
||||
const currentMinute = Math.floor(Date.now() / 60000).toString()
|
||||
const ttl = Math.floor(Date.now() / 1000) + 7 * 24 * 60 * 60
|
||||
|
||||
try {
|
||||
// Try to update assuming same minute (most common case)
|
||||
// Uses condition to ensure we're in the same minute
|
||||
await client.send(
|
||||
new UpdateItemCommand({
|
||||
TableName: TABLE,
|
||||
Key: { PK: { S: `IP#${ip}` } },
|
||||
UpdateExpression:
|
||||
"SET #ttl = :ttl ADD dailyTokenCount :tokens, tpmCount :tokens",
|
||||
ConditionExpression: "lastMinute = :minute",
|
||||
ExpressionAttributeNames: { "#ttl": "ttl" },
|
||||
ExpressionAttributeValues: {
|
||||
":minute": { S: currentMinute },
|
||||
":tokens": { N: String(tokens) },
|
||||
":ttl": { N: String(ttl) },
|
||||
},
|
||||
}),
|
||||
)
|
||||
} catch (e: any) {
|
||||
if (e instanceof ConditionalCheckFailedException) {
|
||||
// Different minute - reset TPM count and set new minute
|
||||
try {
|
||||
await client.send(
|
||||
new UpdateItemCommand({
|
||||
TableName: TABLE,
|
||||
Key: { PK: { S: `IP#${ip}` } },
|
||||
UpdateExpression:
|
||||
"SET lastMinute = :minute, tpmCount = :tokens, #ttl = :ttl ADD dailyTokenCount :tokens",
|
||||
ExpressionAttributeNames: { "#ttl": "ttl" },
|
||||
ExpressionAttributeValues: {
|
||||
":minute": { S: currentMinute },
|
||||
":tokens": { N: String(tokens) },
|
||||
":ttl": { N: String(ttl) },
|
||||
},
|
||||
}),
|
||||
)
|
||||
} catch (retryError: any) {
|
||||
console.error(
|
||||
`[quota] Failed to record tokens (retry), IP prefix: ${ip.slice(0, 8)}..., tokens: ${tokens}, error: ${retryError.message}`,
|
||||
)
|
||||
}
|
||||
} else {
|
||||
console.error(
|
||||
`[quota] Failed to record tokens, IP prefix: ${ip.slice(0, 8)}..., tokens: ${tokens}, error: ${e.message}`,
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,11 +1,10 @@
|
||||
"use client"
|
||||
|
||||
import { useCallback, useMemo } from "react"
|
||||
import { useCallback } from "react"
|
||||
import { toast } from "sonner"
|
||||
import { QuotaLimitToast } from "@/components/quota-limit-toast"
|
||||
import { useDictionary } from "@/hooks/use-dictionary"
|
||||
import { formatMessage } from "@/lib/i18n/utils"
|
||||
import { STORAGE_KEYS } from "@/lib/storage"
|
||||
|
||||
export interface QuotaConfig {
|
||||
dailyRequestLimit: number
|
||||
@@ -13,134 +12,19 @@ export interface QuotaConfig {
|
||||
tpmLimit: number
|
||||
}
|
||||
|
||||
export interface QuotaCheckResult {
|
||||
allowed: boolean
|
||||
remaining: number
|
||||
used: number
|
||||
}
|
||||
|
||||
/**
|
||||
* Hook for managing request/token quotas and rate limiting.
|
||||
* Handles three types of limits:
|
||||
* - Daily request limit
|
||||
* - Daily token limit
|
||||
* - Tokens per minute (TPM) rate limit
|
||||
*
|
||||
* Users with their own API key bypass all limits.
|
||||
* Hook for displaying quota limit toasts.
|
||||
* Server-side handles actual quota enforcement via DynamoDB.
|
||||
* This hook only provides UI feedback when limits are exceeded.
|
||||
*/
|
||||
export function useQuotaManager(config: QuotaConfig): {
|
||||
hasOwnApiKey: () => boolean
|
||||
checkDailyLimit: () => QuotaCheckResult
|
||||
checkTokenLimit: () => QuotaCheckResult
|
||||
checkTPMLimit: () => QuotaCheckResult
|
||||
incrementRequestCount: () => void
|
||||
incrementTokenCount: (tokens: number) => void
|
||||
incrementTPMCount: (tokens: number) => void
|
||||
showQuotaLimitToast: () => void
|
||||
showTokenLimitToast: (used: number) => void
|
||||
showTPMLimitToast: () => void
|
||||
} {
|
||||
const { dailyRequestLimit, dailyTokenLimit, tpmLimit } = config
|
||||
|
||||
const dict = useDictionary()
|
||||
|
||||
// Check if user has their own API key configured (bypass limits)
|
||||
const hasOwnApiKey = useCallback((): boolean => {
|
||||
const provider = localStorage.getItem(STORAGE_KEYS.aiProvider)
|
||||
const apiKey = localStorage.getItem(STORAGE_KEYS.aiApiKey)
|
||||
return !!(provider && apiKey)
|
||||
}, [])
|
||||
|
||||
// Generic helper: Parse count from localStorage with NaN guard
|
||||
const parseStorageCount = (key: string): number => {
|
||||
const count = parseInt(localStorage.getItem(key) || "0", 10)
|
||||
return Number.isNaN(count) ? 0 : count
|
||||
}
|
||||
|
||||
// Generic helper: Create quota checker factory
|
||||
const createQuotaChecker = useCallback(
|
||||
(
|
||||
getTimeKey: () => string,
|
||||
timeStorageKey: string,
|
||||
countStorageKey: string,
|
||||
limit: number,
|
||||
) => {
|
||||
return (): QuotaCheckResult => {
|
||||
if (hasOwnApiKey())
|
||||
return { allowed: true, remaining: -1, used: 0 }
|
||||
if (limit <= 0) return { allowed: true, remaining: -1, used: 0 }
|
||||
|
||||
const currentTime = getTimeKey()
|
||||
const storedTime = localStorage.getItem(timeStorageKey)
|
||||
let count = parseStorageCount(countStorageKey)
|
||||
|
||||
if (storedTime !== currentTime) {
|
||||
count = 0
|
||||
localStorage.setItem(timeStorageKey, currentTime)
|
||||
localStorage.setItem(countStorageKey, "0")
|
||||
}
|
||||
|
||||
return {
|
||||
allowed: count < limit,
|
||||
remaining: limit - count,
|
||||
used: count,
|
||||
}
|
||||
}
|
||||
},
|
||||
[hasOwnApiKey],
|
||||
)
|
||||
|
||||
// Generic helper: Create quota incrementer factory
|
||||
const createQuotaIncrementer = useCallback(
|
||||
(
|
||||
getTimeKey: () => string,
|
||||
timeStorageKey: string,
|
||||
countStorageKey: string,
|
||||
validateInput: boolean = false,
|
||||
) => {
|
||||
return (tokens: number = 1): void => {
|
||||
if (validateInput && (!Number.isFinite(tokens) || tokens <= 0))
|
||||
return
|
||||
|
||||
const currentTime = getTimeKey()
|
||||
const storedTime = localStorage.getItem(timeStorageKey)
|
||||
let count = parseStorageCount(countStorageKey)
|
||||
|
||||
if (storedTime !== currentTime) {
|
||||
count = 0
|
||||
localStorage.setItem(timeStorageKey, currentTime)
|
||||
}
|
||||
|
||||
localStorage.setItem(countStorageKey, String(count + tokens))
|
||||
}
|
||||
},
|
||||
[],
|
||||
)
|
||||
|
||||
// Check daily request limit
|
||||
const checkDailyLimit = useMemo(
|
||||
() =>
|
||||
createQuotaChecker(
|
||||
() => new Date().toDateString(),
|
||||
STORAGE_KEYS.requestDate,
|
||||
STORAGE_KEYS.requestCount,
|
||||
dailyRequestLimit,
|
||||
),
|
||||
[createQuotaChecker, dailyRequestLimit],
|
||||
)
|
||||
|
||||
// Increment request count
|
||||
const incrementRequestCount = useMemo(
|
||||
() =>
|
||||
createQuotaIncrementer(
|
||||
() => new Date().toDateString(),
|
||||
STORAGE_KEYS.requestDate,
|
||||
STORAGE_KEYS.requestCount,
|
||||
false,
|
||||
),
|
||||
[createQuotaIncrementer],
|
||||
)
|
||||
|
||||
// Show quota limit toast (request-based)
|
||||
const showQuotaLimitToast = useCallback(() => {
|
||||
toast.custom(
|
||||
@@ -155,30 +39,6 @@ export function useQuotaManager(config: QuotaConfig): {
|
||||
)
|
||||
}, [dailyRequestLimit])
|
||||
|
||||
// Check daily token limit
|
||||
const checkTokenLimit = useMemo(
|
||||
() =>
|
||||
createQuotaChecker(
|
||||
() => new Date().toDateString(),
|
||||
STORAGE_KEYS.tokenDate,
|
||||
STORAGE_KEYS.tokenCount,
|
||||
dailyTokenLimit,
|
||||
),
|
||||
[createQuotaChecker, dailyTokenLimit],
|
||||
)
|
||||
|
||||
// Increment token count
|
||||
const incrementTokenCount = useMemo(
|
||||
() =>
|
||||
createQuotaIncrementer(
|
||||
() => new Date().toDateString(),
|
||||
STORAGE_KEYS.tokenDate,
|
||||
STORAGE_KEYS.tokenCount,
|
||||
true, // Validate input tokens
|
||||
),
|
||||
[createQuotaIncrementer],
|
||||
)
|
||||
|
||||
// Show token limit toast
|
||||
const showTokenLimitToast = useCallback(
|
||||
(used: number) => {
|
||||
@@ -197,30 +57,6 @@ export function useQuotaManager(config: QuotaConfig): {
|
||||
[dailyTokenLimit],
|
||||
)
|
||||
|
||||
// Check TPM (tokens per minute) limit
|
||||
const checkTPMLimit = useMemo(
|
||||
() =>
|
||||
createQuotaChecker(
|
||||
() => Math.floor(Date.now() / 60000).toString(),
|
||||
STORAGE_KEYS.tpmMinute,
|
||||
STORAGE_KEYS.tpmCount,
|
||||
tpmLimit,
|
||||
),
|
||||
[createQuotaChecker, tpmLimit],
|
||||
)
|
||||
|
||||
// Increment TPM count
|
||||
const incrementTPMCount = useMemo(
|
||||
() =>
|
||||
createQuotaIncrementer(
|
||||
() => Math.floor(Date.now() / 60000).toString(),
|
||||
STORAGE_KEYS.tpmMinute,
|
||||
STORAGE_KEYS.tpmCount,
|
||||
true, // Validate input tokens
|
||||
),
|
||||
[createQuotaIncrementer],
|
||||
)
|
||||
|
||||
// Show TPM limit toast
|
||||
const showTPMLimitToast = useCallback(() => {
|
||||
const limitDisplay =
|
||||
@@ -233,18 +69,6 @@ export function useQuotaManager(config: QuotaConfig): {
|
||||
}, [tpmLimit, dict])
|
||||
|
||||
return {
|
||||
// Check functions
|
||||
hasOwnApiKey,
|
||||
checkDailyLimit,
|
||||
checkTokenLimit,
|
||||
checkTPMLimit,
|
||||
|
||||
// Increment functions
|
||||
incrementRequestCount,
|
||||
incrementTokenCount,
|
||||
incrementTPMCount,
|
||||
|
||||
// Toast functions
|
||||
showQuotaLimitToast,
|
||||
showTokenLimitToast,
|
||||
showTPMLimitToast,
|
||||
|
||||
Reference in New Issue
Block a user