diff --git a/plugin/scripts/post-tool-failure.mjs b/plugin/scripts/post-tool-failure.mjs index 6fdad8d9d..ba5950a8c 100755 --- a/plugin/scripts/post-tool-failure.mjs +++ b/plugin/scripts/post-tool-failure.mjs @@ -1,7 +1,46 @@ #!/usr/bin/env node import { execSync } from "node:child_process"; import { basename } from "node:path"; - +//#region src/hooks/_capture-filter.ts +const DEFAULT_DENY_PATTERNS = [ + "memory_*", + "toolsearch", + "listmcpresources", + "fetchmcpresource" +]; +function parseEnvList(raw) { + if (!raw?.trim()) return void 0; + return raw.split(/[,\s]+/).map((part) => part.trim()).filter(Boolean); +} +function bareToolName(toolName) { + const trimmed = toolName.trim(); + if (/^mcp__/i.test(trimmed)) { + const parts = trimmed.split("__"); + if (parts.length >= 3) return parts[parts.length - 1]; + } + return trimmed; +} +function normalizePattern(pattern) { + return pattern.trim().toLowerCase(); +} +function matchesPattern(toolName, pattern) { + const bare = bareToolName(toolName).toLowerCase(); + const pat = normalizePattern(pattern); + if (!pat.includes("*")) return bare === pat; + const escaped = pat.replace(/[.+?^${}()|[\]\\]/g, "\\$&"); + const re = new RegExp(`^${escaped.replace(/\*/g, ".*")}$`); + return re.test(bare) || re.test(toolName.toLowerCase()); +} +function matchesAny(toolName, patterns) { + return patterns.some((pattern) => matchesPattern(toolName, pattern)); +} +function shouldCaptureTool(toolName) { + if (typeof toolName !== "string" || !toolName.trim()) return true; + const allow = parseEnvList(process.env["AGENTMEMORY_CAPTURE_ALLOW"]); + if (allow) return matchesAny(toolName, allow); + return !matchesAny(toolName, [...DEFAULT_DENY_PATTERNS, ...parseEnvList(process.env["AGENTMEMORY_CAPTURE_DENY"]) ?? []]); +} +//#endregion //#region src/hooks/_project.ts function resolveProject(cwd) { const explicit = process.env["AGENTMEMORY_PROJECT_NAME"]; @@ -21,7 +60,6 @@ function resolveProject(cwd) { } catch {} return basename(dir); } - //#endregion //#region src/hooks/post-tool-failure.ts function isSdkChildContext(payload) { @@ -49,6 +87,7 @@ async function main() { if (data.is_interrupt || data.isInterrupt) return; const sessionId = data.session_id || data.sessionId || "unknown"; const toolName = data.tool_name ?? data.toolName; + if (!shouldCaptureTool(toolName)) return; const toolInput = data.tool_input ?? data.toolArgs; const error = data.error ?? data.errorMessage; fetch(`${REST_URL}/agentmemory/observe`, { @@ -71,7 +110,7 @@ async function main() { setTimeout(() => process.exit(0), 500).unref(); } main(); - //#endregion -export { }; +export {}; + //# sourceMappingURL=post-tool-failure.mjs.map \ No newline at end of file diff --git a/plugin/scripts/post-tool-use.mjs b/plugin/scripts/post-tool-use.mjs index b4aef9c94..9308239c8 100755 --- a/plugin/scripts/post-tool-use.mjs +++ b/plugin/scripts/post-tool-use.mjs @@ -1,7 +1,52 @@ #!/usr/bin/env node import { execSync } from "node:child_process"; import { basename } from "node:path"; - +//#region src/hooks/_capture-filter.ts +const DEFAULT_DENY_PATTERNS = [ + "memory_*", + "toolsearch", + "listmcpresources", + "fetchmcpresource" +]; +function parseEnvList(raw) { + if (!raw?.trim()) return void 0; + return raw.split(/[,\s]+/).map((part) => part.trim()).filter(Boolean); +} +function bareToolName(toolName) { + const trimmed = toolName.trim(); + if (/^mcp__/i.test(trimmed)) { + const parts = trimmed.split("__"); + if (parts.length >= 3) return parts[parts.length - 1]; + } + return trimmed; +} +function normalizePattern(pattern) { + return pattern.trim().toLowerCase(); +} +function matchesPattern(toolName, pattern) { + const bare = bareToolName(toolName).toLowerCase(); + const pat = normalizePattern(pattern); + if (!pat.includes("*")) return bare === pat; + const escaped = pat.replace(/[.+?^${}()|[\]\\]/g, "\\$&"); + const re = new RegExp(`^${escaped.replace(/\*/g, ".*")}$`); + return re.test(bare) || re.test(toolName.toLowerCase()); +} +function matchesAny(toolName, patterns) { + return patterns.some((pattern) => matchesPattern(toolName, pattern)); +} +function shouldCaptureTool(toolName) { + if (typeof toolName !== "string" || !toolName.trim()) return true; + const allow = parseEnvList(process.env["AGENTMEMORY_CAPTURE_ALLOW"]); + if (allow) return matchesAny(toolName, allow); + return !matchesAny(toolName, [...DEFAULT_DENY_PATTERNS, ...parseEnvList(process.env["AGENTMEMORY_CAPTURE_DENY"]) ?? []]); +} +function captureOutputMax() { + const raw = process.env["AGENTMEMORY_CAPTURE_OUTPUT_MAX"]; + if (!raw?.trim()) return 8e3; + const parsed = Number.parseInt(raw, 10); + return Number.isFinite(parsed) && parsed > 0 ? parsed : 8e3; +} +//#endregion //#region src/hooks/_project.ts function resolveProject(cwd) { const explicit = process.env["AGENTMEMORY_PROJECT_NAME"]; @@ -21,7 +66,6 @@ function resolveProject(cwd) { } catch {} return basename(dir); } - //#endregion //#region src/hooks/post-tool-use.ts function isSdkChildContext(payload) { @@ -48,8 +92,10 @@ async function main() { if (isSdkChildContext(data)) return; const sessionId = data.session_id || data.sessionId || "unknown"; const toolName = data.tool_name ?? data.toolName; + if (!shouldCaptureTool(toolName)) return; const toolInput = data.tool_input ?? data.toolArgs; const { imageData, cleanOutput } = extractImageData(toolOutput(data)); + const outputMax = captureOutputMax(); fetch(`${REST_URL}/agentmemory/observe`, { method: "POST", headers: authHeaders(), @@ -62,7 +108,7 @@ async function main() { data: { tool_name: toolName, tool_input: toolInput, - tool_output: truncate(cleanOutput, 8e3), + tool_output: truncate(cleanOutput, outputMax), ...imageData ? { image_data: imageData } : {} } }), @@ -116,7 +162,7 @@ function truncate(value, max) { return value; } main(); - //#endregion -export { }; +export {}; + //# sourceMappingURL=post-tool-use.mjs.map \ No newline at end of file diff --git a/plugin/scripts/pre-compact.mjs b/plugin/scripts/pre-compact.mjs index 0afdcb0b4..519652740 100755 --- a/plugin/scripts/pre-compact.mjs +++ b/plugin/scripts/pre-compact.mjs @@ -1,7 +1,15 @@ #!/usr/bin/env node import { execSync } from "node:child_process"; import { basename } from "node:path"; - +//#region src/hooks/_capture-filter.ts +function preCompactBudget() { + const raw = process.env["AGENTMEMORY_PRE_COMPACT_BUDGET"]; + if (raw?.trim() === "0") return 0; + if (!raw?.trim()) return 1500; + const parsed = Number.parseInt(raw, 10); + return Number.isFinite(parsed) && parsed >= 0 ? parsed : 1500; +} +//#endregion //#region src/hooks/_project.ts function resolveProject(cwd) { const explicit = process.env["AGENTMEMORY_PROJECT_NAME"]; @@ -21,7 +29,6 @@ function resolveProject(cwd) { } catch {} return basename(dir); } - //#endregion //#region src/hooks/pre-compact.ts function isSdkChildContext(payload) { @@ -56,6 +63,8 @@ async function main() { signal: AbortSignal.timeout(5e3) }); } catch {} + const budget = preCompactBudget(); + if (budget === 0) return; try { const res = await fetch(`${REST_URL}/agentmemory/context`, { method: "POST", @@ -63,7 +72,7 @@ async function main() { body: JSON.stringify({ sessionId, project, - budget: 1500 + budget }), signal: AbortSignal.timeout(5e3) }); @@ -74,7 +83,7 @@ async function main() { } catch {} } main(); - //#endregion -export { }; +export {}; + //# sourceMappingURL=pre-compact.mjs.map \ No newline at end of file diff --git a/src/hooks/_capture-filter.ts b/src/hooks/_capture-filter.ts new file mode 100644 index 000000000..ba7493ee9 --- /dev/null +++ b/src/hooks/_capture-filter.ts @@ -0,0 +1,68 @@ +const DEFAULT_DENY_PATTERNS = [ + "memory_*", + "toolsearch", + "listmcpresources", + "fetchmcpresource", +]; + +function parseEnvList(raw: string | undefined): string[] | undefined { + if (!raw?.trim()) return undefined; + return raw + .split(/[,\s]+/) + .map((part) => part.trim()) + .filter(Boolean); +} + +export function bareToolName(toolName: string): string { + const trimmed = toolName.trim(); + if (/^mcp__/i.test(trimmed)) { + const parts = trimmed.split("__"); + if (parts.length >= 3) return parts[parts.length - 1]!; + } + return trimmed; +} + +function normalizePattern(pattern: string): string { + return pattern.trim().toLowerCase(); +} + +function matchesPattern(toolName: string, pattern: string): boolean { + const bare = bareToolName(toolName).toLowerCase(); + const pat = normalizePattern(pattern); + if (!pat.includes("*")) return bare === pat; + const escaped = pat.replace(/[.+?^${}()|[\]\\]/g, "\\$&"); + const re = new RegExp(`^${escaped.replace(/\*/g, ".*")}$`); + return re.test(bare) || re.test(toolName.toLowerCase()); +} + +function matchesAny(toolName: string, patterns: string[]): boolean { + return patterns.some((pattern) => matchesPattern(toolName, pattern)); +} + +export function shouldCaptureTool(toolName: unknown): boolean { + if (typeof toolName !== "string" || !toolName.trim()) return true; + + const allow = parseEnvList(process.env["AGENTMEMORY_CAPTURE_ALLOW"]); + if (allow) return matchesAny(toolName, allow); + + const deny = [ + ...DEFAULT_DENY_PATTERNS, + ...(parseEnvList(process.env["AGENTMEMORY_CAPTURE_DENY"]) ?? []), + ]; + return !matchesAny(toolName, deny); +} + +export function captureOutputMax(): number { + const raw = process.env["AGENTMEMORY_CAPTURE_OUTPUT_MAX"]; + if (!raw?.trim()) return 8000; + const parsed = Number.parseInt(raw, 10); + return Number.isFinite(parsed) && parsed > 0 ? parsed : 8000; +} + +export function preCompactBudget(): number { + const raw = process.env["AGENTMEMORY_PRE_COMPACT_BUDGET"]; + if (raw?.trim() === "0") return 0; + if (!raw?.trim()) return 1500; + const parsed = Number.parseInt(raw, 10); + return Number.isFinite(parsed) && parsed >= 0 ? parsed : 1500; +} diff --git a/src/hooks/post-tool-failure.ts b/src/hooks/post-tool-failure.ts index f0295b3e6..a0f54e907 100644 --- a/src/hooks/post-tool-failure.ts +++ b/src/hooks/post-tool-failure.ts @@ -1,4 +1,5 @@ #!/usr/bin/env node +import { shouldCaptureTool } from "./_capture-filter.js"; import { resolveProject } from "./_project.js"; function isSdkChildContext(payload: unknown): boolean { @@ -34,6 +35,8 @@ async function main() { const sessionId = ((data.session_id || data.sessionId) as string) || "unknown"; const toolName = data.tool_name ?? data.toolName; + if (!shouldCaptureTool(toolName)) return; + const toolInput = data.tool_input ?? data.toolArgs; const error = data.error ?? data.errorMessage; diff --git a/src/hooks/post-tool-use.ts b/src/hooks/post-tool-use.ts index c68a7731d..28e1ae40d 100644 --- a/src/hooks/post-tool-use.ts +++ b/src/hooks/post-tool-use.ts @@ -1,4 +1,5 @@ #!/usr/bin/env node +import { captureOutputMax, shouldCaptureTool } from "./_capture-filter.js"; import { resolveProject } from "./_project.js"; function isSdkChildContext(payload: unknown): boolean { @@ -33,9 +34,12 @@ async function main() { const sessionId = ((data.session_id || data.sessionId) as string) || "unknown"; const toolName = data.tool_name ?? data.toolName; + if (!shouldCaptureTool(toolName)) return; + const toolInput = data.tool_input ?? data.toolArgs; const { imageData, cleanOutput } = extractImageData(toolOutput(data)); + const outputMax = captureOutputMax(); fetch(`${REST_URL}/agentmemory/observe`, { method: "POST", @@ -49,7 +53,7 @@ async function main() { data: { tool_name: toolName, tool_input: toolInput, - tool_output: truncate(cleanOutput, 8000), + tool_output: truncate(cleanOutput, outputMax), ...(imageData ? { image_data: imageData } : {}), }, }), diff --git a/src/hooks/pre-compact.ts b/src/hooks/pre-compact.ts index 8bb2660c4..09cade14e 100644 --- a/src/hooks/pre-compact.ts +++ b/src/hooks/pre-compact.ts @@ -1,4 +1,5 @@ #!/usr/bin/env node +import { preCompactBudget } from "./_capture-filter.js"; import { resolveProject } from "./_project.js"; function isSdkChildContext(payload: unknown): boolean { @@ -47,11 +48,14 @@ async function main() { } } + const budget = preCompactBudget(); + if (budget === 0) return; + try { const res = await fetch(`${REST_URL}/agentmemory/context`, { method: "POST", headers: authHeaders(), - body: JSON.stringify({ sessionId, project, budget: 1500 }), + body: JSON.stringify({ sessionId, project, budget }), signal: AbortSignal.timeout(5000), }); diff --git a/test/capture-filter.test.ts b/test/capture-filter.test.ts new file mode 100644 index 000000000..98bcba48f --- /dev/null +++ b/test/capture-filter.test.ts @@ -0,0 +1,113 @@ +import { describe, it, expect, beforeEach, afterEach } from "vitest"; +import { + bareToolName, + captureOutputMax, + preCompactBudget, + shouldCaptureTool, +} from "../src/hooks/_capture-filter.js"; + +describe("bareToolName", () => { + it("returns plain tool names unchanged", () => { + expect(bareToolName("Bash")).toBe("Bash"); + expect(bareToolName("memory_recall")).toBe("memory_recall"); + }); + + it("strips Claude Code MCP prefixes", () => { + expect(bareToolName("mcp__agentmemory__memory_smart_search")).toBe( + "memory_smart_search", + ); + }); +}); + +describe("shouldCaptureTool (#993)", () => { + const envKeys = [ + "AGENTMEMORY_CAPTURE_ALLOW", + "AGENTMEMORY_CAPTURE_DENY", + ] as const; + const original: Record = {}; + + beforeEach(() => { + for (const key of envKeys) { + original[key] = process.env[key]; + delete process.env[key]; + } + }); + + afterEach(() => { + for (const key of envKeys) { + if (original[key] === undefined) delete process.env[key]; + else process.env[key] = original[key]; + } + }); + + it("captures normal agent tools by default", () => { + expect(shouldCaptureTool("Bash")).toBe(true); + expect(shouldCaptureTool("Edit")).toBe(true); + expect(shouldCaptureTool("Agent")).toBe(true); + }); + + it("skips agentmemory memory_* MCP tools by default", () => { + expect(shouldCaptureTool("memory_recall")).toBe(false); + expect(shouldCaptureTool("memory_smart_search")).toBe(false); + expect(shouldCaptureTool("mcp__agentmemory__memory_recall")).toBe(false); + }); + + it("skips plumbing tools by default", () => { + expect(shouldCaptureTool("ToolSearch")).toBe(false); + expect(shouldCaptureTool("ListMcpResources")).toBe(false); + }); + + it("honors AGENTMEMORY_CAPTURE_DENY", () => { + process.env.AGENTMEMORY_CAPTURE_DENY = "Grep, Glob"; + expect(shouldCaptureTool("Grep")).toBe(false); + expect(shouldCaptureTool("Bash")).toBe(true); + }); + + it("honors AGENTMEMORY_CAPTURE_ALLOW over defaults", () => { + process.env.AGENTMEMORY_CAPTURE_ALLOW = "memory_recall, Bash"; + expect(shouldCaptureTool("memory_recall")).toBe(true); + expect(shouldCaptureTool("memory_smart_search")).toBe(false); + expect(shouldCaptureTool("Bash")).toBe(true); + expect(shouldCaptureTool("Edit")).toBe(false); + }); +}); + +describe("captureOutputMax", () => { + const key = "AGENTMEMORY_CAPTURE_OUTPUT_MAX"; + const original = process.env[key]; + + afterEach(() => { + if (original === undefined) delete process.env[key]; + else process.env[key] = original; + }); + + it("defaults to 8000", () => { + delete process.env[key]; + expect(captureOutputMax()).toBe(8000); + }); + + it("reads a positive override", () => { + process.env[key] = "4096"; + expect(captureOutputMax()).toBe(4096); + }); +}); + +describe("preCompactBudget", () => { + const key = "AGENTMEMORY_PRE_COMPACT_BUDGET"; + const original = process.env[key]; + + afterEach(() => { + if (original === undefined) delete process.env[key]; + else process.env[key] = original; + }); + + it("defaults to 1500", () => { + delete process.env[key]; + expect(preCompactBudget()).toBe(1500); + }); + + it("allows disabling injection with 0", () => { + process.env[key] = "0"; + expect(preCompactBudget()).toBe(0); + }); +}); diff --git a/test/post-tool-use-capture.test.ts b/test/post-tool-use-capture.test.ts new file mode 100644 index 000000000..a4e48302f --- /dev/null +++ b/test/post-tool-use-capture.test.ts @@ -0,0 +1,119 @@ +import { describe, it, expect, afterEach } from "vitest"; +import { spawn } from "node:child_process"; +import { createServer, type Server } from "node:http"; +import { join } from "node:path"; + +const HOOKS_DIR = join(import.meta.dirname, "..", "plugin", "scripts"); + +function runHook( + scriptName: string, + stdin: string, + env: Record, +): Promise<{ exitCode: number | null }> { + return new Promise((resolve, reject) => { + const child = spawn(process.execPath, [join(HOOKS_DIR, scriptName)], { + env: { + PATH: process.env["PATH"] ?? "", + ...env, + }, + stdio: ["pipe", "ignore", "ignore"], + }); + child.on("error", reject); + child.on("close", (exitCode) => resolve({ exitCode })); + child.stdin.write(stdin); + child.stdin.end(); + }); +} + +describe("post-tool-use hook — capture filter (#993)", () => { + let server: Server; + let observeCalls = 0; + let port = 0; + let onObserve: (() => void) | undefined; + + afterEach(async () => { + observeCalls = 0; + onObserve = undefined; + if (server) { + await new Promise((resolve, reject) => { + server.close((err) => (err ? reject(err) : resolve())); + }); + } + }); + + async function assertNoObserveAfterHook( + hook: Promise<{ exitCode: number | null }>, + { settleMs = 300 } = {}, + ): Promise<{ exitCode: number | null }> { + const result = await hook; + const deadline = Date.now() + settleMs; + while (Date.now() < deadline) { + if (observeCalls !== 0) { + throw new Error(`unexpected observe call(s): ${observeCalls}`); + } + await new Promise((r) => setTimeout(r, 25)); + } + return result; + } + + async function startServer() { + server = createServer((req, res) => { + if (req.method === "POST" && req.url === "/agentmemory/observe") { + observeCalls += 1; + onObserve?.(); + } + res.writeHead(200, { "Content-Type": "application/json" }); + res.end("{}"); + }); + await new Promise((resolve) => { + server.listen(0, "127.0.0.1", () => { + const addr = server.address(); + if (addr && typeof addr === "object") port = addr.port; + resolve(); + }); + }); + } + + it("does not observe memory_* MCP tools", async () => { + await startServer(); + const payload = JSON.stringify({ + session_id: "ses_test", + tool_name: "mcp__agentmemory__memory_smart_search", + tool_input: { query: "hooks" }, + tool_output: "results", + }); + const result = await assertNoObserveAfterHook( + runHook("post-tool-use.mjs", payload, { + AGENTMEMORY_URL: `http://127.0.0.1:${port}`, + }), + ); + expect(result.exitCode).toBe(0); + expect(observeCalls).toBe(0); + }); + + it("still observes normal tools", async () => { + await startServer(); + const payload = JSON.stringify({ + session_id: "ses_test", + tool_name: "Bash", + tool_input: { command: "ls" }, + tool_output: "ok", + }); + const observeSeen = new Promise((resolve, reject) => { + const timer = setTimeout( + () => reject(new Error("observe request never arrived")), + 2000, + ); + onObserve = () => { + clearTimeout(timer); + resolve(); + }; + }); + const result = await runHook("post-tool-use.mjs", payload, { + AGENTMEMORY_URL: `http://127.0.0.1:${port}`, + }); + await observeSeen; + expect(result.exitCode).toBe(0); + expect(observeCalls).toBe(1); + }); +});