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,