fix(hooks): add security hardening for message handlers

- Add ReDoS protection using safe-regex validation for contentPattern
- Add token bucket rate limiting (10/min per handler) to prevent cost explosions
- Require at least one match condition in Zod schema (reject empty match)
- Change error logging from logVerbose to logWarn for visibility

Includes tests for rate limiter and ReDoS protection.
This commit is contained in:
Steven Gonsalvez 2026-01-28 21:59:34 +00:00
parent 42c647f2c9
commit 761b7b6849
No known key found for this signature in database
GPG Key ID: 907EC78C72C6AFF6
10 changed files with 334 additions and 122 deletions

View File

@ -366,16 +366,18 @@ Message handlers fix this by immediately processing important messages as they a
### Match Conditions ### Match Conditions
All specified conditions must match (AND logic): All specified conditions must match (AND logic). **At least one condition is required** - empty match objects are rejected to prevent accidental catch-all handlers.
| Field | Type | Description | | Field | Type | Description |
|-------|------|-------------| |-------|------|-------------|
| `channelId` | `string \| string[]` | Channel to match: `"whatsapp"`, `"telegram"`, `["discord", "slack"]`, or `"*"` for all | | `channelId` | `string \| string[]` | Channel to match: `"whatsapp"`, `"telegram"`, `["discord", "slack"]`, or `"*"` for all |
| `conversationId` | `string \| string[]` | Chat/group ID to match | | `conversationId` | `string \| string[]` | Chat/group ID to match |
| `from` | `string \| string[]` | Sender identifier (phone number, user ID) | | `from` | `string \| string[]` | Sender identifier (phone number, user ID) |
| `contentPattern` | `string` | Regex pattern (case-insensitive) | | `contentPattern` | `string` | Regex pattern (case-insensitive). Unsafe patterns (ReDoS vulnerable) are rejected. |
| `contentContains` | `string \| string[]` | Keywords to find in message (case-insensitive, any match) | | `contentContains` | `string \| string[]` | Keywords to find in message (case-insensitive, any match) |
**Security Note**: The `contentPattern` field is validated for ReDoS (Regular Expression Denial of Service) safety before use. Patterns that could cause catastrophic backtracking (e.g., `(a+)+`) are rejected and logged as warnings.
### Handler Options ### Handler Options
| Field | Type | Default | Description | | Field | Type | Default | Description |
@ -491,12 +493,23 @@ All specified conditions must match (AND logic):
This triggers `analytics-bot` AND lets the normal message flow continue (so the user's main agent also processes it). This triggers `analytics-bot` AND lets the normal message flow continue (so the user's main agent also processes it).
### Rate Limiting
Message handlers are rate limited to prevent cost explosions from unbounded agent execution:
- **Default limit**: 10 executions per minute per handler
- Rate-limited messages fall through to normal queue processing
- A warning is logged when rate limiting kicks in
This protects against scenarios like a busy group chat triggering dozens of agent executions per minute.
### Order of Evaluation ### Order of Evaluation
1. Handlers are evaluated in order (first match wins) 1. Handlers are evaluated in order (first match wins)
2. Disabled handlers (`enabled: false`) are skipped 2. Disabled handlers (`enabled: false`) are skipped
3. If a handler matches with `mode: "exclusive"` (default), normal processing stops 3. Rate limiting is checked before execution
4. If a handler matches with `mode: "parallel"`, normal processing continues after the handler fires 4. If a handler matches with `mode: "exclusive"` (default), normal processing stops
5. If a handler matches with `mode: "parallel"`, normal processing continues after the handler fires
## Creating Custom Hooks ## Creating Custom Hooks

View File

