Merge bdf04ab426 into 09be5d45d5
This commit is contained in:
commit
f1cdf9d503
@ -75,6 +75,7 @@ import { prewarmSessionFile, trackSessionManagerAccess } from "../session-manage
|
|||||||
import { prepareSessionManagerForRun } from "../session-manager-init.js";
|
import { prepareSessionManagerForRun } from "../session-manager-init.js";
|
||||||
import { buildEmbeddedSystemPrompt, createSystemPromptOverride } from "../system-prompt.js";
|
import { buildEmbeddedSystemPrompt, createSystemPromptOverride } from "../system-prompt.js";
|
||||||
import { splitSdkTools } from "../tool-split.js";
|
import { splitSdkTools } from "../tool-split.js";
|
||||||
|
import { wrapToolsWithHooks } from "../../pi-tools.hooks.js";
|
||||||
import { toClientToolDefinitions } from "../../pi-tool-definition-adapter.js";
|
import { toClientToolDefinitions } from "../../pi-tool-definition-adapter.js";
|
||||||
import { buildSystemPromptParams } from "../../system-prompt-params.js";
|
import { buildSystemPromptParams } from "../../system-prompt-params.js";
|
||||||
import { describeUnknownError, mapThinkingLevel } from "../utils.js";
|
import { describeUnknownError, mapThinkingLevel } from "../utils.js";
|
||||||
@ -432,6 +433,10 @@ export async function runEmbeddedAttempt(
|
|||||||
model: params.model,
|
model: params.model,
|
||||||
});
|
});
|
||||||
|
|
||||||
|
// Get hook runner once for tool-call hooks, before_agent_start, and agent_end
|
||||||
|
const hookRunner = getGlobalHookRunner();
|
||||||
|
const hookAgentId = params.sessionKey?.split(":")[0] ?? "main";
|
||||||
|
|
||||||
const { builtInTools, customTools } = splitSdkTools({
|
const { builtInTools, customTools } = splitSdkTools({
|
||||||
tools,
|
tools,
|
||||||
sandboxEnabled: !!sandbox?.enabled,
|
sandboxEnabled: !!sandbox?.enabled,
|
||||||
@ -447,6 +452,15 @@ export async function runEmbeddedAttempt(
|
|||||||
|
|
||||||
const allCustomTools = [...customTools, ...clientToolDefs];
|
const allCustomTools = [...customTools, ...clientToolDefs];
|
||||||
|
|
||||||
|
// Wrap tools with before_tool_call / after_tool_call plugin hooks
|
||||||
|
const hookCtx = { agentId: hookAgentId, sessionKey: params.sessionKey };
|
||||||
|
const hookedBuiltInTools = hookRunner
|
||||||
|
? wrapToolsWithHooks(builtInTools, hookRunner, hookCtx)
|
||||||
|
: builtInTools;
|
||||||
|
const hookedCustomTools = hookRunner
|
||||||
|
? wrapToolsWithHooks(allCustomTools, hookRunner, hookCtx)
|
||||||
|
: allCustomTools;
|
||||||
|
|
||||||
({ session } = await createAgentSession({
|
({ session } = await createAgentSession({
|
||||||
cwd: resolvedWorkspace,
|
cwd: resolvedWorkspace,
|
||||||
agentDir,
|
agentDir,
|
||||||
@ -455,8 +469,8 @@ export async function runEmbeddedAttempt(
|
|||||||
model: params.model,
|
model: params.model,
|
||||||
thinkingLevel: mapThinkingLevel(params.thinkLevel),
|
thinkingLevel: mapThinkingLevel(params.thinkLevel),
|
||||||
systemPrompt,
|
systemPrompt,
|
||||||
tools: builtInTools,
|
tools: hookedBuiltInTools,
|
||||||
customTools: allCustomTools,
|
customTools: hookedCustomTools,
|
||||||
sessionManager,
|
sessionManager,
|
||||||
settingsManager,
|
settingsManager,
|
||||||
skills: [],
|
skills: [],
|
||||||
@ -676,9 +690,6 @@ export async function runEmbeddedAttempt(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// Get hook runner once for both before_agent_start and agent_end hooks
|
|
||||||
const hookRunner = getGlobalHookRunner();
|
|
||||||
|
|
||||||
let promptError: unknown = null;
|
let promptError: unknown = null;
|
||||||
try {
|
try {
|
||||||
const promptStartedAt = Date.now();
|
const promptStartedAt = Date.now();
|
||||||
@ -693,7 +704,7 @@ export async function runEmbeddedAttempt(
|
|||||||
messages: activeSession.messages,
|
messages: activeSession.messages,
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
agentId: params.sessionKey?.split(":")[0] ?? "main",
|
agentId: hookAgentId,
|
||||||
sessionKey: params.sessionKey,
|
sessionKey: params.sessionKey,
|
||||||
workspaceDir: params.workspaceDir,
|
workspaceDir: params.workspaceDir,
|
||||||
messageProvider: params.messageProvider ?? undefined,
|
messageProvider: params.messageProvider ?? undefined,
|
||||||
@ -821,7 +832,7 @@ export async function runEmbeddedAttempt(
|
|||||||
durationMs: Date.now() - promptStartedAt,
|
durationMs: Date.now() - promptStartedAt,
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
agentId: params.sessionKey?.split(":")[0] ?? "main",
|
agentId: hookAgentId,
|
||||||
sessionKey: params.sessionKey,
|
sessionKey: params.sessionKey,
|
||||||
workspaceDir: params.workspaceDir,
|
workspaceDir: params.workspaceDir,
|
||||||
messageProvider: params.messageProvider ?? undefined,
|
messageProvider: params.messageProvider ?? undefined,
|
||||||
|
|||||||
115
src/agents/pi-tools.hooks.ts
Normal file
115
src/agents/pi-tools.hooks.ts
Normal file
@ -0,0 +1,115 @@
|
|||||||
|
/**
|
||||||
|
* Tool Hook Wrappers
|
||||||
|
*
|
||||||
|
* Wraps tool execute methods to fire before_tool_call and after_tool_call
|
||||||
|
* plugin hooks around every tool invocation.
|
||||||
|
*/
|
||||||
|
|
||||||
|
import type { HookRunner } from "../plugins/hooks.js";
|
||||||
|
import { createSubsystemLogger } from "../logging/subsystem.js";
|
||||||
|
import type { AnyAgentTool } from "./pi-tools.types.js";
|
||||||
|
|
||||||
|
const log = createSubsystemLogger("tools/hooks");
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Wrap a single tool's execute method to fire before_tool_call and after_tool_call hooks.
|
||||||
|
*/
|
||||||
|
export function wrapToolWithHooks(
|
||||||
|
tool: AnyAgentTool,
|
||||||
|
hookRunner: HookRunner,
|
||||||
|
ctx: { agentId?: string; sessionKey?: string },
|
||||||
|
): AnyAgentTool {
|
||||||
|
const originalExecute = tool.execute;
|
||||||
|
if (!originalExecute) return tool;
|
||||||
|
|
||||||
|
const toolName = tool.name;
|
||||||
|
|
||||||
|
return {
|
||||||
|
...tool,
|
||||||
|
execute: async (
|
||||||
|
toolCallId: string,
|
||||||
|
params: unknown,
|
||||||
|
signal?: AbortSignal,
|
||||||
|
onUpdate?: (data: unknown) => void,
|
||||||
|
) => {
|
||||||
|
const hookCtx = {
|
||||||
|
agentId: ctx.agentId,
|
||||||
|
sessionKey: ctx.sessionKey,
|
||||||
|
toolName,
|
||||||
|
};
|
||||||
|
|
||||||
|
// --- before_tool_call ---
|
||||||
|
let effectiveParams = params;
|
||||||
|
if (hookRunner.hasHooks("before_tool_call")) {
|
||||||
|
try {
|
||||||
|
const beforeResult = await hookRunner.runBeforeToolCall(
|
||||||
|
{
|
||||||
|
toolName,
|
||||||
|
params: (params ?? {}) as Record<string, unknown>,
|
||||||
|
},
|
||||||
|
hookCtx,
|
||||||
|
);
|
||||||
|
if (beforeResult?.block) {
|
||||||
|
const reason = beforeResult.blockReason ?? "Blocked by plugin hook";
|
||||||
|
log.debug(`before_tool_call: blocked ${toolName} — ${reason}`);
|
||||||
|
return `[Tool call blocked] ${reason}`;
|
||||||
|
}
|
||||||
|
if (beforeResult?.params) {
|
||||||
|
effectiveParams = beforeResult.params;
|
||||||
|
}
|
||||||
|
} catch (err) {
|
||||||
|
log.debug(`before_tool_call hook error for ${toolName}: ${String(err)}`);
|
||||||
|
// Hook errors must not break tool execution
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// --- execute ---
|
||||||
|
const startMs = Date.now();
|
||||||
|
let result: unknown;
|
||||||
|
let error: string | undefined;
|
||||||
|
try {
|
||||||
|
result = await originalExecute(toolCallId, effectiveParams, signal, onUpdate);
|
||||||
|
return result;
|
||||||
|
} catch (err) {
|
||||||
|
error = String(err);
|
||||||
|
throw err;
|
||||||
|
} finally {
|
||||||
|
// --- after_tool_call (fire-and-forget) ---
|
||||||
|
if (hookRunner.hasHooks("after_tool_call")) {
|
||||||
|
hookRunner
|
||||||
|
.runAfterToolCall(
|
||||||
|
{
|
||||||
|
toolName,
|
||||||
|
params: (effectiveParams ?? {}) as Record<string, unknown>,
|
||||||
|
result,
|
||||||
|
error,
|
||||||
|
durationMs: Date.now() - startMs,
|
||||||
|
},
|
||||||
|
hookCtx,
|
||||||
|
)
|
||||||
|
.catch((hookErr) => {
|
||||||
|
log.debug(`after_tool_call hook error for ${toolName}: ${String(hookErr)}`);
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
},
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Wrap all tools in an array with before/after tool call hooks.
|
||||||
|
* Returns the original array unchanged if no tool call hooks are registered.
|
||||||
|
*/
|
||||||
|
export function wrapToolsWithHooks(
|
||||||
|
tools: AnyAgentTool[],
|
||||||
|
hookRunner: HookRunner,
|
||||||
|
ctx: { agentId?: string; sessionKey?: string },
|
||||||
|
): AnyAgentTool[] {
|
||||||
|
if (
|
||||||
|
!hookRunner.hasHooks("before_tool_call") &&
|
||||||
|
!hookRunner.hasHooks("after_tool_call")
|
||||||
|
) {
|
||||||
|
return tools;
|
||||||
|
}
|
||||||
|
return tools.map((tool) => wrapToolWithHooks(tool, hookRunner, ctx));
|
||||||
|
}
|
||||||
Loading…
Reference in New Issue
Block a user