From 15fe87e6b714ccb6d6789a50627f4c26d5c8513b Mon Sep 17 00:00:00 2001 From: Parker Todd Brooks Date: Mon, 16 Feb 2026 00:36:48 -0800 Subject: [PATCH] feat: add before_message_write plugin hook Synchronous hook that lets plugins inspect and optionally block messages before they are written to the session JSONL file. Primary use case is private mode... when enabled, the plugin returns { block: true } and the message never gets persisted. The hook runs on the hot path (synchronous, like tool_result_persist). Handlers execute sequentially in priority order. If any handler returns { block: true }, the write is skipped immediately. Handlers can also return a modified message to write instead of the original. Changes: - src/plugins/types.ts: add hook name, event/result types, handler map entry - src/plugins/hooks.ts: add runBeforeMessageWrite() following tool_result_persist pattern - src/agents/session-tool-result-guard.ts: invoke hook before every originalAppend() call - src/agents/session-tool-result-guard-wrapper.ts: wire hook runner to the guard Co-Authored-By: Claude Opus 4.6 --- .../session-tool-result-guard-wrapper.ts | 10 +++ src/agents/session-tool-result-guard.ts | 42 ++++++++-- src/plugins/hooks.ts | 84 +++++++++++++++++++ src/plugins/types.ts | 17 ++++ 4 files changed, 148 insertions(+), 5 deletions(-) diff --git a/src/agents/session-tool-result-guard-wrapper.ts b/src/agents/session-tool-result-guard-wrapper.ts index 32bfd27d35e..896680234c6 100644 --- a/src/agents/session-tool-result-guard-wrapper.ts +++ b/src/agents/session-tool-result-guard-wrapper.ts @@ -29,6 +29,15 @@ export function guardSessionManager( } const hookRunner = getGlobalHookRunner(); + const beforeMessageWrite = hookRunner?.hasHooks("before_message_write") + ? (event: { message: import("@mariozechner/pi-agent-core").AgentMessage }) => { + return hookRunner.runBeforeMessageWrite(event, { + agentId: opts?.agentId, + sessionKey: opts?.sessionKey, + }); + } + : undefined; + const transform = hookRunner?.hasHooks("tool_result_persist") ? // oxlint-disable-next-line typescript/no-explicit-any (message: any, meta: { toolCallId?: string; toolName?: string; isSynthetic?: boolean }) => { @@ -55,6 +64,7 @@ export function guardSessionManager( applyInputProvenanceToUserMessage(message, opts?.inputProvenance), transformToolResultForPersistence: transform, allowSyntheticToolResults: opts?.allowSyntheticToolResults, + beforeMessageWriteHook: beforeMessageWrite, }); (sessionManager as GuardedSessionManager).flushPendingToolResults = guard.flushPendingToolResults; return sessionManager as GuardedSessionManager; diff --git a/src/agents/session-tool-result-guard.ts b/src/agents/session-tool-result-guard.ts index 8a2644dae45..4ff9035a119 100644 --- a/src/agents/session-tool-result-guard.ts +++ b/src/agents/session-tool-result-guard.ts @@ -1,4 +1,8 @@ import type { AgentMessage } from "@mariozechner/pi-agent-core"; +import type { + PluginHookBeforeMessageWriteEvent, + PluginHookBeforeMessageWriteResult, +} from "../plugins/types.js"; import type { TextContent } from "@mariozechner/pi-ai"; import type { SessionManager } from "@mariozechner/pi-coding-agent"; import { emitSessionTranscriptUpdate } from "../sessions/transcript-events.js"; @@ -92,6 +96,14 @@ export function installSessionToolResultGuard( * Defaults to true. */ allowSyntheticToolResults?: boolean; + /** + * Synchronous hook invoked before any message is written to the session JSONL. + * If the hook returns { block: true }, the message is silently dropped. + * If it returns { message }, the modified message is written instead. + */ + beforeMessageWriteHook?: ( + event: PluginHookBeforeMessageWriteEvent, + ) => PluginHookBeforeMessageWriteResult | undefined; }, ): { flushPendingToolResults: () => void; @@ -113,6 +125,19 @@ export function installSessionToolResultGuard( }; const allowSyntheticToolResults = opts?.allowSyntheticToolResults ?? true; + const beforeWrite = opts?.beforeMessageWriteHook; + + /** + * Run the before_message_write hook. Returns the (possibly modified) message, + * or null if the message should be blocked. + */ + const applyBeforeWriteHook = (msg: AgentMessage): AgentMessage | null => { + if (!beforeWrite) return msg; + const result = beforeWrite({ message: msg }); + if (result?.block) return null; + if (result?.message) return result.message; + return msg; + }; const flushPendingToolResults = () => { if (pending.size === 0) { @@ -121,13 +146,16 @@ export function installSessionToolResultGuard( if (allowSyntheticToolResults) { for (const [id, name] of pending.entries()) { const synthetic = makeMissingToolResult({ toolCallId: id, toolName: name }); - originalAppend( + const flushed = applyBeforeWriteHook( persistToolResult(persistMessage(synthetic), { toolCallId: id, toolName: name, isSynthetic: true, - }) as never, + }), ); + if (flushed) { + originalAppend(flushed as never); + } } } pending.clear(); @@ -157,13 +185,15 @@ export function installSessionToolResultGuard( // Apply hard size cap before persistence to prevent oversized tool results // from consuming the entire context window on subsequent LLM calls. const capped = capToolResultSize(persistMessage(nextMessage)); - return originalAppend( + const persisted = applyBeforeWriteHook( persistToolResult(capped, { toolCallId: id ?? undefined, toolName, isSynthetic: false, - }) as never, + }), ); + if (!persisted) return undefined; + return originalAppend(persisted as never); } const toolCalls = @@ -182,7 +212,9 @@ export function installSessionToolResultGuard( } } - const result = originalAppend(persistMessage(nextMessage) as never); + const finalMessage = applyBeforeWriteHook(persistMessage(nextMessage)); + if (!finalMessage) return undefined; + const result = originalAppend(finalMessage as never); const sessionFile = ( sessionManager as { getSessionFile?: () => string | null } diff --git a/src/plugins/hooks.ts b/src/plugins/hooks.ts index d8eab80aed4..006958f71f0 100644 --- a/src/plugins/hooks.ts +++ b/src/plugins/hooks.ts @@ -36,6 +36,8 @@ import type { PluginHookToolResultPersistContext, PluginHookToolResultPersistEvent, PluginHookToolResultPersistResult, + PluginHookBeforeMessageWriteEvent, + PluginHookBeforeMessageWriteResult, } from "./types.js"; // Re-export types for consumers @@ -61,6 +63,8 @@ export type { PluginHookToolResultPersistContext, PluginHookToolResultPersistEvent, PluginHookToolResultPersistResult, + PluginHookBeforeMessageWriteEvent, + PluginHookBeforeMessageWriteResult, PluginHookSessionContext, PluginHookSessionStartEvent, PluginHookSessionEndEvent, @@ -410,6 +414,84 @@ export function createHookRunner(registry: PluginRegistry, options: HookRunnerOp return { message: current }; } + + // ========================================================================= + // Message Write Hooks + // ========================================================================= + + /** + * Run before_message_write hook. + * + * This hook is intentionally synchronous: it runs on the hot path where + * session transcripts are appended synchronously. + * + * Handlers are executed sequentially in priority order (higher first). + * If any handler returns { block: true }, the message is NOT written + * to the session JSONL and we return immediately. + * If a handler returns { message }, the modified message replaces the + * original for subsequent handlers and the final write. + */ + function runBeforeMessageWrite( + event: PluginHookBeforeMessageWriteEvent, + ctx: { agentId?: string; sessionKey?: string }, + ): PluginHookBeforeMessageWriteResult | undefined { + const hooks = getHooksForName(registry, "before_message_write"); + if (hooks.length === 0) { + return undefined; + } + + let current = event.message; + + for (const hook of hooks) { + try { + // oxlint-disable-next-line typescript/no-explicit-any + const out = (hook.handler as any)({ ...event, message: current }, ctx) as + | PluginHookBeforeMessageWriteResult + | void + | Promise; + + // Guard against accidental async handlers (this hook is sync-only). + // oxlint-disable-next-line typescript/no-explicit-any + if (out && typeof (out as any).then === "function") { + const msg = + `[hooks] before_message_write handler from ${hook.pluginId} returned a Promise; ` + + `this hook is synchronous and the result was ignored.`; + if (catchErrors) { + logger?.warn?.(msg); + continue; + } + throw new Error(msg); + } + + const result = out as PluginHookBeforeMessageWriteResult | undefined; + + // If any handler blocks, return immediately. + if (result?.block) { + return { block: true }; + } + + // If handler provided a modified message, use it for subsequent handlers. + if (result?.message) { + current = result.message; + } + } catch (err) { + const msg = `[hooks] before_message_write handler from ${hook.pluginId} failed: ${String(err)}`; + if (catchErrors) { + logger?.error(msg); + } else { + throw new Error(msg, { cause: err }); + } + } + } + + // If message was modified by any handler, return it. + if (current !== event.message) { + return { message: current }; + } + + return undefined; + } + // ========================================================================= // Session Hooks // ========================================================================= @@ -497,6 +579,8 @@ export function createHookRunner(registry: PluginRegistry, options: HookRunnerOp runBeforeToolCall, runAfterToolCall, runToolResultPersist, + // Message write hooks + runBeforeMessageWrite, // Session hooks runSessionStart, runSessionEnd, diff --git a/src/plugins/types.ts b/src/plugins/types.ts index 25ee3ced18a..dd04e6269be 100644 --- a/src/plugins/types.ts +++ b/src/plugins/types.ts @@ -309,6 +309,7 @@ export type PluginHookName = | "before_tool_call" | "after_tool_call" | "tool_result_persist" + | "before_message_write" | "session_start" | "session_end" | "gateway_start" @@ -493,6 +494,18 @@ export type PluginHookToolResultPersistResult = { message?: AgentMessage; }; +// before_message_write hook +export type PluginHookBeforeMessageWriteEvent = { + message: AgentMessage; + sessionKey?: string; + agentId?: string; +}; + +export type PluginHookBeforeMessageWriteResult = { + block?: boolean; // If true, message is NOT written to JSONL + message?: AgentMessage; // Optional: modified message to write instead +}; + // Session context export type PluginHookSessionContext = { agentId?: string; @@ -575,6 +588,10 @@ export type PluginHookHandlerMap = { event: PluginHookToolResultPersistEvent, ctx: PluginHookToolResultPersistContext, ) => PluginHookToolResultPersistResult | void; + before_message_write: ( + event: PluginHookBeforeMessageWriteEvent, + ctx: { agentId?: string; sessionKey?: string }, + ) => PluginHookBeforeMessageWriteResult | void; session_start: ( event: PluginHookSessionStartEvent, ctx: PluginHookSessionContext,