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 <noreply@anthropic.com>
This commit is contained in:
Parker Todd Brooks
2026-02-16 00:36:48 -08:00
committed by Peter Steinberger
parent 94eecaa446
commit 15fe87e6b7
4 changed files with 148 additions and 5 deletions

View File

@@ -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;

View File

@@ -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 }