index.ts
1 /** 2 * Tools Extension 3 * 4 * Provides a /tools command to enable/disable tools interactively. 5 * Tool selection persists: 6 * - Globally in $PI_CODING_AGENT_DIR/extensions/tools/tools.json or ~/.pi/agent/extensions/tools/tools.json 7 * - Per-session via session entries (for branch-specific overrides) 8 * 9 * Usage: 10 * 1. Copy this folder (`tools/`) to ~/.pi/agent/extensions/ or your project's .pi/extensions/ 11 * 2. Use /tools to open the tool selector 12 */ 13 14 import { existsSync, mkdirSync, readFileSync, writeFileSync } from "node:fs"; 15 import { homedir } from "node:os"; 16 import { dirname, join } from "node:path"; 17 import type { ExtensionAPI, ExtensionContext, ToolInfo } from "@mariozechner/pi-coding-agent"; 18 import { getSettingsListTheme } from "@mariozechner/pi-coding-agent"; 19 import { Container, type SettingItem, SettingsList } from "@mariozechner/pi-tui"; 20 21 type ToolOverride = "enabled" | "disabled"; 22 23 type ToolsConfigV2 = { 24 version: 2; 25 overrides: Record<string, ToolOverride>; 26 }; 27 28 type ToolsConfigEntryLike = { 29 type?: unknown; 30 customType?: unknown; 31 data?: unknown; 32 }; 33 34 const TOOLS_CONFIG_TYPE = "tools-config"; 35 const EMPTY_TOOLS_CONFIG: ToolsConfigV2 = { version: 2, overrides: {} }; 36 37 function getConfigPath(): string { 38 const agentDir = process.env.PI_CODING_AGENT_DIR ?? join(homedir(), ".pi", "agent"); 39 return join(agentDir, "extensions", "tools", "tools.json"); 40 } 41 42 function isRecord(value: unknown): value is Record<string, unknown> { 43 return typeof value === "object" && value !== null && !Array.isArray(value); 44 } 45 46 function getToolNames(allTools: readonly ToolInfo[]): string[] { 47 return allTools.map((tool) => tool.name); 48 } 49 50 function serializeOverrides(overrides: ReadonlyMap<string, ToolOverride>): Record<string, ToolOverride> { 51 return Object.fromEntries([...overrides.entries()].sort(([left], [right]) => left.localeCompare(right))); 52 } 53 54 function toPersistedState( 55 overrides: ReadonlyMap<string, ToolOverride>, 56 availableTools: ReadonlySet<string>, 57 ): ToolsConfigV2 { 58 return { 59 version: 2, 60 overrides: serializeOverrides( 61 new Map([...overrides.entries()].filter(([toolName]) => availableTools.has(toolName))), 62 ), 63 }; 64 } 65 66 function overridesToMap(config: ToolsConfigV2): Map<string, ToolOverride> { 67 return new Map(Object.entries(config.overrides)); 68 } 69 70 function updateToolOverride( 71 toolName: string, 72 desiredEnabled: boolean, 73 baseActiveTools: ReadonlySet<string>, 74 overrides: Map<string, ToolOverride>, 75 ): void { 76 const baseEnabled = baseActiveTools.has(toolName); 77 if (desiredEnabled === baseEnabled) { 78 overrides.delete(toolName); 79 return; 80 } 81 82 overrides.set(toolName, desiredEnabled ? "enabled" : "disabled"); 83 } 84 85 function normalizePersistedState( 86 value: unknown, 87 baseActiveTools: ReadonlySet<string>, 88 availableTools: ReadonlySet<string>, 89 ): ToolsConfigV2 | undefined { 90 if (!isRecord(value)) return undefined; 91 92 if (value.version === 2) { 93 if (!isRecord(value.overrides)) return undefined; 94 95 const overrides: Record<string, ToolOverride> = {}; 96 for (const [toolName, state] of Object.entries(value.overrides)) { 97 if (state !== "enabled" && state !== "disabled") { 98 return undefined; 99 } 100 if (availableTools.has(toolName)) { 101 overrides[toolName] = state; 102 } 103 } 104 105 return { version: 2, overrides }; 106 } 107 108 if (!Array.isArray(value.enabledTools) || value.enabledTools.some((toolName) => typeof toolName !== "string")) { 109 return undefined; 110 } 111 112 const overrides: Record<string, ToolOverride> = {}; 113 for (const toolName of value.enabledTools) { 114 if (availableTools.has(toolName) && !baseActiveTools.has(toolName)) { 115 overrides[toolName] = "enabled"; 116 } 117 } 118 119 return { version: 2, overrides }; 120 } 121 122 function stripPreviouslyAppliedOverrides( 123 currentActiveTools: ReadonlySet<string>, 124 overrides: ReadonlyMap<string, ToolOverride>, 125 availableTools: ReadonlySet<string>, 126 ): Set<string> { 127 const baseActiveTools = new Set([...currentActiveTools].filter((toolName) => availableTools.has(toolName))); 128 129 for (const [toolName, state] of overrides.entries()) { 130 if (!availableTools.has(toolName)) continue; 131 if (state === "enabled") { 132 baseActiveTools.delete(toolName); 133 continue; 134 } 135 baseActiveTools.add(toolName); 136 } 137 138 return baseActiveTools; 139 } 140 141 function applyOverrides( 142 baseActiveTools: ReadonlySet<string>, 143 overrides: ReadonlyMap<string, ToolOverride>, 144 allTools: readonly ToolInfo[], 145 ): string[] { 146 const toolNames = getToolNames(allTools); 147 const availableTools = new Set(toolNames); 148 const effectiveTools = new Set([...baseActiveTools].filter((toolName) => availableTools.has(toolName))); 149 150 for (const [toolName, state] of overrides.entries()) { 151 if (!availableTools.has(toolName)) continue; 152 if (state === "enabled") { 153 effectiveTools.add(toolName); 154 continue; 155 } 156 effectiveTools.delete(toolName); 157 } 158 159 return toolNames.filter((toolName) => effectiveTools.has(toolName)); 160 } 161 162 function sameToolMembership(left: readonly string[], right: Iterable<string>): boolean { 163 const leftSet = new Set(left); 164 const rightSet = new Set(right); 165 166 if (leftSet.size !== rightSet.size) return false; 167 for (const toolName of leftSet) { 168 if (!rightSet.has(toolName)) return false; 169 } 170 return true; 171 } 172 173 function loadGlobalConfig( 174 baseActiveTools: ReadonlySet<string>, 175 availableTools: ReadonlySet<string>, 176 ): ToolsConfigV2 | undefined { 177 const configPath = getConfigPath(); 178 if (!existsSync(configPath)) return undefined; 179 180 try { 181 const content = readFileSync(configPath, "utf-8"); 182 return normalizePersistedState(JSON.parse(content), baseActiveTools, availableTools); 183 } catch { 184 return undefined; 185 } 186 } 187 188 function readLatestBranchConfig( 189 ctx: ExtensionContext, 190 baseActiveTools: ReadonlySet<string>, 191 availableTools: ReadonlySet<string>, 192 ): ToolsConfigV2 | undefined { 193 let latestValidConfig: ToolsConfigV2 | undefined; 194 195 for (const entry of ctx.sessionManager.getBranch()) { 196 const candidate = entry as ToolsConfigEntryLike; 197 if (candidate.type !== "custom" || candidate.customType !== TOOLS_CONFIG_TYPE) continue; 198 199 const normalized = normalizePersistedState(candidate.data, baseActiveTools, availableTools); 200 if (normalized) { 201 latestValidConfig = normalized; 202 } 203 } 204 205 return latestValidConfig; 206 } 207 208 export default function toolsExtension(pi: ExtensionAPI) { 209 let allTools: ToolInfo[] = []; 210 let toolOverrides = new Map<string, ToolOverride>(); 211 212 function getAvailableToolSet(): Set<string> { 213 return new Set(getToolNames(allTools)); 214 } 215 216 function saveGlobalConfig() { 217 const configPath = getConfigPath(); 218 const config = toPersistedState(toolOverrides, getAvailableToolSet()); 219 220 try { 221 const configDir = dirname(configPath); 222 if (!existsSync(configDir)) { 223 mkdirSync(configDir, { recursive: true }); 224 } 225 writeFileSync(configPath, JSON.stringify(config, null, 2)); 226 } catch (err) { 227 console.error(`Failed to save tools config: ${err}`); 228 } 229 } 230 231 function persistToSession() { 232 pi.appendEntry<ToolsConfigV2>(TOOLS_CONFIG_TYPE, toPersistedState(toolOverrides, getAvailableToolSet())); 233 } 234 235 function persistState() { 236 saveGlobalConfig(); 237 persistToSession(); 238 } 239 240 function restoreFromBranch(ctx: ExtensionContext) { 241 allTools = pi.getAllTools(); 242 const availableTools = getAvailableToolSet(); 243 const currentActiveTools = new Set(pi.getActiveTools()); 244 const baseActiveTools = stripPreviouslyAppliedOverrides(currentActiveTools, toolOverrides, availableTools); 245 const persistedConfig = 246 readLatestBranchConfig(ctx, baseActiveTools, availableTools) ?? 247 loadGlobalConfig(baseActiveTools, availableTools) ?? 248 EMPTY_TOOLS_CONFIG; 249 250 toolOverrides = overridesToMap(persistedConfig); 251 const desiredTools = applyOverrides(baseActiveTools, toolOverrides, allTools); 252 253 if (!sameToolMembership(desiredTools, currentActiveTools)) { 254 pi.setActiveTools(desiredTools); 255 } 256 } 257 258 pi.registerCommand("tools", { 259 description: "Enable/disable tools", 260 handler: async (_args, ctx) => { 261 allTools = pi.getAllTools(); 262 const availableTools = getAvailableToolSet(); 263 const baseActiveTools = stripPreviouslyAppliedOverrides( 264 new Set(pi.getActiveTools()), 265 toolOverrides, 266 availableTools, 267 ); 268 const initialEffectiveTools = new Set(applyOverrides(baseActiveTools, toolOverrides, allTools)); 269 270 await ctx.ui.custom((tui, theme, _kb, done) => { 271 const items: SettingItem[] = allTools.map((tool) => ({ 272 id: tool.name, 273 label: tool.name, 274 currentValue: initialEffectiveTools.has(tool.name) ? "enabled" : "disabled", 275 values: ["enabled", "disabled"], 276 })); 277 278 const container = new Container(); 279 container.addChild( 280 new (class { 281 render(_width: number) { 282 return [theme.fg("accent", theme.bold("Tool Configuration")), ""]; 283 } 284 invalidate() {} 285 })(), 286 ); 287 288 const settingsList = new SettingsList( 289 items, 290 Math.min(items.length + 2, 15), 291 getSettingsListTheme(), 292 (id, newValue) => { 293 const desiredEnabled = newValue === "enabled"; 294 updateToolOverride(id, desiredEnabled, baseActiveTools, toolOverrides); 295 296 const desiredTools = applyOverrides(baseActiveTools, toolOverrides, allTools); 297 const currentActiveTools = new Set(pi.getActiveTools()); 298 if (!sameToolMembership(desiredTools, currentActiveTools)) { 299 pi.setActiveTools(desiredTools); 300 } 301 persistState(); 302 }, 303 () => { 304 done(undefined); 305 }, 306 ); 307 308 container.addChild(settingsList); 309 310 const component = { 311 render(width: number) { 312 return container.render(width); 313 }, 314 invalidate() { 315 container.invalidate(); 316 }, 317 handleInput(data: string) { 318 settingsList.handleInput?.(data); 319 tui.requestRender(); 320 }, 321 }; 322 323 return component; 324 }); 325 }, 326 }); 327 328 pi.on("session_start", async (_event, ctx) => { 329 restoreFromBranch(ctx); 330 }); 331 332 pi.on("session_tree", async (_event, ctx) => { 333 restoreFromBranch(ctx); 334 }); 335 } 336 337 export const __test__ = { 338 applyOverrides, 339 getConfigPath, 340 normalizePersistedState, 341 serializeOverrides, 342 stripPreviouslyAppliedOverrides, 343 toPersistedState, 344 updateToolOverride, 345 };