dotfiles/pi/files/agent/extensions/slowtool.ts

/**
 * Slow Tool Monitor Extension
 *
 * Notifies you when a tool command is taking too long, useful for detecting
 * interactive commands (like git without --non-interactive, or commands waiting for stdin).
 *
 * Also integrates with Peon for sound notifications.
 *
 * Configuration:
 *   /slowtool:timeout <seconds>  - Set the timeout threshold (default: 30s)
 *   /slowtool:enable              - Enable notifications
 *   /slowtool:disable             - Disable notifications
 *   /slowtool:status              - Show current settings
 */

import type { ExtensionAPI, ExtensionCommandContext } from "@mariozechner/pi-coding-agent";
import * as fs from "node:fs";
import * as path from "node:path";
import * as os from "node:os";

interface ToolTimeout {
	toolCallId: string;
	toolName: string;
	command?: string;
	args: unknown;
	timeoutId: ReturnType<typeof setTimeout>;
	startTime: number;
	notified: boolean;
}

// Configuration
let timeoutSeconds = 30;
let enabled = true;
const SETTINGS_NAMESPACE = "slowtool";
const globalSettingsPath = path.join(os.homedir(), ".pi", "agent", "settings.json");

// Track running tools
const runningTools: Map<string, ToolTimeout> = new Map();

// ============ HELPERS ============

function formatDuration(ms: number): string {
	if (ms < 1000) return `${ms}ms`;
	const seconds = Math.floor(ms / 1000);
	if (seconds < 60) return `${seconds}s`;
	const minutes = Math.floor(seconds / 60);
	const remainingSeconds = seconds % 60;
	return `${minutes}m ${remainingSeconds}s`;
}

function asRecord(value: unknown): Record<string, unknown> | undefined {
	if (!value || typeof value !== "object") return undefined;
	return value as Record<string, unknown>;
}

function readSettingsFile(filePath: string): Record<string, unknown> {
	try {
		if (!fs.existsSync(filePath)) return {};
		const raw = fs.readFileSync(filePath, "utf-8");
		const parsed = JSON.parse(raw) as unknown;
		return asRecord(parsed) ?? {};
	} catch {
		return {};
	}
}

function loadGlobalConfig(): { timeoutSeconds: number; enabled: boolean } {
	const settings = readSettingsFile(globalSettingsPath);
	const slowtoolSettings = asRecord(settings[SETTINGS_NAMESPACE]);

	const configuredTimeout = slowtoolSettings?.timeoutSeconds;
	const nextTimeout =
		typeof configuredTimeout === "number" && Number.isFinite(configuredTimeout) && configuredTimeout >= 1
			? Math.floor(configuredTimeout)
			: 30;

	const configuredEnabled = slowtoolSettings?.enabled;
	const nextEnabled = typeof configuredEnabled === "boolean" ? configuredEnabled : true;

	return { timeoutSeconds: nextTimeout, enabled: nextEnabled };
}

function saveGlobalConfig(next: { timeoutSeconds: number; enabled: boolean }): boolean {
	try {
		const settings = readSettingsFile(globalSettingsPath);
		const existing = asRecord(settings[SETTINGS_NAMESPACE]) ?? {};
		settings[SETTINGS_NAMESPACE] = {
			...existing,
			timeoutSeconds: next.timeoutSeconds,
			enabled: next.enabled,
		};
		fs.mkdirSync(path.dirname(globalSettingsPath), { recursive: true });
		fs.writeFileSync(globalSettingsPath, `${JSON.stringify(settings, null, 2)}\n`, "utf-8");
		return true;
	} catch {
		return false;
	}
}

function getCommandPreview(args: unknown): string | undefined {
	if (!args) return undefined;
	const anyArgs = args as Record<string, unknown>;
	if (typeof anyArgs.command === "string") {
		const cmd = anyArgs.command;
		// Truncate long commands
		if (cmd.length > 60) {
			return cmd.slice(0, 60) + "...";
		}
		return cmd;
	}
	return undefined;
}

// ============ NOTIFICATIONS ============

function notifyTimeout(pi: ExtensionAPI, tool: ToolTimeout): void {
	if (tool.notified) return;
	tool.notified = true;

	const duration = Date.now() - tool.startTime;

	// Emit event for Peon integration and other listeners
	pi.events.emit("slowtool:timeout", {
		toolCallId: tool.toolCallId,
		toolName: tool.toolName,
		duration,
		command: tool.command,
	});
}

// ============ EVENT HANDLERS ============

