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:
parent
42c647f2c9
commit
761b7b6849
@ -366,16 +366,18 @@ Message handlers fix this by immediately processing important messages as they a
|
||||
|
||||
### 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 |
|
||||
|-------|------|-------------|
|
||||
| `channelId` | `string \| string[]` | Channel to match: `"whatsapp"`, `"telegram"`, `["discord", "slack"]`, or `"*"` for all |
|
||||
| `conversationId` | `string \| string[]` | Chat/group ID to match |
|
||||
| `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) |
|
||||
|
||||
**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
|
||||
|
||||
| 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).
|
||||
|
||||
### 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
|
||||
|
||||
1. Handlers are evaluated in order (first match wins)
|
||||
2. Disabled handlers (`enabled: false`) are skipped
|
||||
3. If a handler matches with `mode: "exclusive"` (default), normal processing stops
|
||||
4. If a handler matches with `mode: "parallel"`, normal processing continues after the handler fires
|
||||
3. Rate limiting is checked before execution
|
||||
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
|
||||
|
||||
|
||||
@ -205,6 +205,7 @@
|
||||
"tslog": "^4.10.2",
|
||||
"undici": "^7.19.0",
|
||||
"ws": "^8.19.0",
|
||||
"safe-regex": "^2.1.1",
|
||||
"yaml": "^2.8.2",
|
||||
"zod": "^4.3.6"
|
||||
},
|
||||
@ -223,6 +224,7 @@
|
||||
"@types/node": "^25.0.10",
|
||||
"@types/proper-lockfile": "^4.1.4",
|
||||
"@types/qrcode-terminal": "^0.12.2",
|
||||
"@types/safe-regex": "^1.1.6",
|
||||
"@types/ws": "^8.18.1",
|
||||
"@typescript/native-preview": "7.0.0-dev.20260124.1",
|
||||
"@vitest/coverage-v8": "^4.0.18",
|
||||
|
||||
113
pnpm-lock.yaml
generated
113
pnpm-lock.yaml
generated
@ -148,6 +148,9 @@ importers:
|
||||
qrcode-terminal:
|
||||
specifier: ^0.12.0
|
||||
version: 0.12.0
|
||||
safe-regex:
|
||||
specifier: ^2.1.1
|
||||
version: 2.1.1
|
||||
sharp:
|
||||
specifier: ^0.34.5
|
||||
version: 0.34.5
|
||||
@ -203,6 +206,9 @@ importers:
|
||||
'@types/qrcode-terminal':
|
||||
specifier: ^0.12.2
|
||||
version: 0.12.2
|
||||
'@types/safe-regex':
|
||||
specifier: ^1.1.6
|
||||
version: 1.1.6
|
||||
'@types/ws':
|
||||
specifier: ^8.18.1
|
||||
version: 8.18.1
|
||||
@ -383,12 +389,12 @@ importers:
|
||||
'@microsoft/agents-hosting-extensions-teams':
|
||||
specifier: ^1.2.2
|
||||
version: 1.2.2
|
||||
moltbot:
|
||||
specifier: workspace:*
|
||||
version: link:../..
|
||||
express:
|
||||
specifier: ^5.2.1
|
||||
version: 5.2.1
|
||||
moltbot:
|
||||
specifier: workspace:*
|
||||
version: link:../..
|
||||
proper-lockfile:
|
||||
specifier: ^4.1.2
|
||||
version: 4.1.2
|
||||
@ -2767,6 +2773,9 @@ packages:
|
||||
'@types/retry@0.12.5':
|
||||
resolution: {integrity: sha512-3xSjTp3v03X/lSQLkczaN9UIEwJMoMCA1+Nb5HfbJEQWogdeQIyVtTvxPXDQjZ5zws8rFQfVfRdz03ARihPJgw==}
|
||||
|
||||
'@types/safe-regex@1.1.6':
|
||||
resolution: {integrity: sha512-CQ/uPB9fLOPKwDsrTeVbNIkwfUthTWOx0l6uIGwVFjZxv7e68pCW5gtTYFzdJi3EBJp8h8zYhJbTasAbX7gEMQ==}
|
||||
|
||||
'@types/send@0.17.6':
|
||||
resolution: {integrity: sha512-Uqt8rPBE8SY0RK8JB1EzVOIZ32uqy8HwdxCnoCOsYrvnswqmFZ/k+9Ikidlk/ImhsdvBsloHbAlewb2IEBV/Og==}
|
||||
|
||||
@ -3214,11 +3223,6 @@ packages:
|
||||
class-variance-authority@0.7.1:
|
||||
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:
|
||||
resolution: {integrity: sha512-aCj4O5wKyszjMmDT4tZj93kxyydN/K5zPWSCe6/0AV/AA1pqe5ZBIw0a2ZfPQV7lL5/yb5HsUreJ6UFAF1tEQw==}
|
||||
engines: {node: '>=18'}
|
||||
@ -4864,6 +4868,10 @@ packages:
|
||||
reflect-metadata@0.2.2:
|
||||
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:
|
||||
resolution: {integrity: sha512-TTbAfBBRdWD7aNNOoVOBH4pN/KigV6LyapYNNlAPA8JwbovRti1E88m3sYAwsLi5ryhPKsE9APwnjFTgdUjTpw==}
|
||||
engines: {node: '>=0.10.0'}
|
||||
@ -4940,6 +4948,9 @@ packages:
|
||||
safe-buffer@5.2.1:
|
||||
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:
|
||||
resolution: {integrity: sha512-b3rppTKm9T+PsVCBEOUR46GWI7fdOs00VKZ1+9c1EWDaDMvjQc6tUwuFyIprgGgTcWoVHSKrU8H31ZHA2e0RHA==}
|
||||
engines: {node: '>=10'}
|
||||
@ -8549,6 +8560,8 @@ snapshots:
|
||||
|
||||
'@types/retry@0.12.5': {}
|
||||
|
||||
'@types/safe-regex@1.1.6': {}
|
||||
|
||||
'@types/send@0.17.6':
|
||||
dependencies:
|
||||
'@types/mime': 1.3.5
|
||||
@ -9098,84 +9111,6 @@ snapshots:
|
||||
dependencies:
|
||||
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:
|
||||
dependencies:
|
||||
restore-cursor: 5.1.0
|
||||
@ -11014,6 +10949,8 @@ snapshots:
|
||||
|
||||
reflect-metadata@0.2.2: {}
|
||||
|
||||
regexp-tree@0.1.27: {}
|
||||
|
||||
request-promise-core@1.1.4(request@2.88.2):
|
||||
dependencies:
|
||||
lodash: 4.17.23
|
||||
@ -11147,6 +11084,10 @@ snapshots:
|
||||
|
||||
safe-buffer@5.2.1: {}
|
||||
|
||||
safe-regex@2.1.1:
|
||||
dependencies:
|
||||
regexp-tree: 0.1.27
|
||||
|
||||
safe-stable-stringify@2.5.0: {}
|
||||
|
||||
safer-buffer@2.1.2: {}
|
||||
|
||||
@ -3,10 +3,12 @@ import { createDefaultDeps } from "../../cli/deps.js";
|
||||
import type { MoltbotConfig } from "../../config/config.js";
|
||||
import { loadSessionStore, resolveStorePath } from "../../config/sessions.js";
|
||||
import { logVerbose } from "../../globals.js";
|
||||
import { logWarn } from "../../logger.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 { 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 {
|
||||
logMessageProcessed,
|
||||
@ -263,33 +265,41 @@ export async function dispatchReplyFromConfig(params: {
|
||||
|
||||
const matchedHandler = matchMessageHandler(messageHandlers, handlerContext);
|
||||
if (matchedHandler) {
|
||||
const priority = matchedHandler.priority ?? "immediate";
|
||||
const mode = matchedHandler.mode ?? "exclusive";
|
||||
// Check rate limit before executing
|
||||
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") {
|
||||
// Fire-and-forget for immediate handlers
|
||||
const deps = createDefaultDeps();
|
||||
void runMessageHandler({
|
||||
cfg,
|
||||
deps,
|
||||
handler: matchedHandler,
|
||||
context: handlerContext,
|
||||
}).catch((err) => {
|
||||
logVerbose(
|
||||
`dispatch-from-config: message handler "${matchedHandler.id}" failed: ${String(err)}`,
|
||||
);
|
||||
});
|
||||
if (priority === "immediate") {
|
||||
// Fire-and-forget for immediate handlers
|
||||
const deps = createDefaultDeps();
|
||||
void runMessageHandler({
|
||||
cfg,
|
||||
deps,
|
||||
handler: matchedHandler,
|
||||
context: handlerContext,
|
||||
}).catch((err) => {
|
||||
logWarn(
|
||||
`dispatch-from-config: message handler "${matchedHandler.id}" failed: ${String(err)}`,
|
||||
);
|
||||
});
|
||||
|
||||
// Check mode: exclusive (default) vs parallel
|
||||
if (mode === "exclusive") {
|
||||
// Return early - handler takes over completely
|
||||
return {
|
||||
queuedFinal: false,
|
||||
counts: dispatcher.getQueuedCounts(),
|
||||
handledByMessageHandler: { handlerId: matchedHandler.id },
|
||||
};
|
||||
// Check mode: exclusive (default) vs parallel
|
||||
if (mode === "exclusive") {
|
||||
// Return early - handler takes over completely
|
||||
return {
|
||||
queuedFinal: false,
|
||||
counts: dispatcher.getQueuedCounts(),
|
||||
handledByMessageHandler: { handlerId: matchedHandler.id },
|
||||
};
|
||||
}
|
||||
// parallel mode: continue with normal flow
|
||||
}
|
||||
// parallel mode: continue with normal flow
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@ -60,7 +60,23 @@ const MessageHandlerMatchSchema = z
|
||||
contentPattern: 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
|
||||
.object({
|
||||
|
||||
@ -162,6 +162,23 @@ describe("matchMessageHandler", () => {
|
||||
const result = matchMessageHandler(handlers, ctx);
|
||||
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", () => {
|
||||
@ -304,7 +321,11 @@ describe("matchMessageHandler", () => {
|
||||
});
|
||||
|
||||
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 ctx = createContext();
|
||||
|
||||
|
||||
@ -3,7 +3,9 @@
|
||||
* 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 { logWarn } from "../logger.js";
|
||||
import type { MessageReceivedHookContext } from "./internal-hooks.js";
|
||||
|
||||
/**
|
||||
@ -53,8 +55,13 @@ function matchesConditions(match: MessageHandlerMatch, ctx: MessageReceivedHookC
|
||||
return false;
|
||||
}
|
||||
|
||||
// contentPattern (regex)
|
||||
// contentPattern (regex) - with ReDoS protection
|
||||
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 {
|
||||
const regex = new RegExp(match.contentPattern, "i");
|
||||
if (!regex.test(ctx.content)) {
|
||||
|
||||
115
src/hooks/message-handler-rate-limit.test.ts
Normal file
115
src/hooks/message-handler-rate-limit.test.ts
Normal 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);
|
||||
});
|
||||
});
|
||||
});
|
||||
87
src/hooks/message-handler-rate-limit.ts
Normal file
87
src/hooks/message-handler-rate-limit.ts
Normal 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();
|
||||
}
|
||||
@ -8,7 +8,7 @@ import type { MoltbotConfig } from "../config/config.js";
|
||||
import type { MessageHandlerConfig } from "../config/types.hooks.js";
|
||||
import { runCronIsolatedAgentTurn } from "../cron/isolated-agent/run.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";
|
||||
|
||||
export type RunMessageHandlerParams = {
|
||||
@ -81,14 +81,14 @@ export async function runMessageHandler(
|
||||
});
|
||||
|
||||
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: result.status };
|
||||
} catch (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 };
|
||||
}
|
||||
}
|
||||
|
||||
Loading…
Reference in New Issue
Block a user