feat: add tool_result_persist hook

This commit is contained in:
Doug von Kohorn
2026-01-19 13:11:31 +01:00
parent 9f280454ba
commit c3a34408f3
8 changed files with 299 additions and 7 deletions

View File

@@ -286,7 +286,10 @@ export async function compactEmbeddedPiSession(params: {
});
try {
await prewarmSessionFile(params.sessionFile);
const sessionManager = guardSessionManager(SessionManager.open(params.sessionFile));
const sessionManager = guardSessionManager(SessionManager.open(params.sessionFile), {
agentId: sessionAgentId,
sessionKey: params.sessionKey,
});
trackSessionManagerAccess(params.sessionFile);
const settingsManager = SettingsManager.create(effectiveWorkspace, agentDir);
ensurePiCompactionReserveTokens({

View File

@@ -284,7 +284,10 @@ export async function runEmbeddedAttempt(
.catch(() => false);
await prewarmSessionFile(params.sessionFile);
sessionManager = guardSessionManager(SessionManager.open(params.sessionFile));
sessionManager = guardSessionManager(SessionManager.open(params.sessionFile), {
agentId: sessionAgentId,
sessionKey: params.sessionKey,
});
trackSessionManagerAccess(params.sessionFile);
await prepareSessionManagerForRun({

View File

@@ -1,5 +1,6 @@
import type { SessionManager } from "@mariozechner/pi-coding-agent";
import { getGlobalHookRunner } from "../plugins/hook-runner-global.js";
import { installSessionToolResultGuard } from "./session-tool-result-guard.js";
export type GuardedSessionManager = SessionManager & {
@@ -11,12 +12,38 @@ export type GuardedSessionManager = SessionManager & {
* Apply the tool-result guard to a SessionManager exactly once and expose
* a flush method on the instance for easy teardown handling.
*/
export function guardSessionManager(sessionManager: SessionManager): GuardedSessionManager {
export function guardSessionManager(
sessionManager: SessionManager,
opts?: { agentId?: string; sessionKey?: string },
): GuardedSessionManager {
if (typeof (sessionManager as GuardedSessionManager).flushPendingToolResults === "function") {
return sessionManager as GuardedSessionManager;
}
const guard = installSessionToolResultGuard(sessionManager);
const hookRunner = getGlobalHookRunner();
const transform = hookRunner?.hasHooks("tool_result_persist")
? (message: any, meta: { toolCallId?: string; toolName?: string; isSynthetic?: boolean }) => {
const out = hookRunner.runToolResultPersist(
{
toolName: meta.toolName,
toolCallId: meta.toolCallId,
message,
isSynthetic: meta.isSynthetic,
},
{
agentId: opts?.agentId,
sessionKey: opts?.sessionKey,
toolName: meta.toolName,
toolCallId: meta.toolCallId,
},
);
return out?.message ?? message;
}
: undefined;
const guard = installSessionToolResultGuard(sessionManager, {
transformToolResultForPersistence: transform,
});
(sessionManager as GuardedSessionManager).flushPendingToolResults = guard.flushPendingToolResults;
return sessionManager as GuardedSessionManager;
}

View File

@@ -0,0 +1,133 @@
import fs from "node:fs";
import os from "node:os";
import path from "node:path";
import type { AgentMessage } from "@mariozechner/pi-agent-core";
import { SessionManager } from "@mariozechner/pi-coding-agent";
import { describe, expect, it, afterEach } from "vitest";
import { loadClawdbotPlugins } from "../plugins/loader.js";
import { resetGlobalHookRunner } from "../plugins/hook-runner-global.js";
import { guardSessionManager } from "./session-tool-result-guard-wrapper.js";
const EMPTY_CONFIG_SCHEMA = `configSchema: {
validate: () => ({ ok: true }),
jsonSchema: { type: "object", additionalProperties: true },
uiHints: {}
}`;
function writeTempPlugin(params: { dir: string; id: string; body: string }): string {
const file = path.join(params.dir, `${params.id}.mjs`);
fs.writeFileSync(file, params.body, "utf-8");
return file;
}
afterEach(() => {
resetGlobalHookRunner();
});
describe("tool_result_persist hook", () => {
it("does not modify persisted toolResult messages when no hook is registered", () => {
const sm = guardSessionManager(SessionManager.inMemory(), {
agentId: "main",
sessionKey: "main",
});
sm.appendMessage({
role: "assistant",
content: [{ type: "toolCall", id: "call_1", name: "read", arguments: {} }],
} as AgentMessage);
sm.appendMessage({
role: "toolResult",
toolCallId: "call_1",
isError: false,
content: [{ type: "text", text: "ok" }],
details: { big: "x".repeat(10_000) },
} as any);
const messages = sm
.getEntries()
.filter((e) => e.type === "message")
.map((e) => (e as { message: AgentMessage }).message);
const toolResult = messages.find((m) => (m as any).role === "toolResult") as any;
expect(toolResult).toBeTruthy();
expect(toolResult.details).toBeTruthy();
});
it("composes transforms in priority order and allows stripping toolResult.details", () => {
const tmp = fs.mkdtempSync(path.join(os.tmpdir(), "clawdbot-toolpersist-"));
process.env.CLAWDBOT_BUNDLED_PLUGINS_DIR = "/nonexistent/bundled/plugins";
const pluginA = writeTempPlugin({
dir: tmp,
id: "persist-a",
body: `export default { id: "persist-a", ${EMPTY_CONFIG_SCHEMA}, register(api) {
api.on("tool_result_persist", (event, ctx) => {
const msg = event.message;
// Example: remove large diagnostic payloads before persistence.
const { details: _details, ...rest } = msg;
return { message: { ...rest, persistOrder: ["a"], agentSeen: ctx.agentId ?? null } };
}, { priority: 10 });
} };`,
});
const pluginB = writeTempPlugin({
dir: tmp,
id: "persist-b",
body: `export default { id: "persist-b", ${EMPTY_CONFIG_SCHEMA}, register(api) {
api.on("tool_result_persist", (event) => {
const prior = (event.message && event.message.persistOrder) ? event.message.persistOrder : [];
return { message: { ...event.message, persistOrder: [...prior, "b"] } };
}, { priority: 5 });
} };`,
});
loadClawdbotPlugins({
cache: false,
workspaceDir: tmp,
config: {
plugins: {
load: { paths: [pluginA, pluginB] },
allow: ["persist-a", "persist-b"],
},
},
});
const sm = guardSessionManager(SessionManager.inMemory(), {
agentId: "main",
sessionKey: "main",
});
// Tool call (so the guard can infer tool name -> id mapping).
sm.appendMessage({
role: "assistant",
content: [{ type: "toolCall", id: "call_1", name: "read", arguments: {} }],
} as AgentMessage);
// Tool result containing a large-ish details payload.
sm.appendMessage({
role: "toolResult",
toolCallId: "call_1",
isError: false,
content: [{ type: "text", text: "ok" }],
details: { big: "x".repeat(10_000) },
} as any);
const messages = sm
.getEntries()
.filter((e) => e.type === "message")
.map((e) => (e as { message: AgentMessage }).message);
const toolResult = messages.find((m) => (m as any).role === "toolResult") as any;
expect(toolResult).toBeTruthy();
// Default behavior: strip details.
expect(toolResult.details).toBeUndefined();
// Hook composition: priority 10 runs before priority 5.
expect(toolResult.persistOrder).toEqual(["a", "b"]);
expect(toolResult.agentSeen).toBe("main");
});
});

View File

@@ -68,17 +68,44 @@ function extractToolResultId(msg: Extract<AgentMessage, { role: "toolResult" }>)
return null;
}
export function installSessionToolResultGuard(sessionManager: SessionManager): {
export function installSessionToolResultGuard(
sessionManager: SessionManager,
opts?: {
/**
* Optional, synchronous transform applied to toolResult messages *before* they are
* persisted to the session transcript.
*/
transformToolResultForPersistence?: (
message: AgentMessage,
meta: { toolCallId?: string; toolName?: string; isSynthetic?: boolean },
) => AgentMessage;
},
): {
flushPendingToolResults: () => void;
getPendingIds: () => string[];
} {
const originalAppend = sessionManager.appendMessage.bind(sessionManager);
const pending = new Map<string, string | undefined>();
const persistToolResult = (
message: AgentMessage,
meta: { toolCallId?: string; toolName?: string; isSynthetic?: boolean },
) => {
const transformer = opts?.transformToolResultForPersistence;
return transformer ? transformer(message, meta) : message;
};
const flushPendingToolResults = () => {
if (pending.size === 0) return;
for (const [id, name] of pending.entries()) {
originalAppend(makeMissingToolResult({ toolCallId: id, toolName: name }));
const synthetic = makeMissingToolResult({ toolCallId: id, toolName: name });
originalAppend(
persistToolResult(synthetic, {
toolCallId: id,
toolName: name,
isSynthetic: true,
}) as never,
);
}
pending.clear();
};
@@ -88,8 +115,15 @@ export function installSessionToolResultGuard(sessionManager: SessionManager): {
if (role === "toolResult") {
const id = extractToolResultId(message as Extract<AgentMessage, { role: "toolResult" }>);
const toolName = id ? pending.get(id) : undefined;
if (id) pending.delete(id);
return originalAppend(message as never);
return originalAppend(
persistToolResult(message, {
toolCallId: id ?? undefined,
toolName,
isSynthetic: false,
}) as never,
);
}
const sanitized =