dotfiles/pi/.pi/agent/extensions/pi-ask-tool/index.ts

import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
import { Type, type Static } from "@sinclair/typebox";
import { OTHER_OPTION, type AskQuestion } from "./ask-logic";
import { askSingleQuestionWithInlineNote } from "./ask-inline-ui";
import { askQuestionsWithTabs } from "./ask-tabs-ui";

const OptionItemSchema = Type.Object({
	label: Type.String({ description: "Display label" }),
});

const QuestionItemSchema = Type.Object({
	id: Type.String({ description: "Question id (e.g. auth, cache, priority)" }),
	question: Type.String({ description: "Question text" }),
	options: Type.Array(OptionItemSchema, {
		description: "Available options. Do not include 'Other'.",
		minItems: 1,
	}),
	multi: Type.Optional(Type.Boolean({ description: "Allow multi-select" })),
	recommended: Type.Optional(
		Type.Number({ description: "0-indexed recommended option. '(Recommended)' is shown automatically." }),
	),
});

const AskParamsSchema = Type.Object({
	questions: Type.Array(QuestionItemSchema, { description: "Questions to ask", minItems: 1 }),
});

type AskParams = Static<typeof AskParamsSchema>;

interface QuestionResult {
	id: string;
	question: string;
	options: string[];
	multi: boolean;
	selectedOptions: string[];
	customInput?: string;
}

interface AskToolDetails {
	id?: string;
	question?: string;
	options?: string[];
	multi?: boolean;
	selectedOptions?: string[];
	customInput?: string;
	results?: QuestionResult[];
}

function sanitizeForSessionText(value: string): string {
	return value
		.replace(/[\r\n\t]/g, " ")
		.replace(/[\x00-\x08\x0B\x0C\x0E-\x1F\x7F]/g, "")
		.replace(/\s{2,}/g, " ")
		.trim();
}

function sanitizeOptionForSessionText(option: string): string {
	const sanitizedOption = sanitizeForSessionText(option);
	return sanitizedOption.length > 0 ? sanitizedOption : "(empty option)";
}

function toSessionSafeQuestionResult(result: QuestionResult): QuestionResult {
	const selectedOptions = result.selectedOptions
		.map((selectedOption) => sanitizeForSessionText(selectedOption))
		.filter((selectedOption) => selectedOption.length > 0);

	const rawCustomInput = result.customInput;
	const customInput = rawCustomInput == null ? undefined : sanitizeForSessionText(rawCustomInput);

	return {
		id: sanitizeForSessionText(result.id) || "(unknown)",
		question: sanitizeForSessionText(result.question) || "(empty question)",
		options: result.options.map(sanitizeOptionForSessionText),
		multi: result.multi,
		selectedOptions,
		customInput: customInput && customInput.length > 0 ? customInput : undefined,
	};
}

function formatSelectionForSummary(result: QuestionResult): string {
	const hasSelectedOptions = result.selectedOptions.length > 0;
	const hasCustomInput = Boolean(result.customInput);

	if (!hasSelectedOptions && !hasCustomInput) {
		return "(cancelled)";
	}

	if (hasSelectedOptions && hasCustomInput) {
		const selectedPart = result.multi
			? `[${result.selectedOptions.join(", ")}]`
			: result.selectedOptions[0];
		return `${selectedPart} + Other: "${result.customInput}"`;
	}

	if (hasCustomInput) {
		return `"${result.customInput}"`;
	}

	if (result.multi) {
		return `[${result.selectedOptions.join(", ")}]`;
	}

	return result.selectedOptions[0];
}

function formatQuestionResult(result: QuestionResult): string {
	return `${result.id}: ${formatSelectionForSummary(result)}`;
}

