diff --git a/.env.example b/.env.example index 2e2ed29..05908c2 100644 --- a/.env.example +++ b/.env.example @@ -49,6 +49,13 @@ INDEX_ALL_CHANNELS=false INDEXED_CHANNEL_IDS= EXCLUDED_CHANNEL_IDS= EXCLUDED_CHANNEL_NAMES= +VOID_USAGE_COOLDOWN_SECONDS=30 +VOID_USAGE_DAILY_LIMIT=100 +VOID_USAGE_UNLIMITED_USER_IDS= +VOID_USAGE_UNLIMITED_ROLE_IDS= +VOID_USAGE_BOOST_USER_IDS= +VOID_USAGE_BOOST_ROLE_IDS= +VOID_USAGE_BOOST_MULTIPLIER=10 WORKER_POLL_INTERVAL_MS=5000 OFFSITE_BACKUP_SSH_TARGET= OFFSITE_BACKUP_REMOTE_WINDOWS_DIR=C:\Meta\voidbot-backups diff --git a/README.md b/README.md index ccc75e1..1097496 100644 --- a/README.md +++ b/README.md @@ -87,6 +87,18 @@ If you want the remote Ollama generation path too, configure: If you want a private local persona instead, point `STYLE_PACK_PATH` and `SYSTEM_MESSAGES_PATH` at your own ignored local files. +Void usage limits are configurable too: + +- `VOID_USAGE_COOLDOWN_SECONDS` +- `VOID_USAGE_DAILY_LIMIT` +- `VOID_USAGE_UNLIMITED_USER_IDS` +- `VOID_USAGE_UNLIMITED_ROLE_IDS` +- `VOID_USAGE_BOOST_USER_IDS` +- `VOID_USAGE_BOOST_ROLE_IDS` +- `VOID_USAGE_BOOST_MULTIPLIER` + +The default shape is simple on purpose: everyone gets the base cooldown and daily cap, listed users or roles can be made unlimited, and listed users or roles can get a boosted lane that cuts cooldowns and multiplies daily allowance by `VOID_USAGE_BOOST_MULTIPLIER` (default `10`). + ### 2. Install dependencies ```bash diff --git a/apps/bot/src/discord-bot.ts b/apps/bot/src/discord-bot.ts index ebedc07..57e1964 100644 --- a/apps/bot/src/discord-bot.ts +++ b/apps/bot/src/discord-bot.ts @@ -26,6 +26,7 @@ import { loadStylePack, loadSystemMessageCatalog, type SystemMessageCatalog, + VoidUsageRateLimiter, } from "@voidbot/core"; import { LocalLlmProvider, @@ -147,13 +148,19 @@ export async function startBot(): Promise { jobQueue, auditLog, interactionMemory, + voidUsageRateLimits, } = await createStateStorage({ backend: config.stateStorageBackend, databaseDsn: config.databaseDsn, jobsFile: config.jobsFile, auditLogFile: config.auditLogFile, interactionMemoryFile: config.interactionMemoryFile, + rateLimitStateFile: config.rateLimitStateFile, }); + const voidUsageRateLimiter = new VoidUsageRateLimiter( + voidUsageRateLimits, + config.rateLimits, + ); const contextBuilder = new ContextBuilder(); let activeStylePack = await loadStylePack(config.stylePackPath); let activeSystemMessages = await loadSystemMessageCatalog(config.systemMessagesPath); @@ -339,6 +346,7 @@ export async function startBot(): Promise { prompt, command: "ask", actor: buildActorFromMessage(message), + roleIds: getRoleIdsFromMessage(message), guildContext: buildGuildContextFromMessage(message), outputChannelId: message.channelId, requestMessageId: message.id, @@ -354,6 +362,7 @@ export async function startBot(): Promise { jobQueue, auditLog, interactionMemory, + voidUsageRateLimiter, providerRegistry, stylePack: activeStylePack, systemMessages: activeSystemMessages, @@ -406,6 +415,7 @@ export async function startBot(): Promise { try { const actor = buildActorFromInteraction(interaction); + const roleIds = getRoleIdsFromInteraction(interaction); const guildContext = buildGuildContextFromInteraction(interaction); switch (interaction.commandName) { @@ -417,6 +427,7 @@ export async function startBot(): Promise { prompt: interaction.options.getString("question", true), command: "ask", actor, + roleIds, guildContext, outputChannelId: interaction.channelId, requestMessageId: interaction.id, @@ -432,6 +443,7 @@ export async function startBot(): Promise { jobQueue, auditLog, interactionMemory, + voidUsageRateLimiter, providerRegistry, stylePack: activeStylePack, systemMessages: activeSystemMessages, @@ -443,6 +455,7 @@ export async function startBot(): Promise { prompt: interaction.options.getString("prompt", true), command: "queue-codex", actor, + roleIds, guildContext, outputChannelId: interaction.channelId, requestMessageId: interaction.id, @@ -458,6 +471,7 @@ export async function startBot(): Promise { jobQueue, auditLog, interactionMemory, + voidUsageRateLimiter, providerRegistry, stylePack: activeStylePack, systemMessages: activeSystemMessages, @@ -562,6 +576,7 @@ interface PromptHandlerOptions { prompt: string; command: CommandName; actor: Actor; + roleIds: string[]; guildContext: GuildContext; outputChannelId: string; requestMessageId?: string; @@ -575,6 +590,7 @@ interface PromptHandlerOptions { jobQueue: JobQueue; auditLog: AuditLog; interactionMemory: InteractionMemoryBank; + voidUsageRateLimiter: VoidUsageRateLimiter; providerRegistry: ProviderRegistry; stylePack?: StylePack; systemMessages: SystemMessageCatalog; @@ -583,19 +599,6 @@ interface PromptHandlerOptions { } async function handlePrompt(options: PromptHandlerOptions): Promise { - const rememberedInteraction = await options.interactionMemory.recordInteraction({ - actorId: options.actor.id, - actorName: options.actor.displayName, - sourceKind: "direct_prompt", - guildId: options.guildContext.guildId, - channelId: options.guildContext.channelId, - channelName: options.guildContext.channelName, - command: options.command, - prompt: options.prompt, - eventId: options.requestMessageId, - }); - const interactionMemory = - rememberedInteraction.totalInteractions > 0 ? rememberedInteraction : undefined; const providerName = options.forceProvider ?? pickProvider(options.actor, options.guildContext, options.providerRegistry); @@ -628,6 +631,50 @@ async function handlePrompt(options: PromptHandlerOptions): Promise { return; } + const rateLimitDecision = await options.voidUsageRateLimiter.consume({ + actorId: options.actor.id, + roleIds: options.roleIds, + command: options.command, + provider: providerName, + guildId: options.guildContext.guildId, + channelId: options.guildContext.channelId, + }); + + if (!rateLimitDecision.allowed) { + await options.auditLog.record({ + type: "rate_limit.denied", + actorId: options.actor.id, + provider: providerName, + details: { + command: options.command, + reason: rateLimitDecision.reason, + dailyCount: rateLimitDecision.dailyCount, + dailyLimit: rateLimitDecision.policy.dailyLimit ?? null, + cooldownSeconds: rateLimitDecision.policy.cooldownSeconds ?? null, + retryAfterSeconds: rateLimitDecision.retryAfterSeconds ?? null, + resetsAt: rateLimitDecision.resetsAt ?? null, + modifier: rateLimitDecision.policy.modifier, + matchedSubjects: rateLimitDecision.policy.matchedSubjects, + }, + }); + await options.respond(renderRateLimitMessage(options.systemMessages, rateLimitDecision)); + return; + } + + const rememberedInteraction = await options.interactionMemory.recordInteraction({ + actorId: options.actor.id, + actorName: options.actor.displayName, + sourceKind: "direct_prompt", + guildId: options.guildContext.guildId, + channelId: options.guildContext.channelId, + channelName: options.guildContext.channelName, + command: options.command, + prompt: options.prompt, + eventId: options.requestMessageId, + }); + const interactionMemory = + rememberedInteraction.totalInteractions > 0 ? rememberedInteraction : undefined; + const recentMessages = await getRecentMessages(options.channel, 10); const sourceGrounding = inferSourceGroundingHint( options.prompt, @@ -966,6 +1013,32 @@ function buildActorFromInteraction(interaction: ChatInputCommandInteraction, +): string[] { + const member = interaction.member; + + if (!member || !("roles" in member)) { + return []; + } + + const { roles } = member; + + if (Array.isArray(roles)) { + return roles.filter((roleId): roleId is string => typeof roleId === "string"); + } + + if (roles && typeof roles === "object" && "cache" in roles) { + return [...roles.cache.keys()]; + } + + return []; +} + function buildGuildContextFromMessage(message: Message): GuildContext { return { guildId: message.guildId ?? undefined, @@ -1347,6 +1420,22 @@ function renderSystemMessage( return systemMessages.render(key, variables); } +function renderRateLimitMessage( + systemMessages: SystemMessageCatalog, + decision: Awaited>, +): string { + if (decision.reason === "daily_limit") { + return renderSystemMessage(systemMessages, "rate_limit.daily_limit", { + dailyLimit: decision.policy.dailyLimit ?? decision.dailyCount, + resetsAt: decision.resetsAt ? formatUtcTimestamp(decision.resetsAt) : "the next UTC midnight", + }); + } + + return renderSystemMessage(systemMessages, "rate_limit.cooldown", { + retryAfter: formatDurationSeconds(decision.retryAfterSeconds ?? 1), + }); +} + function formatHistoryResults( results: Awaited>, ): Array> { @@ -1436,6 +1525,40 @@ function truncate(input: string, limit: number): string { return `${input.slice(0, limit - 3)}...`; } +function formatDurationSeconds(totalSeconds: number): string { + const seconds = Math.max(1, Math.ceil(totalSeconds)); + + if (seconds < 60) { + return `${seconds} second${seconds === 1 ? "" : "s"}`; + } + + const minutes = Math.floor(seconds / 60); + const remainderSeconds = seconds % 60; + + if (minutes < 60) { + return remainderSeconds === 0 + ? `${minutes} minute${minutes === 1 ? "" : "s"}` + : `${minutes}m ${remainderSeconds}s`; + } + + const hours = Math.floor(minutes / 60); + const remainderMinutes = minutes % 60; + + return remainderMinutes === 0 + ? `${hours} hour${hours === 1 ? "" : "s"}` + : `${hours}h ${remainderMinutes}m`; +} + +function formatUtcTimestamp(timestamp: string): string { + const date = new Date(timestamp); + const year = date.getUTCFullYear(); + const month = `${date.getUTCMonth() + 1}`.padStart(2, "0"); + const day = `${date.getUTCDate()}`.padStart(2, "0"); + const hours = `${date.getUTCHours()}`.padStart(2, "0"); + const minutes = `${date.getUTCMinutes()}`.padStart(2, "0"); + return `${year}-${month}-${day} ${hours}:${minutes} UTC`; +} + function canSendMessages( channel: Awaited>, ): channel is TextBasedChannel & { send: (content: string) => Promise } { diff --git a/apps/worker/src/index.ts b/apps/worker/src/index.ts index d0b4da7..fc2e4f1 100644 --- a/apps/worker/src/index.ts +++ b/apps/worker/src/index.ts @@ -89,6 +89,7 @@ async function main(): Promise { jobsFile: config.jobsFile, auditLogFile: config.auditLogFile, interactionMemoryFile: config.interactionMemoryFile, + rateLimitStateFile: config.rateLimitStateFile, }); jobQueue = stateJobQueue; auditLog = stateAuditLog; diff --git a/config/system-messages.json b/config/system-messages.json index 93ecb1f..7cd9004 100644 --- a/config/system-messages.json +++ b/config/system-messages.json @@ -39,6 +39,14 @@ "That provider is not available to you.", "Access denied for that provider." ], + "rate_limit.cooldown": [ + "Easy. Wait {retryAfter} before asking again.", + "Not yet. Give it {retryAfter} before the next request." + ], + "rate_limit.daily_limit": [ + "You hit today's cap of {dailyLimit} requests. It resets at {resetsAt}.", + "Today's limit is spent: {dailyLimit} requests. Try again after {resetsAt}." + ], "permission.owner_codex.denied": [ "The owner Codex provider is restricted to the configured owner account.", "That lane is reserved for the configured owner account." diff --git a/packages/config/src/index.ts b/packages/config/src/index.ts index 5a4e5a1..d65b5fd 100644 --- a/packages/config/src/index.ts +++ b/packages/config/src/index.ts @@ -94,6 +94,13 @@ const envSchema = z.object({ INDEXED_CHANNEL_IDS: z.string().default(""), EXCLUDED_CHANNEL_IDS: z.string().default(""), EXCLUDED_CHANNEL_NAMES: z.string().default(""), + VOID_USAGE_COOLDOWN_SECONDS: z.coerce.number().int().nonnegative().default(30), + VOID_USAGE_DAILY_LIMIT: z.coerce.number().int().nonnegative().default(100), + VOID_USAGE_UNLIMITED_USER_IDS: z.string().default(""), + VOID_USAGE_UNLIMITED_ROLE_IDS: z.string().default(""), + VOID_USAGE_BOOST_USER_IDS: z.string().default(""), + VOID_USAGE_BOOST_ROLE_IDS: z.string().default(""), + VOID_USAGE_BOOST_MULTIPLIER: z.coerce.number().int().positive().default(10), WORKER_POLL_INTERVAL_MS: z.coerce.number().int().positive().default(5000), }); @@ -124,6 +131,7 @@ export interface AppConfig { auditLogFile: string; interactionMemoryFile: string; artifactsDir: string; + rateLimitStateFile: string; ragArchivePath: string; ragSourceArchivePath: string; ragImportStatePath: string; @@ -159,6 +167,15 @@ export interface AppConfig { historyCollection: string; sourceCollection: string; }; + rateLimits: { + defaultCooldownSeconds: number; + defaultDailyLimit: number; + unlimitedUserIds: string[]; + unlimitedRoleIds: string[]; + boostedUserIds: string[]; + boostedRoleIds: string[]; + boostedMultiplier: number; + }; } function parseList(value: string): string[] { @@ -287,6 +304,7 @@ export function loadConfig(env: NodeJS.ProcessEnv = process.env): AppConfig { jobsFile: resolve(storageRoot, "jobs", "jobs.json"), auditLogFile: resolve(storageRoot, "audit", "events.jsonl"), interactionMemoryFile: resolve(storageRoot, "memory", "interactions.json"), + rateLimitStateFile: resolve(storageRoot, "limits", "usage.json"), artifactsDir: resolve(storageRoot, "artifacts"), ragArchivePath: resolve(parsed.RAG_ARCHIVE_PATH), ragSourceArchivePath: resolve(parsed.RAG_SOURCE_ARCHIVE_PATH), @@ -323,5 +341,14 @@ export function loadConfig(env: NodeJS.ProcessEnv = process.env): AppConfig { historyCollection: parsed.QDRANT_HISTORY_COLLECTION, sourceCollection: parsed.QDRANT_SOURCE_COLLECTION, }, + rateLimits: { + defaultCooldownSeconds: parsed.VOID_USAGE_COOLDOWN_SECONDS, + defaultDailyLimit: parsed.VOID_USAGE_DAILY_LIMIT, + unlimitedUserIds: parseList(parsed.VOID_USAGE_UNLIMITED_USER_IDS), + unlimitedRoleIds: parseList(parsed.VOID_USAGE_UNLIMITED_ROLE_IDS), + boostedUserIds: parseList(parsed.VOID_USAGE_BOOST_USER_IDS), + boostedRoleIds: parseList(parsed.VOID_USAGE_BOOST_ROLE_IDS), + boostedMultiplier: parsed.VOID_USAGE_BOOST_MULTIPLIER, + }, }; } diff --git a/packages/core/sql/bootstrap.sql b/packages/core/sql/bootstrap.sql index 28a700c..a1d0e59 100644 --- a/packages/core/sql/bootstrap.sql +++ b/packages/core/sql/bootstrap.sql @@ -108,3 +108,15 @@ create table if not exists interaction_memory_events ( create index if not exists interaction_memory_events_actor_timestamp_idx on interaction_memory_events (actor_id, event_timestamp desc); + +create table if not exists void_usage_rate_limit_state ( + actor_id text primary key, + daily_bucket date not null, + daily_count integer not null default 0, + last_request_at timestamptz, + updated_at timestamptz not null default now(), + state_json jsonb not null default '{}'::jsonb +); + +create index if not exists void_usage_rate_limit_state_bucket_idx + on void_usage_rate_limit_state (daily_bucket, updated_at desc); diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index fb10cb4..563f8fc 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -7,3 +7,4 @@ export * from "./permission-engine"; export * from "./state-storage"; export * from "./style-pack-loader"; export * from "./system-message-catalog"; +export * from "./void-usage-rate-limiter"; diff --git a/packages/core/src/state-storage.ts b/packages/core/src/state-storage.ts index 118ca65..d7dee1e 100644 --- a/packages/core/src/state-storage.ts +++ b/packages/core/src/state-storage.ts @@ -1,6 +1,6 @@ import { randomUUID } from "node:crypto"; -import { readFile } from "node:fs/promises"; -import { resolve } from "node:path"; +import { mkdir, readFile, writeFile } from "node:fs/promises"; +import { dirname, resolve } from "node:path"; import { Pool, type PoolClient } from "pg"; @@ -28,6 +28,11 @@ import { type JobQueue, FileBackedJobQueue, } from "./job-queue"; +import { + type ConsumeVoidUsageRateLimitInput, + type VoidUsageRateLimitDecision, + type VoidUsageRateLimitStore, +} from "./void-usage-rate-limiter"; const LEGACY_FILE_STATE_IMPORT = "legacy_file_state_import_v1"; @@ -37,6 +42,7 @@ export interface StateStorageConfig { jobsFile: string; auditLogFile: string; interactionMemoryFile: string; + rateLimitStateFile: string; } export interface StateStorage { @@ -44,6 +50,7 @@ export interface StateStorage { jobQueue: JobQueue; auditLog: AuditLog; interactionMemory: InteractionMemoryBank; + voidUsageRateLimits: VoidUsageRateLimitStore; close(): Promise; } @@ -56,6 +63,7 @@ export async function createStateStorage( jobQueue: new FileBackedJobQueue(config.jobsFile), auditLog: new FileAuditLog(config.auditLogFile), interactionMemory: new FileInteractionMemoryBank(config.interactionMemoryFile), + voidUsageRateLimits: new FileBackedVoidUsageRateLimitStore(config.rateLimitStateFile), close: async () => undefined, }; } @@ -77,6 +85,7 @@ export async function createStateStorage( jobQueue: new PostgresJobQueue(pool), auditLog: new PostgresAuditLog(pool), interactionMemory: new PostgresInteractionMemoryBank(pool), + voidUsageRateLimits: new PostgresVoidUsageRateLimitStore(pool), close: async () => { await pool.end(); }, @@ -495,6 +504,247 @@ class PostgresInteractionMemoryBank implements InteractionMemoryBank { } } +interface VoidUsageRateLimitStateRow { + actor_id: string; + daily_bucket: string; + daily_count: number; + last_request_at: string | null; +} + +class PostgresVoidUsageRateLimitStore implements VoidUsageRateLimitStore { + public constructor(private readonly pool: Pool) {} + + public async consume( + input: ConsumeVoidUsageRateLimitInput, + ): Promise { + const timestamp = input.timestamp ?? new Date().toISOString(); + const now = new Date(timestamp); + const bucket = timestamp.slice(0, 10); + const client = await this.pool.connect(); + + try { + await client.query("begin"); + await client.query( + `insert into void_usage_rate_limit_state ( + actor_id, + daily_bucket, + daily_count, + state_json + ) values ($1, $2, 0, '{}'::jsonb) + on conflict (actor_id) do nothing`, + [input.actorId, bucket], + ); + const stateResult = await client.query( + `select actor_id, daily_bucket::text, daily_count, last_request_at::text + from void_usage_rate_limit_state + where actor_id = $1 + for update`, + [input.actorId], + ); + const state = stateResult.rows[0]; + + if (!state) { + throw new Error(`Missing rate limit state for actor ${input.actorId}.`); + } + + const dailyCount = state.daily_bucket === bucket ? state.daily_count : 0; + + if (input.dailyLimit !== undefined && dailyCount >= input.dailyLimit) { + await client.query("commit"); + return { + allowed: false, + reason: "daily_limit", + dailyCount, + resetsAt: nextUtcDay(timestamp), + policy: { + modifier: "default", + matchedSubjects: [], + }, + }; + } + + if (input.cooldownSeconds !== undefined && state.last_request_at) { + const lastRequestAt = new Date(state.last_request_at); + const elapsedSeconds = Math.floor((now.getTime() - lastRequestAt.getTime()) / 1000); + const retryAfterSeconds = input.cooldownSeconds - elapsedSeconds; + + if (retryAfterSeconds > 0) { + await client.query("commit"); + return { + allowed: false, + reason: "cooldown", + dailyCount, + retryAfterSeconds, + resetsAt: now.toISOString(), + policy: { + modifier: "default", + matchedSubjects: [], + }, + }; + } + } + + const nextDailyCount = dailyCount + 1; + await client.query( + `update void_usage_rate_limit_state + set daily_bucket = $2, + daily_count = $3, + last_request_at = $4, + updated_at = now(), + state_json = $5::jsonb + where actor_id = $1`, + [ + input.actorId, + bucket, + nextDailyCount, + timestamp, + JSON.stringify({ + lastCommand: input.command, + lastProvider: input.provider, + guildId: input.guildId ?? null, + channelId: input.channelId, + }), + ], + ); + await client.query("commit"); + + return { + allowed: true, + dailyCount: nextDailyCount, + resetsAt: + input.dailyLimit === undefined ? undefined : nextUtcDay(timestamp), + policy: { + modifier: "default", + matchedSubjects: [], + }, + }; + } catch (error) { + await client.query("rollback").catch(() => undefined); + throw error; + } finally { + client.release(); + } + } +} + +interface VoidUsageRateLimitFileStore { + version: 1; + actors: Record< + string, + { + dailyBucket: string; + dailyCount: number; + lastRequestAt?: string; + } + >; +} + +class FileBackedVoidUsageRateLimitStore implements VoidUsageRateLimitStore { + private writeChain: Promise = Promise.resolve(); + + public constructor(private readonly filePath: string) {} + + public async consume( + input: ConsumeVoidUsageRateLimitInput, + ): Promise { + return this.serialize(async () => { + const timestamp = input.timestamp ?? new Date().toISOString(); + const now = new Date(timestamp); + const bucket = timestamp.slice(0, 10); + const store = await this.readUnlocked(); + const existing = store.actors[input.actorId] ?? { + dailyBucket: bucket, + dailyCount: 0, + }; + const dailyCount = existing.dailyBucket === bucket ? existing.dailyCount : 0; + + if (input.dailyLimit !== undefined && dailyCount >= input.dailyLimit) { + return { + allowed: false, + reason: "daily_limit", + dailyCount, + resetsAt: nextUtcDay(timestamp), + policy: { + modifier: "default", + matchedSubjects: [], + }, + }; + } + + if (input.cooldownSeconds !== undefined && existing.lastRequestAt) { + const lastRequestAt = new Date(existing.lastRequestAt); + const elapsedSeconds = Math.floor((now.getTime() - lastRequestAt.getTime()) / 1000); + const retryAfterSeconds = input.cooldownSeconds - elapsedSeconds; + + if (retryAfterSeconds > 0) { + return { + allowed: false, + reason: "cooldown", + dailyCount, + retryAfterSeconds, + resetsAt: now.toISOString(), + policy: { + modifier: "default", + matchedSubjects: [], + }, + }; + } + } + + const nextDailyCount = dailyCount + 1; + store.actors[input.actorId] = { + dailyBucket: bucket, + dailyCount: nextDailyCount, + lastRequestAt: timestamp, + }; + await this.writeUnlocked(store); + + return { + allowed: true, + dailyCount: nextDailyCount, + resetsAt: + input.dailyLimit === undefined ? undefined : nextUtcDay(timestamp), + policy: { + modifier: "default", + matchedSubjects: [], + }, + }; + }); + } + + private async serialize(operation: () => Promise): Promise { + const pending = this.writeChain.then(operation, operation); + this.writeChain = pending.then( + () => undefined, + () => undefined, + ); + return pending; + } + + private async readUnlocked(): Promise { + try { + const raw = await readFile(this.filePath, "utf8"); + return JSON.parse(raw) as VoidUsageRateLimitFileStore; + } catch (error) { + const code = (error as NodeJS.ErrnoException).code; + + if (code === "ENOENT") { + return { + version: 1, + actors: {}, + }; + } + + throw error; + } + } + + private async writeUnlocked(store: VoidUsageRateLimitFileStore): Promise { + await mkdir(dirname(this.filePath), { recursive: true }); + await writeFile(this.filePath, `${JSON.stringify(store, null, 2)}\n`, "utf8"); + } +} + async function bootstrapPostgres(pool: Pool): Promise { const sqlPath = resolve(process.cwd(), "packages", "core", "sql", "bootstrap.sql"); const sql = await readFile(sqlPath, "utf8"); @@ -772,3 +1022,11 @@ function dedupeLegacyJobs(jobs: JobRecord[]): JobRecord[] { left.createdAt.localeCompare(right.createdAt), ); } + +function nextUtcDay(timestamp: string): string { + const date = new Date(timestamp); + const next = new Date( + Date.UTC(date.getUTCFullYear(), date.getUTCMonth(), date.getUTCDate() + 1, 0, 0, 0, 0), + ); + return next.toISOString(); +} diff --git a/packages/core/src/void-usage-rate-limiter.ts b/packages/core/src/void-usage-rate-limiter.ts new file mode 100644 index 0000000..050ffbc --- /dev/null +++ b/packages/core/src/void-usage-rate-limiter.ts @@ -0,0 +1,169 @@ +import { type CommandName, type ProviderName } from "@voidbot/shared"; + +export interface VoidUsageRateLimitConfig { + defaultCooldownSeconds: number; + defaultDailyLimit: number; + unlimitedUserIds: string[]; + unlimitedRoleIds: string[]; + boostedUserIds: string[]; + boostedRoleIds: string[]; + boostedMultiplier: number; +} + +export interface VoidUsageSubject { + actorId: string; + roleIds: string[]; +} + +export interface EffectiveVoidUsagePolicy { + cooldownSeconds?: number; + dailyLimit?: number; + modifier: "default" | "boosted" | "unlimited"; + matchedSubjects: string[]; +} + +export interface VoidUsageRateLimitDecision { + allowed: boolean; + reason?: "cooldown" | "daily_limit"; + policy: EffectiveVoidUsagePolicy; + dailyCount: number; + retryAfterSeconds?: number; + resetsAt?: string; +} + +export interface ConsumeVoidUsageRateLimitInput { + actorId: string; + command: CommandName; + provider: ProviderName; + guildId?: string; + channelId: string; + cooldownSeconds?: number; + dailyLimit?: number; + timestamp?: string; +} + +export interface VoidUsageRateLimitStore { + consume(input: ConsumeVoidUsageRateLimitInput): Promise; +} + +export class VoidUsageRateLimiter { + public constructor( + private readonly store: VoidUsageRateLimitStore, + private readonly config: VoidUsageRateLimitConfig, + ) {} + + public async consume(input: { + actorId: string; + roleIds: string[]; + command: CommandName; + provider: ProviderName; + guildId?: string; + channelId: string; + timestamp?: string; + }): Promise { + const policy = resolveEffectivePolicy(this.config, { + actorId: input.actorId, + roleIds: input.roleIds, + }); + + if (policy.cooldownSeconds === undefined && policy.dailyLimit === undefined) { + return { + allowed: true, + policy, + dailyCount: 0, + }; + } + + const decision = await this.store.consume({ + actorId: input.actorId, + command: input.command, + provider: input.provider, + guildId: input.guildId, + channelId: input.channelId, + cooldownSeconds: policy.cooldownSeconds, + dailyLimit: policy.dailyLimit, + timestamp: input.timestamp, + }); + + return { + ...decision, + policy, + }; + } +} + +export function resolveEffectivePolicy( + config: VoidUsageRateLimitConfig, + subject: VoidUsageSubject, +): EffectiveVoidUsagePolicy { + const baseCooldownSeconds = normalizeLimitValue(config.defaultCooldownSeconds); + const baseDailyLimit = normalizeLimitValue(config.defaultDailyLimit); + const unlimitedMatches = collectMatches(subject, { + userIds: config.unlimitedUserIds, + roleIds: config.unlimitedRoleIds, + }); + + if (unlimitedMatches.length > 0) { + return { + modifier: "unlimited", + matchedSubjects: unlimitedMatches, + }; + } + + const boostedMatches = collectMatches(subject, { + userIds: config.boostedUserIds, + roleIds: config.boostedRoleIds, + }); + + if (boostedMatches.length === 0) { + return { + cooldownSeconds: baseCooldownSeconds, + dailyLimit: baseDailyLimit, + modifier: "default", + matchedSubjects: [], + }; + } + + const multiplier = config.boostedMultiplier > 1 ? config.boostedMultiplier : 10; + + return { + cooldownSeconds: + baseCooldownSeconds === undefined + ? undefined + : Math.max(1, Math.floor(baseCooldownSeconds / multiplier)), + dailyLimit: + baseDailyLimit === undefined ? undefined : Math.max(1, baseDailyLimit * multiplier), + modifier: "boosted", + matchedSubjects: boostedMatches, + }; +} + +function normalizeLimitValue(value: number | undefined): number | undefined { + if (value === undefined || !Number.isFinite(value) || value <= 0) { + return undefined; + } + + return Math.floor(value); +} + +function collectMatches( + subject: VoidUsageSubject, + targets: { + userIds: string[]; + roleIds: string[]; + }, +): string[] { + const matches: string[] = []; + + if (targets.userIds.includes(subject.actorId)) { + matches.push(`user:${subject.actorId}`); + } + + for (const roleId of subject.roleIds) { + if (targets.roleIds.includes(roleId)) { + matches.push(`role:${roleId}`); + } + } + + return matches; +}