Add configurable Void usage rate limits
This commit is contained in:
Meta Vens
parent
6c17968fc3
commit
e53e24261e
10 changed files with 633 additions and 15 deletions
|
|
@ -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
|
||||
|
|
|
|||
12
README.md
12
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
|
||||
|
|
|
|||
|
|
@ -26,6 +26,7 @@ import {
|
|||
loadStylePack,
|
||||
loadSystemMessageCatalog,
|
||||
type SystemMessageCatalog,
|
||||
VoidUsageRateLimiter,
|
||||
} from "@voidbot/core";
|
||||
import {
|
||||
LocalLlmProvider,
|
||||
|
|
@ -147,13 +148,19 @@ export async function startBot(): Promise<void> {
|
|||
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<void> {
|
|||
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<void> {
|
|||
jobQueue,
|
||||
auditLog,
|
||||
interactionMemory,
|
||||
voidUsageRateLimiter,
|
||||
providerRegistry,
|
||||
stylePack: activeStylePack,
|
||||
systemMessages: activeSystemMessages,
|
||||
|
|
@ -406,6 +415,7 @@ export async function startBot(): Promise<void> {
|
|||
|
||||
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<void> {
|
|||
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<void> {
|
|||
jobQueue,
|
||||
auditLog,
|
||||
interactionMemory,
|
||||
voidUsageRateLimiter,
|
||||
providerRegistry,
|
||||
stylePack: activeStylePack,
|
||||
systemMessages: activeSystemMessages,
|
||||
|
|
@ -443,6 +455,7 @@ export async function startBot(): Promise<void> {
|
|||
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<void> {
|
|||
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<void> {
|
||||
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<void> {
|
|||
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<Cach
|
|||
};
|
||||
}
|
||||
|
||||
function getRoleIdsFromMessage(message: Message): string[] {
|
||||
return message.member ? [...message.member.roles.cache.keys()] : [];
|
||||
}
|
||||
|
||||
function getRoleIdsFromInteraction(
|
||||
interaction: ChatInputCommandInteraction<CacheType>,
|
||||
): 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<ReturnType<VoidUsageRateLimiter["consume"]>>,
|
||||
): 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<ReturnType<RetrievalService["searchHistory"]>>,
|
||||
): Array<Record<string, unknown>> {
|
||||
|
|
@ -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<ReturnType<Client["channels"]["fetch"]>>,
|
||||
): channel is TextBasedChannel & { send: (content: string) => Promise<unknown> } {
|
||||
|
|
|
|||
|
|
@ -89,6 +89,7 @@ async function main(): Promise<void> {
|
|||
jobsFile: config.jobsFile,
|
||||
auditLogFile: config.auditLogFile,
|
||||
interactionMemoryFile: config.interactionMemoryFile,
|
||||
rateLimitStateFile: config.rateLimitStateFile,
|
||||
});
|
||||
jobQueue = stateJobQueue;
|
||||
auditLog = stateAuditLog;
|
||||
|
|
|
|||
|
|
@ -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."
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
},
|
||||
};
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
|
|
|||
|
|
@ -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";
|
||||
|
|
|
|||
|
|
@ -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<void>;
|
||||
}
|
||||
|
||||
|
|
@ -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<VoidUsageRateLimitDecision> {
|
||||
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<VoidUsageRateLimitStateRow>(
|
||||
`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<void> = Promise.resolve();
|
||||
|
||||
public constructor(private readonly filePath: string) {}
|
||||
|
||||
public async consume(
|
||||
input: ConsumeVoidUsageRateLimitInput,
|
||||
): Promise<VoidUsageRateLimitDecision> {
|
||||
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<T>(operation: () => Promise<T>): Promise<T> {
|
||||
const pending = this.writeChain.then(operation, operation);
|
||||
this.writeChain = pending.then(
|
||||
() => undefined,
|
||||
() => undefined,
|
||||
);
|
||||
return pending;
|
||||
}
|
||||
|
||||
private async readUnlocked(): Promise<VoidUsageRateLimitFileStore> {
|
||||
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<void> {
|
||||
await mkdir(dirname(this.filePath), { recursive: true });
|
||||
await writeFile(this.filePath, `${JSON.stringify(store, null, 2)}\n`, "utf8");
|
||||
}
|
||||
}
|
||||
|
||||
async function bootstrapPostgres(pool: Pool): Promise<void> {
|
||||
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();
|
||||
}
|
||||
|
|
|
|||
169
packages/core/src/void-usage-rate-limiter.ts
Normal file
169
packages/core/src/void-usage-rate-limiter.ts
Normal file
|
|
@ -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<VoidUsageRateLimitDecision>;
|
||||
}
|
||||
|
||||
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<VoidUsageRateLimitDecision> {
|
||||
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;
|
||||
}
|
||||
Loading…
Reference in a new issue