diff --git a/bun.lock b/bun.lock index 51cc671ad1..740bbf0c45 100644 --- a/bun.lock +++ b/bun.lock @@ -102,6 +102,7 @@ "@solidjs/router": "catalog:", "@solidjs/start": "catalog:", "@stripe/stripe-js": "8.6.1", + "@upstash/redis": "1.38.0", "chart.js": "4.5.1", "nitro": "3.0.1-alpha.1", "solid-js": "catalog:", @@ -2469,6 +2470,8 @@ "@ungap/structured-clone": ["@ungap/structured-clone@1.3.0", "", {}, "sha512-WmoN8qaIAo7WTYWbAZuG8PYEhn5fkz7dZrqTBZ7dtt//lL2Gwms1IcnQ5yHqjDfX8Ft5j4YzDM23f87zBfDe9g=="], + "@upstash/redis": ["@upstash/redis@1.38.0", "", { "dependencies": { "uncrypto": "^0.1.3" } }, "sha512-wu+dZBptlLy0+MCUEoHmzrY/TnmgDey3+c7EbIGwrLqAvkP8yi5MWZHYGIFtAygmL4Bkz2TdFu+eU0vFPncIcg=="], + "@valibot/to-json-schema": ["@valibot/to-json-schema@1.6.0", "", { "peerDependencies": { "valibot": "^1.3.0" } }, "sha512-d6rYyK5KVa2XdqamWgZ4/Nr+cXhxjy7lmpe6Iajw15J/jmU+gyxl2IEd1Otg1d7Rl3gOQL5reulnSypzBtYy1A=="], "@vercel/oidc": ["@vercel/oidc@3.2.0", "", {}, "sha512-UycprH3T6n3jH0k44NHMa7pnFHGu/N05MjojYr+Mc6I7obkoLIJujSWwin1pCvdy/eOxrI/l3uDLQsmcrOb4ug=="], diff --git a/infra/console.ts b/infra/console.ts index 2a74c12705..cf1fba8237 100644 --- a/infra/console.ts +++ b/infra/console.ts @@ -250,6 +250,8 @@ new sst.cloudflare.x.SolidStart("Console", { bucket, bucketNew, database, + SECRET.UpstashRedisRestUrl, + SECRET.UpstashRedisRestToken, AUTH_API_URL, STRIPE_WEBHOOK_SECRET, DISCORD_INCIDENT_WEBHOOK_URL, diff --git a/infra/secret.ts b/infra/secret.ts index d4e8b148fc..8a89c01f8b 100644 --- a/infra/secret.ts +++ b/infra/secret.ts @@ -8,4 +8,7 @@ export const SECRET = { R2AccessKey: new sst.Secret("R2AccessKey", "unknown"), R2SecretKey: new sst.Secret("R2SecretKey", "unknown"), HoneycombWebhookSecret: new random.RandomPassword("HoneycombWebhookSecret", { length: 24 }), + UpstashRedisRestUrl: new sst.Secret("UpstashRedisRestUrl"), + UpstashRedisRestToken: new sst.Secret("UpstashRedisRestToken"), } + diff --git a/packages/console/app/package.json b/packages/console/app/package.json index 86e7da2531..b93c15cdce 100644 --- a/packages/console/app/package.json +++ b/packages/console/app/package.json @@ -26,6 +26,7 @@ "@solidjs/router": "catalog:", "@solidjs/start": "catalog:", "@stripe/stripe-js": "8.6.1", + "@upstash/redis": "1.38.0", "chart.js": "4.5.1", "nitro": "3.0.1-alpha.1", "solid-js": "catalog:", diff --git a/packages/console/app/src/routes/zen/util/ipRateLimiter.ts b/packages/console/app/src/routes/zen/util/ipRateLimiter.ts index 177d78514a..81f73a4e5a 100644 --- a/packages/console/app/src/routes/zen/util/ipRateLimiter.ts +++ b/packages/console/app/src/routes/zen/util/ipRateLimiter.ts @@ -2,6 +2,7 @@ import { Database, eq, and, sql, inArray } from "@opencode-ai/console-core/drizz import { IpRateLimitTable } from "@opencode-ai/console-core/schema/ip.sql.js" import { FreeUsageLimitError } from "./error" import { logger } from "./logger" +import { buildRateLimitKey, getRedis } from "./redis" import { i18n } from "~/i18n" import { localeFromRequest } from "~/lib/language" import { Subscription } from "@opencode-ai/console-core/subscription.js" @@ -23,43 +24,62 @@ export function createRateLimiter(modelId: string, rateLimit: number | undefined const now = Date.now() const lifetimeInterval = "" const dailyInterval = rateLimit ? `${buildYYYYMMDD(now)}${modelId.substring(0, 2)}` : buildYYYYMMDD(now) - - let _isNew: boolean + const retryAfter = getRetryAfterDay(now) + const redis = getRedis() + const lifetimeKey = buildRateLimitKey("ip", ip) + const dailyKey = buildRateLimitKey("ip", ip, dailyInterval) + let isNew = false return { check: async () => { - const rows = await Database.use((tx) => - tx - .select({ interval: IpRateLimitTable.interval, count: IpRateLimitTable.count }) - .from(IpRateLimitTable) - .where( - and( - eq(IpRateLimitTable.ip, ip), - isDefaultModel - ? inArray(IpRateLimitTable.interval, [lifetimeInterval, dailyInterval]) - : inArray(IpRateLimitTable.interval, [dailyInterval]), + const [counts, rows] = await Promise.all([ + redis.mget<(string | number | null)[]>(isDefaultModel ? [lifetimeKey, dailyKey] : [dailyKey]).catch(() => []), + Database.use((tx) => + tx + .select({ interval: IpRateLimitTable.interval, count: IpRateLimitTable.count }) + .from(IpRateLimitTable) + .where( + and( + eq(IpRateLimitTable.ip, ip), + isDefaultModel + ? inArray(IpRateLimitTable.interval, [lifetimeInterval, dailyInterval]) + : inArray(IpRateLimitTable.interval, [dailyInterval]), + ), ), - ), - ) - const lifetimeCount = rows.find((r) => r.interval === lifetimeInterval)?.count ?? 0 - const dailyCount = rows.find((r) => r.interval === dailyInterval)?.count ?? 0 + ), + ]) + const redisLifetimeCount = isDefaultModel ? Number(counts[0] ?? 0) : 0 + const redisDailyCount = Number(counts[isDefaultModel ? 1 : 0] ?? 0) + const databaseLifetimeCount = rows.find((r) => r.interval === lifetimeInterval)?.count ?? 0 + const databaseDailyCount = rows.find((r) => r.interval === dailyInterval)?.count ?? 0 + const lifetimeCount = Math.max(redisLifetimeCount, databaseLifetimeCount) + const dailyCount = Math.max(redisDailyCount, databaseDailyCount) logger.debug(`rate limit lifetime: ${lifetimeCount}, daily: ${dailyCount}`) - _isNew = isDefaultModel && lifetimeCount < dailyLimit * 7 + isNew = isDefaultModel && lifetimeCount < dailyLimit * 7 + if (isDefaultModel && databaseLifetimeCount > redisLifetimeCount) + await redis.set(lifetimeKey, databaseLifetimeCount).catch(() => {}) - if ((_isNew && dailyCount >= dailyLimit * 2) || (!_isNew && dailyCount >= dailyLimit)) - throw new FreeUsageLimitError(dict["zen.api.error.rateLimitExceeded"], getRetryAfterDay(now)) + if ((isNew && dailyCount >= dailyLimit * 2) || (!isNew && dailyCount >= dailyLimit)) + throw new FreeUsageLimitError(dict["zen.api.error.rateLimitExceeded"], retryAfter) }, track: async () => { - await Database.use((tx) => - tx - .insert(IpRateLimitTable) - .values([ - { ip, interval: dailyInterval, count: 1 }, - ...(_isNew ? [{ ip, interval: lifetimeInterval, count: 1 }] : []), - ]) - .onDuplicateKeyUpdate({ set: { count: sql`${IpRateLimitTable.count} + 1` } }), - ) + const pipeline = redis.pipeline() + pipeline.incr(dailyKey) + pipeline.expire(dailyKey, retryAfter) + if (isNew) pipeline.incr(lifetimeKey) + await Promise.all([ + pipeline.exec().catch(() => {}), + Database.use((tx) => + tx + .insert(IpRateLimitTable) + .values([ + { ip, interval: dailyInterval, count: 1 }, + ...(isNew ? [{ ip, interval: lifetimeInterval, count: 1 }] : []), + ]) + .onDuplicateKeyUpdate({ set: { count: sql`${IpRateLimitTable.count} + 1` } }), + ), + ]) }, } } diff --git a/packages/console/app/src/routes/zen/util/redis.ts b/packages/console/app/src/routes/zen/util/redis.ts new file mode 100644 index 0000000000..512523298a --- /dev/null +++ b/packages/console/app/src/routes/zen/util/redis.ts @@ -0,0 +1,18 @@ +import { Resource } from "@opencode-ai/console-resource" +import { Redis } from "@upstash/redis/cloudflare" + +let redis: Redis | undefined + +export function getRedis() { + if (redis) return redis + redis = new Redis({ + url: Resource.UpstashRedisRestUrl.value, + token: Resource.UpstashRedisRestToken.value, + enableTelemetry: false, + }) + return redis +} + +export function buildRateLimitKey(kind: string, identifier: string, interval?: string) { + return `${Resource.App.stage}:ratelimit:${kind}:${identifier}${interval ? `:${interval}` : ""}` +} diff --git a/sst-env.d.ts b/sst-env.d.ts index 1eaebd1e59..aa79ec87d3 100644 --- a/sst-env.d.ts +++ b/sst-env.d.ts @@ -137,6 +137,14 @@ declare module "sst" { "type": "sst.cloudflare.SolidStart" "url": string } + "UpstashRedisRestToken": { + "type": "sst.sst.Secret" + "value": string + } + "UpstashRedisRestUrl": { + "type": "sst.sst.Secret" + "value": string + } "Web": { "type": "sst.cloudflare.Astro" "url": string