export default function(pi: ExtensionAPI) {
	const applyPersistedConfig = () => {
		const persisted = loadGlobalConfig();
		timeoutSeconds = persisted.timeoutSeconds;
		enabled = persisted.enabled;
	};

	const persistCurrentConfig = (ctx: ExtensionCommandContext): void => {
		const ok = saveGlobalConfig({ timeoutSeconds, enabled });
		if (!ok) {
			ctx.ui.notify("Failed to persist slowtool settings", "warning");
		}
	};

	applyPersistedConfig();

	pi.on("session_start", async (_event, _ctx) => {
		applyPersistedConfig();
	});

	pi.on("session_switch", async (_event, _ctx) => {
		applyPersistedConfig();
	});

	// Register commands
	pi.registerCommand("slowtool:timeout", {
		description: "Set timeout threshold in seconds (default: 30)",
		handler: async (args: string, ctx: ExtensionCommandContext) => {
			if (!args.trim()) {
				ctx.ui.notify(`Current timeout: ${timeoutSeconds}s`, "info");
				return;
			}
			const newTimeout = parseInt(args.trim(), 10);
			if (isNaN(newTimeout) || newTimeout < 1) {
				ctx.ui.notify("Usage: /slowtool:timeout <seconds> (minimum 1)", "error");
				return;
			}
			timeoutSeconds = newTimeout;
			persistCurrentConfig(ctx);
			ctx.ui.notify(`Timeout set to ${timeoutSeconds}s`, "info");
		},
	});

	pi.registerCommand("slowtool:enable", {
		description: "Enable slow tool notifications",
		handler: async (_args: string, ctx: ExtensionCommandContext) => {
			enabled = true;
			persistCurrentConfig(ctx);
			ctx.ui.notify("Slow tool notifications enabled", "info");
		},
	});

	pi.registerCommand("slowtool:disable", {
		description: "Disable slow tool notifications",
		handler: async (_args: string, ctx: ExtensionCommandContext) => {
			enabled = false;
			persistCurrentConfig(ctx);
			ctx.ui.notify("Slow tool notifications disabled", "info");
		},
	});

	pi.registerCommand("slowtool:status", {
		description: "Show current settings",
		handler: async (_args: string, ctx: ExtensionCommandContext) => {
			const status = enabled ? "enabled" : "disabled";
			ctx.ui.notify(
				`Slow tool monitor: ${status}\nTimeout: ${timeoutSeconds}s`,
				"info"
			);
		},
	});

	// Track tool execution start
	pi.on("tool_execution_start", async (event, ctx) => {
		if (!enabled) return;

		const toolCallId = event.toolCallId;
		const toolName = event.toolName;
		const command = getCommandPreview(event.args);
		const startTime = Date.now();

		// Check if this is an interactive command that likely needs input
		// These are prime candidates for getting stuck
		const interactivePatterns = [
			// Git interactive commands
			/^git\s+(push|pull|fetch|rebase|merge|cherry-pick|restore|checkout)\b/i,
			/^git\b.*-i\b/i,
			/^git\b.*--interactive\b/i,
			// SSH
			/^ssh\b/i,
			// Commands that wait for stdin
			/^cat\s*>?\s*/i,
			/^tee\b/i,
			// Package managers that might prompt
			/^(sudo\s+)?apt(?:\s+get)?\s+(install|upgrade|remove)/i,
			/^(sudo\s+)?yum\s+(install|update|remove)/i,
			/^(sudo\s+)?brew\s+(install|update)/i,
			/^(sudo\s+)?pnpm\s+(install|add)/i,
			/^(sudo\s+)?npm\s+(install|add)/i,
			/^docker\s+(run|pull|build)/i,
			// Password prompts
			/^(sudo\s+)?passwd\b/i,
			// Yes/no prompts
			/\|\s*yes\b/i,
			// Editor launches
			/^(vim|nano|emacs|code\s+-w)\b/i,
			// Less/more pager
			/^\s*(less|more)\b/i,
		];

		// If it's an interactive command, use a shorter timeout
		const isLikelyInteractive = command && interactivePatterns.some(p => p.test(command));
		const effectiveTimeout = isLikelyInteractive ? Math.min(timeoutSeconds, 10) : timeoutSeconds;

		// Set up the timeout
		const timeoutId = setTimeout(() => {
			// Create a minimal context for notification
			notifyTimeout(pi, runningTools.get(toolCallId)!);
		}, effectiveTimeout * 1000);

		const toolTimeout: ToolTimeout = {
			toolCallId,
			toolName,
			command,
			args: event.args,
			timeoutId,
			startTime,
			notified: false,
		};

		runningTools.set(toolCallId, toolTimeout);
	});

	// Handle tool execution updates (reset timeout on each update)
	pi.on("tool_execution_update", async (event, ctx) => {
		const tool = runningTools.get(event.toolCallId);
		if (!tool) return;

		// If we've notified, don't reset - let it continue running
		// But we could extend the timeout if there's activity
		// For now, just track that there's progress
	});

	// Clean up when tool finishes
	pi.on("tool_execution_end", async (event, ctx) => {
		const tool = runningTools.get(event.toolCallId);
		if (!tool) return;

		// Clear the timeout
		clearTimeout(tool.timeoutId);

		// If we notified about a slow tool, send a completion notification
		if (tool.notified) {
			const duration = Date.now() - tool.startTime;
			const toolName = event.toolName;
			const preview = tool.command ? `\nCommand: ${tool.command}` : "";
			const resultText = event.isError ? "\n\n⚠️ Tool execution failed" : "\n\n✓ Tool completed";

			ctx.ui.notify(
				`${toolName} completed after ${formatDuration(duration)}${preview}${resultText}`,
				event.isError ? "warning" : "info"
			);

			// Emit completion event for Peon integration
			if (!event.isError) {
				pi.events.emit("slowtool:completed", {
					toolCallId: event.toolCallId,
					toolName: event.toolName,
					duration,
					command: tool.command,
				});
			}
		}

		runningTools.delete(event.toolCallId);
	});

	// Handle the slowtool:timeout event - show notification
	pi.events.on("slowtool:timeout", async (data) => {
		const timeoutData = data as { toolCallId: string; toolName: string; duration: number; command?: string };
		// This is handled in tool_execution_end for context
		// But we can also emit for Peon integration
		pi.events.emit("peon:input_required", {
			source: "slowtool",
			action: "slow_command",
			toolName: timeoutData.toolName,
			command: timeoutData.command,
			duration: timeoutData.duration,
		});
	});
}