@ -205,6 +205,7 @@
"tslog": "^4.10.2", "tslog": "^4.10.2",
"undici": "^7.19.0", "undici": "^7.19.0",
"ws": "^8.19.0", "ws": "^8.19.0",
"safe-regex": "^2.1.1",
"yaml": "^2.8.2", "yaml": "^2.8.2",
"zod": "^4.3.6" "zod": "^4.3.6"
}, },
@ -223,6 +224,7 @@
"@types/node": "^25.0.10", "@types/node": "^25.0.10",
"@types/proper-lockfile": "^4.1.4", "@types/proper-lockfile": "^4.1.4",
"@types/qrcode-terminal": "^0.12.2", "@types/qrcode-terminal": "^0.12.2",
"@types/safe-regex": "^1.1.6",
"@types/ws": "^8.18.1", "@types/ws": "^8.18.1",
"@typescript/native-preview": "7.0.0-dev.20260124.1", "@typescript/native-preview": "7.0.0-dev.20260124.1",
"@vitest/coverage-v8": "^4.0.18", "@vitest/coverage-v8": "^4.0.18",

113
pnpm-lock.yaml generated
View File

@ -148,6 +148,9 @@ importers:
qrcode-terminal: qrcode-terminal:
specifier: ^0.12.0 specifier: ^0.12.0
version: 0.12.0 version: 0.12.0
safe-regex:
specifier: ^2.1.1
version: 2.1.1
sharp: sharp:
specifier: ^0.34.5 specifier: ^0.34.5
version: 0.34.5 version: 0.34.5
@ -203,6 +206,9 @@ importers:
'@types/qrcode-terminal': '@types/qrcode-terminal':
specifier: ^0.12.2 specifier: ^0.12.2
version: 0.12.2 version: 0.12.2
'@types/safe-regex':
specifier: ^1.1.6
version: 1.1.6
'@types/ws': '@types/ws':
specifier: ^8.18.1 specifier: ^8.18.1
version: 8.18.1 version: 8.18.1
@ -383,12 +389,12 @@ importers:
'@microsoft/agents-hosting-extensions-teams': '@microsoft/agents-hosting-extensions-teams':
specifier: ^1.2.2 specifier: ^1.2.2
version: 1.2.2 version: 1.2.2
moltbot:
specifier: workspace:*
version: link:../..
express: express:
specifier: ^5.2.1 specifier: ^5.2.1
version: 5.2.1 version: 5.2.1
moltbot:
specifier: workspace:*
version: link:../..
proper-lockfile: proper-lockfile:
specifier: ^4.1.2 specifier: ^4.1.2
version: 4.1.2 version: 4.1.2
@ -2767,6 +2773,9 @@ packages:
'@types/retry@0.12.5': '@types/retry@0.12.5':
resolution: {integrity: sha512-3xSjTp3v03X/lSQLkczaN9UIEwJMoMCA1+Nb5HfbJEQWogdeQIyVtTvxPXDQjZ5zws8rFQfVfRdz03ARihPJgw==} resolution: {integrity: sha512-3xSjTp3v03X/lSQLkczaN9UIEwJMoMCA1+Nb5HfbJEQWogdeQIyVtTvxPXDQjZ5zws8rFQfVfRdz03ARihPJgw==}
'@types/safe-regex@1.1.6':
resolution: {integrity: sha512-CQ/uPB9fLOPKwDsrTeVbNIkwfUthTWOx0l6uIGwVFjZxv7e68pCW5gtTYFzdJi3EBJp8h8zYhJbTasAbX7gEMQ==}
'@types/send@0.17.6': '@types/send@0.17.6':
resolution: {integrity: sha512-Uqt8rPBE8SY0RK8JB1EzVOIZ32uqy8HwdxCnoCOsYrvnswqmFZ/k+9Ikidlk/ImhsdvBsloHbAlewb2IEBV/Og==} resolution: {integrity: sha512-Uqt8rPBE8SY0RK8JB1EzVOIZ32uqy8HwdxCnoCOsYrvnswqmFZ/k+9Ikidlk/ImhsdvBsloHbAlewb2IEBV/Og==}
@ -3214,11 +3223,6 @@ packages:
class-variance-authority@0.7.1: class-variance-authority@0.7.1:
resolution: {integrity: sha512-Ka+9Trutv7G8M6WT6SeiRWz792K5qEqIGEGzXKhAE6xOWAY6pPH8U+9IY3oCMv6kqTmLsv7Xh/2w2RigkePMsg==} resolution: {integrity: sha512-Ka+9Trutv7G8M6WT6SeiRWz792K5qEqIGEGzXKhAE6xOWAY6pPH8U+9IY3oCMv6kqTmLsv7Xh/2w2RigkePMsg==}
clawdbot@2026.1.24-3:
resolution: {integrity: sha512-zt9BzhWXduq8ZZR4rfzQDurQWAgmijTTyPZCQGrn5ew6wCEwhxxEr2/NHG7IlCwcfRsKymsY4se9KMhoNz0JtQ==}
engines: {node: '>=22.12.0'}
hasBin: true
cli-cursor@5.0.0: cli-cursor@5.0.0:
resolution: {integrity: sha512-aCj4O5wKyszjMmDT4tZj93kxyydN/K5zPWSCe6/0AV/AA1pqe5ZBIw0a2ZfPQV7lL5/yb5HsUreJ6UFAF1tEQw==} resolution: {integrity: sha512-aCj4O5wKyszjMmDT4tZj93kxyydN/K5zPWSCe6/0AV/AA1pqe5ZBIw0a2ZfPQV7lL5/yb5HsUreJ6UFAF1tEQw==}
engines: {node: '>=18'} engines: {node: '>=18'}
@ -4864,6 +4868,10 @@ packages:
reflect-metadata@0.2.2: reflect-metadata@0.2.2:
resolution: {integrity: sha512-urBwgfrvVP/eAyXx4hluJivBKzuEbSQs9rKWCrCkbSxNv8mxPcUZKeuoF3Uy4mJl3Lwprp6yy5/39VWigZ4K6Q==} resolution: {integrity: sha512-urBwgfrvVP/eAyXx4hluJivBKzuEbSQs9rKWCrCkbSxNv8mxPcUZKeuoF3Uy4mJl3Lwprp6yy5/39VWigZ4K6Q==}
regexp-tree@0.1.27:
resolution: {integrity: sha512-iETxpjK6YoRWJG5o6hXLwvjYAoW+FEZn9os0PD/b6AP6xQwsa/Y7lCVgIixBbUPMfhu+i2LtdeAqVTgGlQarfA==}
hasBin: true
request-promise-core@1.1.4: request-promise-core@1.1.4:
resolution: {integrity: sha512-TTbAfBBRdWD7aNNOoVOBH4pN/KigV6LyapYNNlAPA8JwbovRti1E88m3sYAwsLi5ryhPKsE9APwnjFTgdUjTpw==} resolution: {integrity: sha512-TTbAfBBRdWD7aNNOoVOBH4pN/KigV6LyapYNNlAPA8JwbovRti1E88m3sYAwsLi5ryhPKsE9APwnjFTgdUjTpw==}
engines: {node: '>=0.10.0'} engines: {node: '>=0.10.0'}
@ -4940,6 +4948,9 @@ packages:
safe-buffer@5.2.1: safe-buffer@5.2.1:
resolution: {integrity: sha512-rp3So07KcdmmKbGvgaNxQSJr7bGVSVk5S9Eq1F+ppbRo70+YeaDxkw5Dd8NPN+GD6bjnYm2VuPuCXmpuYvmCXQ==} resolution: {integrity: sha512-rp3So07KcdmmKbGvgaNxQSJr7bGVSVk5S9Eq1F+ppbRo70+YeaDxkw5Dd8NPN+GD6bjnYm2VuPuCXmpuYvmCXQ==}
safe-regex@2.1.1:
resolution: {integrity: sha512-rx+x8AMzKb5Q5lQ95Zoi6ZbJqwCLkqi3XuJXp5P3rT8OEc6sZCJG5AE5dU3lsgRr/F4Bs31jSlVN+j5KrsGu9A==}
safe-stable-stringify@2.5.0: safe-stable-stringify@2.5.0:
resolution: {integrity: sha512-b3rppTKm9T+PsVCBEOUR46GWI7fdOs00VKZ1+9c1EWDaDMvjQc6tUwuFyIprgGgTcWoVHSKrU8H31ZHA2e0RHA==} resolution: {integrity: sha512-b3rppTKm9T+PsVCBEOUR46GWI7fdOs00VKZ1+9c1EWDaDMvjQc6tUwuFyIprgGgTcWoVHSKrU8H31ZHA2e0RHA==}
engines: {node: '>=10'} engines: {node: '>=10'}
@ -8549,6 +8560,8 @@ snapshots:
'@types/retry@0.12.5': {} '@types/retry@0.12.5': {}
'@types/safe-regex@1.1.6': {}
'@types/send@0.17.6': '@types/send@0.17.6':
dependencies: dependencies:
'@types/mime': 1.3.5 '@types/mime': 1.3.5
@ -9098,84 +9111,6 @@ snapshots:
dependencies: dependencies:
clsx: 2.1.1 clsx: 2.1.1
clawdbot@2026.1.24-3(@types/express@5.0.6)(audio-decode@2.2.3)(devtools-protocol@0.0.1561482)(typescript@5.9.3):
dependencies:
'@agentclientprotocol/sdk': 0.13.1(zod@4.3.6)
'@aws-sdk/client-bedrock': 3.975.0
'@buape/carbon': 0.14.0(hono@4.11.4)
'@clack/prompts': 0.11.0
'@grammyjs/runner': 2.0.3(grammy@1.39.3)
'@grammyjs/transformer-throttler': 1.2.1(grammy@1.39.3)
'@homebridge/ciao': 1.3.4
'@line/bot-sdk': 10.6.0
'@lydell/node-pty': 1.2.0-beta.3
'@mariozechner/pi-agent-core': 0.49.3(ws@8.19.0)(zod@4.3.6)
'@mariozechner/pi-ai': 0.49.3(ws@8.19.0)(zod@4.3.6)
'@mariozechner/pi-coding-agent': 0.49.3(ws@8.19.0)(zod@4.3.6)
'@mariozechner/pi-tui': 0.49.3
'@mozilla/readability': 0.6.0
'@sinclair/typebox': 0.34.47
'@slack/bolt': 4.6.0(@types/express@5.0.6)
'@slack/web-api': 7.13.0
'@whiskeysockets/baileys': 7.0.0-rc.9(audio-decode@2.2.3)(sharp@0.34.5)
ajv: 8.17.1
body-parser: 2.2.2
chalk: 5.6.2
chokidar: 5.0.0
chromium-bidi: 13.0.1(devtools-protocol@0.0.1561482)
cli-highlight: 2.1.11
commander: 14.0.2
croner: 9.1.0
detect-libc: 2.1.2
discord-api-types: 0.38.37
dotenv: 17.2.3
express: 5.2.1
file-type: 21.3.0
grammy: 1.39.3
hono: 4.11.4
jiti: 2.6.1
json5: 2.2.3
jszip: 3.10.1
linkedom: 0.18.12
long: 5.3.2
markdown-it: 14.1.0
node-edge-tts: 1.2.9
osc-progress: 0.3.0
pdfjs-dist: 5.4.530
playwright-core: 1.58.0
proper-lockfile: 4.1.2
qrcode-terminal: 0.12.0
sharp: 0.34.5
sqlite-vec: 0.1.7-alpha.2
tar: 7.5.4
tslog: 4.10.2
undici: 7.19.0
ws: 8.19.0
yaml: 2.8.2
zod: 4.3.6
optionalDependencies:
'@napi-rs/canvas': 0.1.88
node-llama-cpp: 3.15.0(typescript@5.9.3)
transitivePeerDependencies:
- '@discordjs/opus'
- '@modelcontextprotocol/sdk'
- '@types/express'
- audio-decode
- aws-crt
- bufferutil
- canvas
- debug
- devtools-protocol
- encoding
- ffmpeg-static
- jimp
- link-preview-js
- node-opus
- opusscript
- supports-color
- typescript
- utf-8-validate
cli-cursor@5.0.0: cli-cursor@5.0.0:
dependencies: dependencies:
restore-cursor: 5.1.0 restore-cursor: 5.1.0
@ -11014,6 +10949,8 @@ snapshots:
reflect-metadata@0.2.2: {} reflect-metadata@0.2.2: {}
regexp-tree@0.1.27: {}
request-promise-core@1.1.4(request@2.88.2): request-promise-core@1.1.4(request@2.88.2):
dependencies: dependencies:
lodash: 4.17.23 lodash: 4.17.23
@ -11147,6 +11084,10 @@ snapshots:
safe-buffer@5.2.1: {} safe-buffer@5.2.1: {}
safe-regex@2.1.1:
dependencies:
regexp-tree: 0.1.27
safe-stable-stringify@2.5.0: {} safe-stable-stringify@2.5.0: {}
safer-buffer@2.1.2: {} safer-buffer@2.1.2: {}

