WIP: Add hook API for dynamic tool control with plan-mode hook example

- Add pi.getTools() and pi.setTools(toolNames) to HookAPI
- Hooks can now enable/disable tools dynamically
- Changes take effect on next agent turn

New example hook: plan-mode.ts
- Claude Code-style read-only exploration mode
- /plan command toggles plan mode on/off
- Plan mode tools: read, bash, grep, find, ls
- Edit/write tools disabled in plan mode
- Injects context telling agent about restrictions
- After each response, prompts to execute/stay/refine
- State persists across sessions
This commit is contained in:
Helmut Januschka 2026-01-03 09:31:39 +01:00
parent 6ddfd1be13
commit 57bba4e32b
14 changed files with 304 additions and 8 deletions

View file

@ -13,7 +13,14 @@
* Modes use this class and add their own I/O layer on top.
*/
import type { Agent, AgentEvent, AgentMessage, AgentState, ThinkingLevel } from "@mariozechner/pi-agent-core";
import type {
Agent,
AgentEvent,
AgentMessage,
AgentState,
AgentTool,
ThinkingLevel,
} from "@mariozechner/pi-agent-core";
import type { AssistantMessage, ImageContent, Message, Model, TextContent } from "@mariozechner/pi-ai";
import { isContextOverflow, modelsAreEqual, supportsXhigh } from "@mariozechner/pi-ai";
import { getAuthPath } from "../config.js";
@ -75,6 +82,8 @@ export interface AgentSessionConfig {
skillsSettings?: Required<SkillsSettings>;
/** Model registry for API key resolution and model discovery */
modelRegistry: ModelRegistry;
/** Tool registry for hook getTools/setTools - maps name to tool */
toolRegistry?: Map<string, AgentTool>;
}
/** Options for AgentSession.prompt() */
@ -174,6 +183,9 @@ export class AgentSession {
// Model registry for API key resolution
private _modelRegistry: ModelRegistry;
// Tool registry for hook getTools/setTools
private _toolRegistry: Map<string, AgentTool>;
constructor(config: AgentSessionConfig) {
this.agent = config.agent;
this.sessionManager = config.sessionManager;
@ -184,6 +196,7 @@ export class AgentSession {
this._customTools = config.customTools ?? [];
this._skillsSettings = config.skillsSettings;
this._modelRegistry = config.modelRegistry;
this._toolRegistry = config.toolRegistry ?? new Map();
// Always subscribe to agent events for internal handling
// (session persistence, hooks, auto-compaction, retry logic)
@ -417,6 +430,30 @@ export class AgentSession {
return this.agent.state.isStreaming;
}
/**
* Get the names of currently active tools.
* Returns the names of tools currently set on the agent.
*/
getActiveToolNames(): string[] {
return this.agent.state.tools.map((t) => t.name);
}
/**
* Set active tools by name.
* Only tools in the registry can be enabled. Unknown tool names are ignored.
* Changes take effect on the next agent turn.
*/
setActiveToolsByName(toolNames: string[]): void {
const tools: AgentTool[] = [];
for (const name of toolNames) {
const tool = this._toolRegistry.get(name);
if (tool) {
tools.push(tool);
}
}
this.agent.setTools(tools);
}
/** Whether auto-compaction is currently running */
get isCompacting(): boolean {
return this._autoCompactionAbortController !== undefined || this._compactionAbortController !== undefined;

View file

@ -4,11 +4,13 @@ export {
loadHooks,
type AppendEntryHandler,
type BranchHandler,
type GetToolsHandler,
type LoadedHook,
type LoadHooksResult,
type NavigateTreeHandler,
type NewSessionHandler,
type SendMessageHandler,
type SetToolsHandler,
} from "./loader.js";
export { execCommand, HookRunner, type HookErrorListener } from "./runner.js";
export { wrapToolsWithHooks, wrapToolWithHooks } from "./tool-wrapper.js";

View file

@ -61,6 +61,16 @@ export type SendMessageHandler = <T = unknown>(
*/
export type AppendEntryHandler = <T = unknown>(customType: string, data?: T) => void;
/**
* Get tools handler type for pi.getTools().
*/
export type GetToolsHandler = () => string[];
/**
* Set tools handler type for pi.setTools().
*/
export type SetToolsHandler = (toolNames: string[]) => void;
/**
* New session handler type for ctx.newSession() in HookCommandContext.
*/
@ -100,6 +110,10 @@ export interface LoadedHook {
setSendMessageHandler: (handler: SendMessageHandler) => void;
/** Set the append entry handler for this hook's pi.appendEntry() */
setAppendEntryHandler: (handler: AppendEntryHandler) => void;
/** Set the get tools handler for this hook's pi.getTools() */
setGetToolsHandler: (handler: GetToolsHandler) => void;
/** Set the set tools handler for this hook's pi.setTools() */
setSetToolsHandler: (handler: SetToolsHandler) => void;
}
/**
@ -159,6 +173,8 @@ function createHookAPI(
commands: Map<string, RegisteredCommand>;
setSendMessageHandler: (handler: SendMessageHandler) => void;
setAppendEntryHandler: (handler: AppendEntryHandler) => void;
setGetToolsHandler: (handler: GetToolsHandler) => void;
setSetToolsHandler: (handler: SetToolsHandler) => void;
} {
let sendMessageHandler: SendMessageHandler = () => {
// Default no-op until mode sets the handler
@ -166,6 +182,10 @@ function createHookAPI(
let appendEntryHandler: AppendEntryHandler = () => {
// Default no-op until mode sets the handler
};
let getToolsHandler: GetToolsHandler = () => [];
let setToolsHandler: SetToolsHandler = () => {
// Default no-op until mode sets the handler
};
const messageRenderers = new Map<string, HookMessageRenderer>();
const commands = new Map<string, RegisteredCommand>();
@ -195,6 +215,12 @@ function createHookAPI(
exec(command: string, args: string[], options?: ExecOptions) {
return execCommand(command, args, options?.cwd ?? cwd, options);
},
getTools(): string[] {
return getToolsHandler();
},
setTools(toolNames: string[]): void {
setToolsHandler(toolNames);
},
} as HookAPI;
return {
@ -207,6 +233,12 @@ function createHookAPI(
setAppendEntryHandler: (handler: AppendEntryHandler) => {
appendEntryHandler = handler;
},
setGetToolsHandler: (handler: GetToolsHandler) => {
getToolsHandler = handler;
},
setSetToolsHandler: (handler: SetToolsHandler) => {
setToolsHandler = handler;
},
};
}
@ -234,10 +266,15 @@ async function loadHook(hookPath: string, cwd: string): Promise<{ hook: LoadedHo
// Create handlers map and API
const handlers = new Map<string, HandlerFn[]>();
const { api, messageRenderers, commands, setSendMessageHandler, setAppendEntryHandler } = createHookAPI(
handlers,
cwd,
);
const {
api,
messageRenderers,
commands,
setSendMessageHandler,
setAppendEntryHandler,
setGetToolsHandler,
setSetToolsHandler,
} = createHookAPI(handlers, cwd);
// Call factory to register handlers
factory(api);
@ -251,6 +288,8 @@ async function loadHook(hookPath: string, cwd: string): Promise<{ hook: LoadedHo
commands,
setSendMessageHandler,
setAppendEntryHandler,
setGetToolsHandler,
setSetToolsHandler,
},
error: null,
};

View file

@ -98,6 +98,10 @@ export class HookRunner {
sendMessageHandler: SendMessageHandler;
/** Handler for hooks to append entries */
appendEntryHandler: AppendEntryHandler;
/** Handler for getting current tools */
getToolsHandler: () => string[];
/** Handler for setting tools */
setToolsHandler: (toolNames: string[]) => void;
/** Handler for creating new sessions (for HookCommandContext) */
newSessionHandler?: NewSessionHandler;
/** Handler for branching sessions (for HookCommandContext) */
@ -132,10 +136,12 @@ export class HookRunner {
if (options.navigateTreeHandler) {
this.navigateTreeHandler = options.navigateTreeHandler;
}
// Set per-hook handlers for pi.sendMessage() and pi.appendEntry()
// Set per-hook handlers for pi.sendMessage(), pi.appendEntry(), pi.getTools(), pi.setTools()
for (const hook of this.hooks) {
hook.setSendMessageHandler(options.sendMessageHandler);
hook.setAppendEntryHandler(options.appendEntryHandler);
hook.setGetToolsHandler(options.getToolsHandler);
hook.setSetToolsHandler(options.setToolsHandler);
}
this.uiContext = options.uiContext ?? noOpUIContext;
this.hasUI = options.hasUI ?? false;

View file

@ -749,6 +749,28 @@ export interface HookAPI {
* Supports timeout and abort signal.
*/
exec(command: string, args: string[], options?: ExecOptions): Promise<ExecResult>;
/**
* Get the list of currently active tool names.
* @returns Array of tool names (e.g., ["read", "bash", "edit", "write"])
*/
getTools(): string[];
/**
* Set the active tools by name.
* Only built-in tools can be enabled/disabled. Custom tools are always active.
* Changes take effect on the next agent turn.
*
* @param toolNames - Array of tool names to enable (e.g., ["read", "bash", "grep", "find", "ls"])
*
* @example
* // Switch to read-only mode (plan mode)
* pi.setTools(["read", "bash", "grep", "find", "ls"]);
*
* // Restore full access
* pi.setTools(["read", "bash", "edit", "write"]);
*/
setTools(toolNames: string[]): void;
}
/**

View file

@ -29,7 +29,7 @@
* ```
*/
import { Agent, type ThinkingLevel } from "@mariozechner/pi-agent-core";
import { Agent, type AgentTool, type ThinkingLevel } from "@mariozechner/pi-agent-core";
import type { Model } from "@mariozechner/pi-ai";
import { join } from "path";
import { getAgentDir } from "../config.js";
@ -349,6 +349,8 @@ function createLoadedHooksFromDefinitions(definitions: Array<{ path?: string; fa
options?: { triggerTurn?: boolean; deliverAs?: "steer" | "followUp" },
) => void = () => {};
let appendEntryHandler: (customType: string, data?: any) => void = () => {};
let getToolsHandler: () => string[] = () => [];
let setToolsHandler: (toolNames: string[]) => void = () => {};
let newSessionHandler: (options?: any) => Promise<{ cancelled: boolean }> = async () => ({ cancelled: false });
let branchHandler: (entryId: string) => Promise<{ cancelled: boolean }> = async () => ({ cancelled: false });
let navigateTreeHandler: (targetId: string, options?: any) => Promise<{ cancelled: boolean }> = async () => ({
@ -376,6 +378,8 @@ function createLoadedHooksFromDefinitions(definitions: Array<{ path?: string; fa
newSession: (options?: any) => newSessionHandler(options),
branch: (entryId: string) => branchHandler(entryId),
navigateTree: (targetId: string, options?: any) => navigateTreeHandler(targetId, options),
getTools: () => getToolsHandler(),
setTools: (toolNames: string[]) => setToolsHandler(toolNames),
};
def.factory(api as any);
@ -403,6 +407,12 @@ function createLoadedHooksFromDefinitions(definitions: Array<{ path?: string; fa
setNavigateTreeHandler: (handler: (targetId: string, options?: any) => Promise<{ cancelled: boolean }>) => {
navigateTreeHandler = handler;
},
setGetToolsHandler: (handler: () => string[]) => {
getToolsHandler = handler;
},
setSetToolsHandler: (handler: (toolNames: string[]) => void) => {
setToolsHandler = handler;
},
};
});
}
@ -588,10 +598,28 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
},
}));
// Create tool registry mapping name -> tool (for hook getTools/setTools)
// Cast to AgentTool since createCodingTools actually returns AgentTool[] (type is just Tool[])
const toolRegistry = new Map<string, AgentTool>();
for (const tool of builtInTools as AgentTool[]) {
toolRegistry.set(tool.name, tool);
}
for (const tool of wrappedCustomTools as AgentTool[]) {
toolRegistry.set(tool.name, tool);
}
let allToolsArray: Tool[] = [...builtInTools, ...wrappedCustomTools];
time("combineTools");
// Wrap tools with hooks if available
let wrappedToolRegistry: Map<string, AgentTool> | undefined;
if (hookRunner) {
allToolsArray = wrapToolsWithHooks(allToolsArray, hookRunner) as Tool[];
allToolsArray = wrapToolsWithHooks(allToolsArray as AgentTool[], hookRunner);
// Also create a wrapped version of the registry for setTools
wrappedToolRegistry = new Map<string, AgentTool>();
for (const tool of allToolsArray as AgentTool[]) {
wrappedToolRegistry.set(tool.name, tool);
}
}
let systemPrompt: string;
@ -670,6 +698,7 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
customTools: customToolsResult.tools,
skillsSettings: settingsManager.getSkillsSettings(),
modelRegistry,
toolRegistry: wrappedToolRegistry ?? toolRegistry,
});
time("createAgentSession");

View file

@ -457,6 +457,8 @@ export class InteractiveMode {
appendEntryHandler: (customType, data) => {
this.sessionManager.appendCustomEntry(customType, data);
},
getToolsHandler: () => this.session.getActiveToolNames(),
setToolsHandler: (toolNames) => this.session.setActiveToolsByName(toolNames),
newSessionHandler: async (options) => {
// Stop any loading animation
if (this.loadingAnimation) {

View file

@ -40,6 +40,8 @@ export async function runPrintMode(
appendEntryHandler: (customType, data) => {
session.sessionManager.appendCustomEntry(customType, data);
},
getToolsHandler: () => session.getActiveToolNames(),
setToolsHandler: (toolNames) => session.setActiveToolsByName(toolNames),
});
hookRunner.onError((err) => {
console.error(`Hook error (${err.hookPath}): ${err.error}`);

View file

@ -189,6 +189,8 @@ export async function runRpcMode(session: AgentSession): Promise<never> {
appendEntryHandler: (customType, data) => {
session.sessionManager.appendCustomEntry(customType, data);
},
getToolsHandler: () => session.getActiveToolNames(),
setToolsHandler: (toolNames) => session.setActiveToolsByName(toolNames),
uiContext: createHookUIContext(),
hasUI: false,
});