co-mono/packages/coding-agent/examples/hooks/tools.ts

/**
 * Tools Hook
 *
 * Provides a /tools command to enable/disable tools interactively.
 * Tool selection persists across session reloads and respects branch navigation.
 *
 * Usage:
 * 1. Copy this file to ~/.pi/agent/hooks/ or your project's .pi/hooks/
 * 2. Use /tools to open the tool selector
 */

import type { HookAPI, HookContext } from "@mariozechner/pi-coding-agent/hooks";
import { Container, type SettingItem, SettingsList, type SettingsListTheme } from "@mariozechner/pi-tui";

// State persisted to session
interface ToolsState {
	enabledTools: string[];
}

export default function toolsHook(pi: HookAPI) {
	// Track enabled tools
	let enabledTools: Set<string> = new Set();
	let allTools: string[] = [];

	// Persist current state
	function persistState() {
		pi.appendEntry<ToolsState>("tools-config", {
			enabledTools: Array.from(enabledTools),
		});
	}

	// Apply current tool selection
	function applyTools() {
		pi.setActiveTools(Array.from(enabledTools));
	}

	// Find the last tools-config entry in the current branch
	function restoreFromBranch(ctx: HookContext) {
		allTools = pi.getAllTools();

		// Get entries in current branch only
		const branchEntries = ctx.sessionManager.getBranch();
		let savedTools: string[] | undefined;

		for (const entry of branchEntries) {
			if (entry.type === "custom" && entry.customType === "tools-config") {
				const data = entry.data as ToolsState | undefined;
				if (data?.enabledTools) {
					savedTools = data.enabledTools;
				}
			}
		}

		if (savedTools) {
			// Restore saved tool selection (filter to only tools that still exist)
			enabledTools = new Set(savedTools.filter((t: string) => allTools.includes(t)));
			applyTools();
		} else {
			// No saved state - enable all tools by default
			enabledTools = new Set(allTools);
		}
	}

	// Register /tools command
	pi.registerCommand("tools", {
		description: "Enable/disable tools",
		handler: async (_args, ctx) => {
			// Refresh tool list
			allTools = pi.getAllTools();

			await ctx.ui.custom((tui, theme, done) => {
				// Build settings items for each tool
				const items: SettingItem[] = allTools.map((tool) => ({
					id: tool,
					label: tool,
					currentValue: enabledTools.has(tool) ? "enabled" : "disabled",
					values: ["enabled", "disabled"],
				}));

				const container = new Container();
				container.addChild(
					new (class {
						render(_width: number) {
							return [theme.fg("accent", theme.bold("Tool Configuration")), ""];
						}
						invalidate() {}
					})(),
				);

				// Create theme for SettingsList using the passed theme
				const settingsTheme: SettingsListTheme = {
					label: (text: string, selected: boolean) => (selected ? theme.fg("accent", text) : text),
					value: (text: string, selected: boolean) =>
						selected ? theme.fg("accent", text) : theme.fg("muted", text),
					description: (text: string) => theme.fg("dim", text),
					cursor: theme.fg("accent", "→ "),
					hint: (text: string) => theme.fg("dim", text),
				};

				const settingsList = new SettingsList(
					items,
					Math.min(items.length + 2, 15),
					settingsTheme,
					(id, newValue) => {
						// Update enabled state and apply immediately
						if (newValue === "enabled") {
							enabledTools.add(id);
						} else {
							enabledTools.delete(id);
						}
						applyTools();
						persistState();
					},
					() => {
						// Close dialog
						done(undefined);
					},
				);

				container.addChild(settingsList);

				const component = {
					render(width: number) {
						return container.render(width);
					},
					invalidate() {
						container.invalidate();
					},
					handleInput(data: string) {
						settingsList.handleInput?.(data);
						tui.requestRender();
					},
				};

				return component;
			});
		},
	});

	// Restore state on session start
	pi.on("session_start", async (_event, ctx) => {
		restoreFromBranch(ctx);
	});

	// Restore state when navigating the session tree
	pi.on("session_tree", async (_event, ctx) => {
		restoreFromBranch(ctx);
	});

	// Restore state after branching
	pi.on("session_branch", async (_event, ctx) => {
		restoreFromBranch(ctx);
	});
}