View File

@ -3,10 +3,12 @@ import { createDefaultDeps } from "../../cli/deps.js";
import type { MoltbotConfig } from "../../config/config.js"; import type { MoltbotConfig } from "../../config/config.js";
import { loadSessionStore, resolveStorePath } from "../../config/sessions.js"; import { loadSessionStore, resolveStorePath } from "../../config/sessions.js";
import { logVerbose } from "../../globals.js"; import { logVerbose } from "../../globals.js";
import { logWarn } from "../../logger.js";
import { createInternalHookEvent, triggerInternalHook } from "../../hooks/hooks.js"; import { createInternalHookEvent, triggerInternalHook } from "../../hooks/hooks.js";
import { matchMessageHandler } from "../../hooks/message-handler-match.js";
import { runMessageHandler } from "../../hooks/message-handler-run.js";
import type { MessageReceivedHookContext } from "../../hooks/internal-hooks.js"; import type { MessageReceivedHookContext } from "../../hooks/internal-hooks.js";
import { matchMessageHandler } from "../../hooks/message-handler-match.js";
import { isHandlerRateLimited } from "../../hooks/message-handler-rate-limit.js";
import { runMessageHandler } from "../../hooks/message-handler-run.js";
import { isDiagnosticsEnabled } from "../../infra/diagnostic-events.js"; import { isDiagnosticsEnabled } from "../../infra/diagnostic-events.js";
import { import {
logMessageProcessed, logMessageProcessed,
@ -263,33 +265,41 @@ export async function dispatchReplyFromConfig(params: {
const matchedHandler = matchMessageHandler(messageHandlers, handlerContext); const matchedHandler = matchMessageHandler(messageHandlers, handlerContext);
if (matchedHandler) { if (matchedHandler) {
const priority = matchedHandler.priority ?? "immediate"; // Check rate limit before executing
const mode = matchedHandler.mode ?? "exclusive"; if (isHandlerRateLimited(matchedHandler.id)) {
logWarn(
`dispatch-from-config: message handler "${matchedHandler.id}" rate limited, skipping`,
);
// Don't return early for rate-limited handlers - fall through to normal queue flow
} else {
const priority = matchedHandler.priority ?? "immediate";
const mode = matchedHandler.mode ?? "exclusive";
if (priority === "immediate") { if (priority === "immediate") {
// Fire-and-forget for immediate handlers // Fire-and-forget for immediate handlers
const deps = createDefaultDeps(); const deps = createDefaultDeps();
void runMessageHandler({ void runMessageHandler({
cfg, cfg,
deps, deps,
handler: matchedHandler, handler: matchedHandler,
context: handlerContext, context: handlerContext,
}).catch((err) => { }).catch((err) => {
logVerbose( logWarn(
`dispatch-from-config: message handler "${matchedHandler.id}" failed: ${String(err)}`, `dispatch-from-config: message handler "${matchedHandler.id}" failed: ${String(err)}`,
); );
}); });
// Check mode: exclusive (default) vs parallel // Check mode: exclusive (default) vs parallel
if (mode === "exclusive") { if (mode === "exclusive") {
// Return early - handler takes over completely // Return early - handler takes over completely
return { return {
queuedFinal: false, queuedFinal: false,
counts: dispatcher.getQueuedCounts(), counts: dispatcher.getQueuedCounts(),
handledByMessageHandler: { handlerId: matchedHandler.id }, handledByMessageHandler: { handlerId: matchedHandler.id },
}; };
}
// parallel mode: continue with normal flow
} }
// parallel mode: continue with normal flow
} }
} }
} }

View File

@ -60,7 +60,23 @@ const MessageHandlerMatchSchema = z
contentPattern: z.string().optional(), contentPattern: z.string().optional(),
contentContains: z.union([z.string(), z.array(z.string())]).optional(), contentContains: z.union([z.string(), z.array(z.string())]).optional(),
}) })
.strict(); .strict()
.refine(
(match) => {
// Require at least one match condition to prevent accidental catch-all handlers
return (
match.channelId !== undefined ||
match.conversationId !== undefined ||
match.from !== undefined ||
match.contentPattern !== undefined ||
match.contentContains !== undefined
);
},
{
message:
"At least one match condition is required (channelId, conversationId, from, contentPattern, or contentContains)",
},
);
const MessageHandlerConfigSchema = z const MessageHandlerConfigSchema = z
.object({ .object({

View File

@ -162,6 +162,23 @@ describe("matchMessageHandler", () => {
const result = matchMessageHandler(handlers, ctx); const result = matchMessageHandler(handlers, ctx);
expect(result).toBeNull(); expect(result).toBeNull();
}); });
it("rejects unsafe regex patterns (ReDoS protection)", () => {
// This pattern causes catastrophic backtracking: (a+)+ on "aaaaaaaaaaaaaaaaaaaaaaaa!"
const handlers = [createHandler({ match: { contentPattern: "(a+)+" } })];
const ctx = createContext({ content: "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaa!" });
const result = matchMessageHandler(handlers, ctx);
expect(result).toBeNull(); // Unsafe pattern should be rejected
});
it("allows safe regex patterns", () => {
const handlers = [createHandler({ match: { contentPattern: "bug|error|broken" } })];
const ctx = createContext({ content: "Found a bug" });
const result = matchMessageHandler(handlers, ctx);
expect(result).toBe(handlers[0]);
});
}); });
describe("contentContains matching", () => { describe("contentContains matching", () => {
@ -304,7 +321,11 @@ describe("matchMessageHandler", () => {
}); });
describe("empty match conditions", () => { describe("empty match conditions", () => {
it("matches any message when match is empty", () => { // NOTE: At the matching logic level, empty match still matches everything.
// However, Zod config validation now requires at least one condition,
// so this scenario won't occur in practice with config-driven handlers.
// This test verifies the matching logic behavior for completeness.
it("matches any message when match is empty (prevented by config validation)", () => {
const handlers = [createHandler({ match: {} })]; const handlers = [createHandler({ match: {} })];
const ctx = createContext(); const ctx = createContext();

View File

@ -3,7 +3,9 @@
* Matches inbound messages against configured handlers based on channel, sender, content, etc. * Matches inbound messages against configured handlers based on channel, sender, content, etc.
*/ */
import safeRegex from "safe-regex";
import type { MessageHandlerConfig, MessageHandlerMatch } from "../config/types.hooks.js"; import type { MessageHandlerConfig, MessageHandlerMatch } from "../config/types.hooks.js";
import { logWarn } from "../logger.js";
import type { MessageReceivedHookContext } from "./internal-hooks.js"; import type { MessageReceivedHookContext } from "./internal-hooks.js";
/** /**
@ -53,8 +55,13 @@ function matchesConditions(match: MessageHandlerMatch, ctx: MessageReceivedHookC
return false; return false;
} }
// contentPattern (regex) // contentPattern (regex) - with ReDoS protection
if (match.contentPattern !== undefined) { if (match.contentPattern !== undefined) {
// Validate regex safety before compilation to prevent catastrophic backtracking
if (!safeRegex(match.contentPattern)) {
logWarn(`message-handler-match: unsafe regex pattern rejected: ${match.contentPattern}`);
return false;
}
try { try {
const regex = new RegExp(match.contentPattern, "i"); const regex = new RegExp(match.contentPattern, "i");
if (!regex.test(ctx.content)) { if (!regex.test(ctx.content)) {

View File

@ -0,0 +1,115 @@
import { afterEach, describe, expect, it, vi } from "vitest";
import {
getHandlerRateLimitInfo,
isHandlerRateLimited,
resetAllRateLimits,
resetHandlerRateLimit,
} from "./message-handler-rate-limit.js";
describe("message-handler-rate-limit", () => {
afterEach(() => {
resetAllRateLimits();
vi.useRealTimers();
});
describe("isHandlerRateLimited", () => {
it("allows first request", () => {
const isLimited = isHandlerRateLimited("test-handler");
expect(isLimited).toBe(false);
});
it("allows up to rate limit requests", () => {
const limit = 5;
for (let i = 0; i < limit; i++) {
expect(isHandlerRateLimited("test-handler", limit)).toBe(false);
}
});
it("blocks requests beyond rate limit", () => {
const limit = 3;
// Use up all tokens
for (let i = 0; i < limit; i++) {
expect(isHandlerRateLimited("test-handler", limit)).toBe(false);
}
// Next request should be blocked
expect(isHandlerRateLimited("test-handler", limit)).toBe(true);
});
it("tracks handlers independently", () => {
// Handler A uses all tokens
for (let i = 0; i < 3; i++) {
isHandlerRateLimited("handler-a", 3);
}
expect(isHandlerRateLimited("handler-a", 3)).toBe(true);
// Handler B should still have tokens
expect(isHandlerRateLimited("handler-b", 3)).toBe(false);
});
it("refills tokens after window expires", () => {
vi.useFakeTimers();
const limit = 2;
const windowMs = 60_000;
// Use up all tokens
expect(isHandlerRateLimited("test-handler", limit, windowMs)).toBe(false);
expect(isHandlerRateLimited("test-handler", limit, windowMs)).toBe(false);
expect(isHandlerRateLimited("test-handler", limit, windowMs)).toBe(true);
// Advance time past window
vi.advanceTimersByTime(windowMs);
// Should have tokens again
expect(isHandlerRateLimited("test-handler", limit, windowMs)).toBe(false);
});
});
describe("getHandlerRateLimitInfo", () => {
it("returns null for unknown handler", () => {
const info = getHandlerRateLimitInfo("unknown");
expect(info).toBeNull();
});
it("returns remaining tokens after usage", () => {
isHandlerRateLimited("test-handler", 5);
isHandlerRateLimited("test-handler", 5);
const info = getHandlerRateLimitInfo("test-handler");
expect(info?.remaining).toBe(3);
});
});
describe("resetHandlerRateLimit", () => {
it("clears rate limit for specific handler", () => {
// Use up tokens
for (let i = 0; i < 3; i++) {
isHandlerRateLimited("test-handler", 3);
}
expect(isHandlerRateLimited("test-handler", 3)).toBe(true);
// Reset
resetHandlerRateLimit("test-handler");
// Should have tokens again
expect(isHandlerRateLimited("test-handler", 3)).toBe(false);
});
});
describe("resetAllRateLimits", () => {
it("clears all rate limits", () => {
// Use up tokens for multiple handlers
for (let i = 0; i < 3; i++) {
isHandlerRateLimited("handler-a", 3);
isHandlerRateLimited("handler-b", 3);
}
expect(isHandlerRateLimited("handler-a", 3)).toBe(true);
expect(isHandlerRateLimited("handler-b", 3)).toBe(true);
// Reset all
resetAllRateLimits();
// Both should have tokens again
expect(isHandlerRateLimited("handler-a", 3)).toBe(false);
expect(isHandlerRateLimited("handler-b", 3)).toBe(false);
});
});
});

View File

@ -0,0 +1,87 @@
/**
* Simple token bucket rate limiter for message handlers.
* Limits execution to prevent cost explosions from unbounded agent execution.
*/
const DEFAULT_RATE_LIMIT = 10; // 10 executions per minute per handler
const DEFAULT_WINDOW_MS = 60_000; // 1 minute window
type HandlerBucket = {
tokens: number;
lastRefill: number;
};
const buckets = new Map<string, HandlerBucket>();
/**
* Check if a handler is rate limited.
* Returns true if the handler should be blocked (rate limited).
* Returns false and consumes a token if the handler can proceed.
*
* @param handlerId - Unique identifier for the handler
* @param rateLimit - Max executions per window (default: 10)
* @param windowMs - Window duration in ms (default: 60000)
* @returns true if rate limited (should skip), false if allowed
*/
export function isHandlerRateLimited(
handlerId: string,
rateLimit: number = DEFAULT_RATE_LIMIT,
windowMs: number = DEFAULT_WINDOW_MS,
): boolean {
const now = Date.now();
let bucket = buckets.get(handlerId);
if (!bucket) {
bucket = { tokens: rateLimit, lastRefill: now };
buckets.set(handlerId, bucket);
}
// Refill tokens based on time elapsed
const elapsed = now - bucket.lastRefill;
if (elapsed >= windowMs) {
bucket.tokens = rateLimit;
bucket.lastRefill = now;
}
// Check if we have tokens
if (bucket.tokens <= 0) {
return true; // Rate limited
}
// Consume a token
bucket.tokens -= 1;
return false; // Not rate limited
}
/**
* Get rate limit info for a handler (for debugging/monitoring).
*
* @param handlerId - Unique identifier for the handler
* @returns Current rate limit state or null if no bucket exists
*/
export function getHandlerRateLimitInfo(
handlerId: string,
): { remaining: number; resetMs: number } | null {
const bucket = buckets.get(handlerId);
if (!bucket) return null;
return {
remaining: bucket.tokens,
resetMs: DEFAULT_WINDOW_MS - (Date.now() - bucket.lastRefill),
};
}
/**
* Reset rate limit for a handler (for testing).
*
* @param handlerId - Unique identifier for the handler
*/
export function resetHandlerRateLimit(handlerId: string): void {
buckets.delete(handlerId);
}
/**
* Reset all rate limits (for testing).
*/
export function resetAllRateLimits(): void {
buckets.clear();
}

View File

@ -8,7 +8,7 @@ import type { MoltbotConfig } from "../config/config.js";
import type { MessageHandlerConfig } from "../config/types.hooks.js"; import type { MessageHandlerConfig } from "../config/types.hooks.js";
import { runCronIsolatedAgentTurn } from "../cron/isolated-agent/run.js"; import { runCronIsolatedAgentTurn } from "../cron/isolated-agent/run.js";
import type { CronJob, CronMessageChannel } from "../cron/types.js"; import type { CronJob, CronMessageChannel } from "../cron/types.js";
import { logVerbose } from "../globals.js"; import { logWarn } from "../logger.js";
import type { MessageReceivedHookContext } from "./internal-hooks.js"; import type { MessageReceivedHookContext } from "./internal-hooks.js";
export type RunMessageHandlerParams = { export type RunMessageHandlerParams = {
@ -81,14 +81,14 @@ export async function runMessageHandler(
}); });
if (result.status === "error") { if (result.status === "error") {
logVerbose(`message-handler-run: handler "${handler.id}" failed: ${result.error}`); logWarn(`message-handler-run: handler "${handler.id}" failed: ${result.error}`);
return { status: "error", error: result.error }; return { status: "error", error: result.error };
} }
return { status: result.status }; return { status: result.status };
} catch (err) { } catch (err) {
const errorMsg = err instanceof Error ? err.message : String(err); const errorMsg = err instanceof Error ? err.message : String(err);
logVerbose(`message-handler-run: handler "${handler.id}" threw: ${errorMsg}`); logWarn(`message-handler-run: handler "${handler.id}" threw: ${errorMsg}`);
return { status: "error", error: errorMsg }; return { status: "error", error: errorMsg };
} }
} }