diff --git a/docs/concepts/agent-loop.md b/docs/concepts/agent-loop.md index 7052cf662..cd71384f2 100644 --- a/docs/concepts/agent-loop.md +++ b/docs/concepts/agent-loop.md @@ -73,6 +73,7 @@ These run inside the agent loop or gateway pipeline: - **`agent_end`**: inspect the final message list and run metadata after completion. - **`before_compaction` / `after_compaction`**: observe or annotate compaction cycles. - **`before_tool_call` / `after_tool_call`**: intercept tool params/results. +- **`tool_result_persist`**: synchronously transform tool results before they are written to the session transcript. - **`message_received` / `message_sending` / `message_sent`**: inbound + outbound message hooks. - **`session_start` / `session_end`**: session lifecycle boundaries. - **`gateway_start` / `gateway_stop`**: gateway lifecycle events. diff --git a/src/agents/pi-embedded-runner/compact.ts b/src/agents/pi-embedded-runner/compact.ts index 56a595892..e70ab1d0f 100644 --- a/src/agents/pi-embedded-runner/compact.ts +++ b/src/agents/pi-embedded-runner/compact.ts @@ -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({ diff --git a/src/agents/pi-embedded-runner/run/attempt.ts b/src/agents/pi-embedded-runner/run/attempt.ts index 06272d7ce..9c679771a 100644 --- a/src/agents/pi-embedded-runner/run/attempt.ts +++ b/src/agents/pi-embedded-runner/run/attempt.ts @@ -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({ diff --git a/src/agents/session-tool-result-guard-wrapper.ts b/src/agents/session-tool-result-guard-wrapper.ts index d386f85af..d83d14829 100644 --- a/src/agents/session-tool-result-guard-wrapper.ts +++ b/src/agents/session-tool-result-guard-wrapper.ts @@ -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; } diff --git a/src/agents/session-tool-result-guard.tool-result-persist-hook.test.ts b/src/agents/session-tool-result-guard.tool-result-persist-hook.test.ts new file mode 100644 index 000000000..ac6aea396 --- /dev/null +++ b/src/agents/session-tool-result-guard.tool-result-persist-hook.test.ts @@ -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"); + }); +}); diff --git a/src/agents/session-tool-result-guard.ts b/src/agents/session-tool-result-guard.ts index 1d3c5f846..35a00e0e3 100644 --- a/src/agents/session-tool-result-guard.ts +++ b/src/agents/session-tool-result-guard.ts @@ -68,17 +68,44 @@ function extractToolResultId(msg: Extract) 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(); + 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); + 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 = diff --git a/src/plugins/hooks.ts b/src/plugins/hooks.ts index fa0591e12..41fd501f7 100644 --- a/src/plugins/hooks.ts +++ b/src/plugins/hooks.ts @@ -30,6 +30,9 @@ import type { PluginHookSessionEndEvent, PluginHookSessionStartEvent, PluginHookToolContext, + PluginHookToolResultPersistContext, + PluginHookToolResultPersistEvent, + PluginHookToolResultPersistResult, } from "./types.js"; // Re-export types for consumers @@ -49,6 +52,9 @@ export type { PluginHookBeforeToolCallEvent, PluginHookBeforeToolCallResult, PluginHookAfterToolCallEvent, + PluginHookToolResultPersistContext, + PluginHookToolResultPersistEvent, + PluginHookToolResultPersistResult, PluginHookSessionContext, PluginHookSessionStartEvent, PluginHookSessionEndEvent, @@ -302,6 +308,59 @@ export function createHookRunner(registry: PluginRegistry, options: HookRunnerOp return runVoidHook("after_tool_call", event, ctx); } + /** + * Run tool_result_persist hook. + * + * This hook is intentionally synchronous: it runs in hot paths where session + * transcripts are appended synchronously. + * + * Handlers are executed sequentially in priority order (higher first). Each + * handler may return `{ message }` to replace the message passed to the next + * handler. + */ + function runToolResultPersist( + event: PluginHookToolResultPersistEvent, + ctx: PluginHookToolResultPersistContext, + ): PluginHookToolResultPersistResult | undefined { + const hooks = getHooksForName(registry, "tool_result_persist"); + if (hooks.length === 0) return undefined; + + let current = event.message; + + for (const hook of hooks) { + try { + const out = (hook.handler as any)({ ...event, message: current }, ctx) as + | PluginHookToolResultPersistResult + | void + | Promise; + + // Guard against accidental async handlers (this hook is sync-only). + if (out && typeof (out as any).then === "function") { + const msg = + `[hooks] tool_result_persist 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 next = (out as PluginHookToolResultPersistResult | undefined)?.message; + if (next) current = next; + } catch (err) { + const msg = `[hooks] tool_result_persist handler from ${hook.pluginId} failed: ${String(err)}`; + if (catchErrors) { + logger?.error(msg); + } else { + throw new Error(msg); + } + } + } + + return { message: current }; + } + // ========================================================================= // Session Hooks // ========================================================================= @@ -385,6 +444,7 @@ export function createHookRunner(registry: PluginRegistry, options: HookRunnerOp // Tool hooks runBeforeToolCall, runAfterToolCall, + runToolResultPersist, // Session hooks runSessionStart, runSessionEnd, diff --git a/src/plugins/types.ts b/src/plugins/types.ts index ba503338b..d13f6c944 100644 --- a/src/plugins/types.ts +++ b/src/plugins/types.ts @@ -1,6 +1,8 @@ import type { IncomingMessage, ServerResponse } from "node:http"; import type { Command } from "commander"; +import type { AgentMessage } from "@mariozechner/pi-agent-core"; + import type { AuthProfileCredential, OAuthCredential } from "../agents/auth-profiles/types.js"; import type { AnyAgentTool } from "../agents/tools/common.js"; import type { ChannelDock } from "../channels/dock.js"; @@ -231,6 +233,7 @@ export type PluginHookName = | "message_sent" | "before_tool_call" | "after_tool_call" + | "tool_result_persist" | "session_start" | "session_end" | "gateway_start" @@ -338,6 +341,30 @@ export type PluginHookAfterToolCallEvent = { durationMs?: number; }; +// tool_result_persist hook +export type PluginHookToolResultPersistContext = { + agentId?: string; + sessionKey?: string; + toolName?: string; + toolCallId?: string; +}; + +export type PluginHookToolResultPersistEvent = { + toolName?: string; + toolCallId?: string; + /** + * The toolResult message about to be written to the session transcript. + * Handlers may return a modified message (e.g. drop non-essential fields). + */ + message: AgentMessage; + /** True when the tool result was synthesized by a guard/repair step. */ + isSynthetic?: boolean; +}; + +export type PluginHookToolResultPersistResult = { + message?: AgentMessage; +}; + // Session context export type PluginHookSessionContext = { agentId?: string; @@ -407,6 +434,10 @@ export type PluginHookHandlerMap = { event: PluginHookAfterToolCallEvent, ctx: PluginHookToolContext, ) => Promise | void; + tool_result_persist: ( + event: PluginHookToolResultPersistEvent, + ctx: PluginHookToolResultPersistContext, + ) => PluginHookToolResultPersistResult | void; session_start: ( event: PluginHookSessionStartEvent, ctx: PluginHookSessionContext,