function formatQuestionContext(result: QuestionResult, questionIndex: number): string {
	const lines: string[] = [
		`Question ${questionIndex + 1} (${result.id})`,
		`Prompt: ${result.question}`,
		"Options:",
		...result.options.map((option, optionIndex) => `  ${optionIndex + 1}. ${option}`),
		"Response:",
	];

	const hasSelectedOptions = result.selectedOptions.length > 0;
	const hasCustomInput = Boolean(result.customInput);

	if (!hasSelectedOptions && !hasCustomInput) {
		lines.push("  Selected: (cancelled)");
		return lines.join("\n");
	}

	if (hasSelectedOptions) {
		const selectedText = result.multi
			? `[${result.selectedOptions.join(", ")}]`
			: result.selectedOptions[0];
		lines.push(`  Selected: ${selectedText}`);
	}

	if (hasCustomInput) {
		if (!hasSelectedOptions) {
			lines.push(`  Selected: ${OTHER_OPTION}`);
		}
		lines.push(`  Custom input: ${result.customInput}`);
	}

	return lines.join("\n");
}

function buildAskSessionContent(results: QuestionResult[]): string {
	const safeResults = results.map(toSessionSafeQuestionResult);
	const summaryLines = safeResults.map(formatQuestionResult).join("\n");
	const contextBlocks = safeResults.map((result, index) => formatQuestionContext(result, index)).join("\n\n");
	return `User answers:\n${summaryLines}\n\nAnswer context:\n${contextBlocks}`;
}

const ASK_TOOL_DESCRIPTION = `
Ask the user for clarification when a choice materially affects the outcome.

- Use when multiple valid approaches have different trade-offs.
- Prefer 2-5 concise options.
- Use multi=true when multiple answers are valid.
- Use recommended=<index> (0-indexed) to mark the default option.
- You can ask multiple related questions in one call using questions[].
- Do NOT include an 'Other' option; UI adds it automatically.
`.trim();

export default function askExtension(pi: ExtensionAPI) {
	pi.registerTool({
		name: "ask",
		label: "Ask",
		description: ASK_TOOL_DESCRIPTION,
		parameters: AskParamsSchema,

		async execute(_toolCallId, params: AskParams, _signal, _onUpdate, ctx) {
			if (!ctx.hasUI) {
				return {
					content: [{ type: "text", text: "Error: ask tool requires interactive mode" }],
					details: {},
				};
			}

			if (params.questions.length === 0) {
				return {
					content: [{ type: "text", text: "Error: questions must not be empty" }],
					details: {},
				};
			}

			if (params.questions.length === 1) {
				const [q] = params.questions;
				const selection = q.multi
					? (await askQuestionsWithTabs(ctx.ui, [q as AskQuestion])).selections[0] ?? { selectedOptions: [] }
					: await askSingleQuestionWithInlineNote(ctx.ui, q as AskQuestion);
				const optionLabels = q.options.map((option) => option.label);

				const result: QuestionResult = {
					id: q.id,
					question: q.question,
					options: optionLabels,
					multi: q.multi ?? false,
					selectedOptions: selection.selectedOptions,
					customInput: selection.customInput,
				};

				const details: AskToolDetails = {
					id: q.id,
					question: q.question,
					options: optionLabels,
					multi: q.multi ?? false,
					selectedOptions: selection.selectedOptions,
					customInput: selection.customInput,
					results: [result],
				};

				return {
					content: [{ type: "text", text: buildAskSessionContent([result]) }],
					details,
				};
			}

			const results: QuestionResult[] = [];
			const tabResult = await askQuestionsWithTabs(ctx.ui, params.questions as AskQuestion[]);
			for (let i = 0; i < params.questions.length; i++) {
				const q = params.questions[i];
				const selection = tabResult.selections[i] ?? { selectedOptions: [] };
				results.push({
					id: q.id,
					question: q.question,
					options: q.options.map((option) => option.label),
					multi: q.multi ?? false,
					selectedOptions: selection.selectedOptions,
					customInput: selection.customInput,
				});
			}

			return {
				content: [{ type: "text", text: buildAskSessionContent(results) }],
				details: { results } satisfies AskToolDetails,
			};
		},
	});
}