50 lines
1.8 KiB
TypeScript
50 lines
1.8 KiB
TypeScript
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 & {
|
|
/** Flush any synthetic tool results for pending tool calls. Idempotent. */
|
|
flushPendingToolResults?: () => void;
|
|
};
|
|
|
|
/**
|
|
* 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,
|
|
opts?: { agentId?: string; sessionKey?: string },
|
|
): GuardedSessionManager {
|
|
if (typeof (sessionManager as GuardedSessionManager).flushPendingToolResults === "function") {
|
|
return sessionManager as GuardedSessionManager;
|
|
}
|
|
|
|
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;
|
|
}
|