1// SPDX-FileCopyrightText: Mario Zechner <https://mariozechner.at>
2//
3// SPDX-License-Identifier: MIT
4
5/**
6 * Questionnaire Tool - Unified tool for asking single or multiple questions
7 *
8 * Single question: simple options list
9 * Multiple questions: tab bar navigation between questions
10 */
11
12import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
13import {
14 Editor,
15 type EditorTheme,
16 Key,
17 matchesKey,
18 Text,
19 truncateToWidth,
20 wrapTextWithAnsi,
21} from "@mariozechner/pi-tui";
22import { Type } from "@sinclair/typebox";
23
24// Types
25interface QuestionOption {
26 value: string;
27 label: string;
28 description?: string;
29}
30
31type RenderOption = QuestionOption & { isOther?: boolean };
32
33interface Question {
34 id: string;
35 label: string;
36 prompt: string;
37 options: QuestionOption[];
38 allowOther: boolean;
39}
40
41interface Answer {
42 id: string;
43 value: string;
44 label: string;
45 wasCustom: boolean;
46 index?: number;
47}
48
49interface QuestionnaireResult {
50 questions: Question[];
51 answers: Answer[];
52 cancelled: boolean;
53}
54
55// Schema
56const QuestionOptionSchema = Type.Object({
57 value: Type.String({ description: "The value returned when selected" }),
58 label: Type.String({ description: "Display label for the option" }),
59 description: Type.Optional(Type.String({ description: "Optional description shown below label" })),
60});
61
62const QuestionSchema = Type.Object({
63 id: Type.String({ description: "Unique identifier for this question" }),
64 label: Type.Optional(
65 Type.String({
66 description: "Short contextual label for tab bar, e.g. 'Scope', 'Priority' (defaults to Q1, Q2)",
67 }),
68 ),
69 prompt: Type.String({ description: "The full question text to display" }),
70 options: Type.Array(QuestionOptionSchema, { description: "Available options to choose from" }),
71 allowOther: Type.Optional(Type.Boolean({ description: "Allow 'Type something' option (default: true)" })),
72});
73
74const QuestionnaireParams = Type.Object({
75 questions: Type.Array(QuestionSchema, { description: "Questions to ask the user" }),
76});
77
78function errorResult(
79 message: string,
80 questions: Question[] = [],
81): { content: { type: "text"; text: string }[]; details: QuestionnaireResult } {
82 return {
83 content: [{ type: "text", text: message }],
84 details: { questions, answers: [], cancelled: true },
85 };
86}
87
88export default function questionnaire(pi: ExtensionAPI) {
89 pi.registerTool({
90 name: "questionnaire",
91 label: "Questionnaire",
92 description:
93 "Ask the user one or more questions. Use for clarifying requirements, getting preferences, or confirming decisions. For single questions, shows a simple option list. For multiple questions, shows a tab-based interface. Users can always add a free-text note alongside their answers.",
94 parameters: QuestionnaireParams,
95
96 async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
97 if (!ctx.hasUI) {
98 return errorResult("Error: UI not available (running in non-interactive mode)");
99 }
100 if (params.questions.length === 0) {
101 return errorResult("Error: No questions provided");
102 }
103
104 // Normalize questions with defaults and deduplicate IDs
105 const seenIds = new Set<string>();
106 const questions: Question[] = params.questions.map((q, i) => {
107 let id = q.id;
108 if (seenIds.has(id)) {
109 let suffix = 2;
110 while (seenIds.has(`${id}_${suffix}`)) suffix++;
111 id = `${id}_${suffix}`;
112 }
113 seenIds.add(id);
114 return {
115 ...q,
116 id,
117 label: q.label || `Q${i + 1}`,
118 allowOther: q.allowOther !== false,
119 };
120 });
121
122 const isMulti = questions.length > 1;
123 const totalTabs = questions.length + 1; // questions + Submit
124
125 const result = await ctx.ui.custom<QuestionnaireResult>((tui, theme, _kb, done) => {
126 // State
127 let currentTab = 0;
128 let optionIndex = 0;
129 let inputMode = false;
130 let inputQuestionId: string | null = null;
131 let cachedLines: string[] | undefined;
132 const answers = new Map<string, Answer>();
133
134 // Editor for "Type something" option
135 const editorTheme: EditorTheme = {
136 borderColor: (s) => theme.fg("accent", s),
137 selectList: {
138 selectedPrefix: (t) => theme.fg("accent", t),
139 selectedText: (t) => theme.fg("accent", t),
140 description: (t) => theme.fg("muted", t),
141 scrollInfo: (t) => theme.fg("dim", t),
142 noMatch: (t) => theme.fg("warning", t),
143 },
144 };
145 const editor = new Editor(tui, editorTheme);
146
147 // Helpers
148 function refresh() {
149 cachedLines = undefined;
150 tui.requestRender();
151 }
152
153 function submit(cancelled: boolean) {
154 const allAnswers = Array.from(answers.values());
155 if (!cancelled && notesText) {
156 allAnswers.push({ id: "_note", value: notesText, label: notesText, wasCustom: true });
157 }
158 done({ questions, answers: allAnswers, cancelled });
159 }
160
161 function currentQuestion(): Question | undefined {
162 return questions[currentTab];
163 }
164
165 function currentOptions(): RenderOption[] {
166 const q = currentQuestion();
167 if (!q) return [];
168 const opts: RenderOption[] = [...q.options];
169 if (q.allowOther) {
170 opts.push({ value: "__other__", label: "Type something.", isOther: true });
171 }
172 return opts;
173 }
174
175 function allAnswered(): boolean {
176 return questions.every((q) => answers.has(q.id));
177 }
178
179 function advanceAfterAnswer() {
180 if (!isMulti) {
181 submit(false);
182 return;
183 }
184 if (currentTab < questions.length - 1) {
185 currentTab++;
186 } else {
187 currentTab = questions.length; // Submit tab
188 }
189 optionIndex = 0;
190 refresh();
191 }
192
193 function saveAnswer(questionId: string, value: string, label: string, wasCustom: boolean, index?: number) {
194 answers.set(questionId, { id: questionId, value, label, wasCustom, index });
195 }
196
197 // Editor submit callback
198 editor.onSubmit = (value) => {
199 if (!inputQuestionId) return;
200 const trimmed = value.trim() || "(no response)";
201 saveAnswer(inputQuestionId, trimmed, trimmed, true);
202 inputMode = false;
203 inputQuestionId = null;
204 editor.setText("");
205 advanceAfterAnswer();
206 };
207
208 // Free-text notes (always available via 'n' key)
209 let notesMode = false;
210 let notesText = "";
211 const notesEditor = new Editor(tui, editorTheme);
212 notesEditor.onSubmit = (value) => {
213 notesText = value.trim();
214 notesMode = false;
215 refresh();
216 };
217
218 function handleInput(data: string) {
219 // Input mode: route to editor
220 if (inputMode) {
221 if (matchesKey(data, Key.escape)) {
222 inputMode = false;
223 inputQuestionId = null;
224 editor.setText("");
225 refresh();
226 return;
227 }
228 editor.handleInput(data);
229 refresh();
230 return;
231 }
232
233 // Notes mode: route to notes editor
234 if (notesMode) {
235 if (matchesKey(data, Key.escape)) {
236 notesMode = false;
237 notesEditor.setText(notesText);
238 refresh();
239 return;
240 }
241 notesEditor.handleInput(data);
242 refresh();
243 return;
244 }
245
246 // Activate notes editor (available on any tab)
247 if (data === "n" || data === "N") {
248 notesMode = true;
249 notesEditor.setText(notesText);
250 refresh();
251 return;
252 }
253
254 const q = currentQuestion();
255 const opts = currentOptions();
256
257 // Tab navigation (multi-question only)
258 if (isMulti) {
259 if (matchesKey(data, Key.tab) || matchesKey(data, Key.right)) {
260 currentTab = (currentTab + 1) % totalTabs;
261 optionIndex = 0;
262 refresh();
263 return;
264 }
265 if (matchesKey(data, Key.shift("tab")) || matchesKey(data, Key.left)) {
266 currentTab = (currentTab - 1 + totalTabs) % totalTabs;
267 optionIndex = 0;
268 refresh();
269 return;
270 }
271 }
272
273 // Submit tab
274 if (currentTab === questions.length) {
275 if (matchesKey(data, Key.enter) && allAnswered()) {
276 submit(false);
277 } else if (matchesKey(data, Key.escape)) {
278 submit(true);
279 }
280 return;
281 }
282
283 // Option navigation
284 if (matchesKey(data, Key.up)) {
285 optionIndex = Math.max(0, optionIndex - 1);
286 refresh();
287 return;
288 }
289 if (matchesKey(data, Key.down)) {
290 optionIndex = Math.min(opts.length - 1, optionIndex + 1);
291 refresh();
292 return;
293 }
294
295 // Select option
296 if (matchesKey(data, Key.enter) && q) {
297 const opt = opts[optionIndex];
298 if (!opt) return; // no options available
299 if (opt.isOther) {
300 inputMode = true;
301 inputQuestionId = q.id;
302 editor.setText("");
303 refresh();
304 return;
305 }
306 saveAnswer(q.id, opt.value, opt.label, false, optionIndex + 1);
307 advanceAfterAnswer();
308 return;
309 }
310
311 // Cancel
312 if (matchesKey(data, Key.escape)) {
313 submit(true);
314 }
315 }
316
317 function render(width: number): string[] {
318 if (cachedLines) return cachedLines;
319
320 const lines: string[] = [];
321 const q = currentQuestion();
322 const opts = currentOptions();
323
324 // Helper to add truncated line
325 const add = (s: string) => lines.push(truncateToWidth(s, width));
326
327 add(theme.fg("accent", "─".repeat(width)));
328
329 // Tab bar (multi-question only)
330 if (isMulti) {
331 const tabs: string[] = ["← "];
332 for (let i = 0; i < questions.length; i++) {
333 const isActive = i === currentTab;
334 const isAnswered = answers.has(questions[i].id);
335 const lbl = questions[i].label;
336 const box = isAnswered ? "■" : "□";
337 const color = isAnswered ? "success" : "muted";
338 const text = ` ${box} ${lbl} `;
339 const styled = isActive ? theme.bg("selectedBg", theme.fg("text", text)) : theme.fg(color, text);
340 tabs.push(`${styled} `);
341 }
342 const canSubmit = allAnswered();
343 const isSubmitTab = currentTab === questions.length;
344 const submitText = " ✓ Submit ";
345 const submitStyled = isSubmitTab
346 ? theme.bg("selectedBg", theme.fg("text", submitText))
347 : theme.fg(canSubmit ? "success" : "dim", submitText);
348 tabs.push(`${submitStyled} →`);
349 add(` ${tabs.join("")}`);
350 lines.push("");
351 }
352
353 // Helper to render options list
354 function renderOptions() {
355 for (let i = 0; i < opts.length; i++) {
356 const opt = opts[i];
357 const selected = i === optionIndex;
358 const isOther = opt.isOther === true;
359 const prefix = selected ? theme.fg("accent", "> ") : " ";
360 const color = selected ? "accent" : "text";
361 // Mark "Type something" differently when in input mode
362 if (isOther && inputMode) {
363 add(prefix + theme.fg("accent", `${i + 1}. ${opt.label} ✎`));
364 } else {
365 add(prefix + theme.fg(color, `${i + 1}. ${opt.label}`));
366 }
367 if (opt.description) {
368 add(` ${theme.fg("muted", opt.description)}`);
369 }
370 }
371 }
372
373 // Helper to add word-wrapped lines (for long prompts)
374 const addWrapped = (s: string) => {
375 for (const line of wrapTextWithAnsi(s, width)) {
376 lines.push(line);
377 }
378 };
379
380 // Content
381 if (inputMode && q) {
382 addWrapped(theme.fg("text", ` ${q.prompt}`));
383 lines.push("");
384 // Show options for reference
385 renderOptions();
386 lines.push("");
387 add(theme.fg("muted", " Your answer:"));
388 for (const line of editor.render(width - 2)) {
389 add(` ${line}`);
390 }
391 lines.push("");
392 add(theme.fg("dim", " Enter to submit • Esc to cancel"));
393 } else if (currentTab === questions.length) {
394 add(theme.fg("accent", theme.bold(" Ready to submit")));
395 lines.push("");
396 for (const question of questions) {
397 const answer = answers.get(question.id);
398 if (answer) {
399 const prefix = answer.wasCustom ? "(wrote) " : "";
400 add(`${theme.fg("muted", ` ${question.label}: `)}${theme.fg("text", prefix + answer.label)}`);
401 }
402 }
403 lines.push("");
404 if (allAnswered()) {
405 add(theme.fg("success", " Press Enter to submit"));
406 } else {
407 const missing = questions
408 .filter((q) => !answers.has(q.id))
409 .map((q) => q.label)
410 .join(", ");
411 add(theme.fg("warning", ` Unanswered: ${missing}`));
412 }
413 } else if (q) {
414 addWrapped(theme.fg("text", ` ${q.prompt}`));
415 lines.push("");
416 renderOptions();
417 }
418
419 // Notes section (visible unless typing an answer)
420 if (!inputMode) {
421 if (notesMode) {
422 lines.push("");
423 add(theme.fg("accent", " ✎ Note:"));
424 for (const line of notesEditor.render(width - 2)) {
425 add(` ${line}`);
426 }
427 lines.push("");
428 add(theme.fg("dim", " Enter to save • Esc to discard"));
429 } else if (notesText) {
430 lines.push("");
431 add(theme.fg("muted", " ✎ Note: ") + theme.fg("text", notesText));
432 add(theme.fg("dim", " Press 'n' to edit"));
433 } else {
434 lines.push("");
435 add(theme.fg("dim", " Press 'n' to add a note"));
436 }
437 }
438
439 lines.push("");
440 if (!inputMode && !notesMode) {
441 const help = isMulti
442 ? " Tab/←→ navigate • ↑↓ select • Enter confirm • Esc cancel"
443 : " ↑↓ navigate • Enter select • Esc cancel";
444 add(theme.fg("dim", help));
445 }
446 add(theme.fg("accent", "─".repeat(width)));
447
448 cachedLines = lines;
449 return lines;
450 }
451
452 return {
453 render,
454 invalidate: () => {
455 cachedLines = undefined;
456 },
457 handleInput,
458 };
459 });
460
461 if (result.cancelled) {
462 return {
463 content: [{ type: "text", text: "User cancelled the questionnaire" }],
464 details: result,
465 };
466 }
467
468 // Check if any answers were custom (user typed instead of selecting)
469 const customAnswers = result.answers.filter((a) => a.wasCustom);
470 const selectedAnswers = result.answers.filter((a) => !a.wasCustom);
471
472 // If there are custom answers, send them as a user message
473 if (customAnswers.length > 0) {
474 // Build user message with identifiers for each custom response
475 const messageLines = customAnswers.map((a) => `[${a.id}] ${a.label}`);
476 const userMessage = messageLines.join("\n\n");
477 pi.sendUserMessage(userMessage, { deliverAs: "steer" });
478
479 // Build tool result showing what happened
480 const toolResultLines: string[] = [];
481
482 // First show any selected answers normally
483 for (const a of selectedAnswers) {
484 const qLabel = questions.find((q) => q.id === a.id)?.label || a.id;
485 toolResultLines.push(`${qLabel}: user selected: ${a.index}. ${a.label}`);
486 }
487
488 // Then indicate which questions got user message responses
489 for (const a of customAnswers) {
490 if (a.id === "_note") {
491 toolResultLines.push(`Note: (see [_note] in following message)`);
492 } else {
493 const qLabel = questions.find((q) => q.id === a.id)?.label || a.id;
494 toolResultLines.push(`${qLabel}: (see [${a.id}] in following message)`);
495 }
496 }
497
498 return {
499 content: [{ type: "text", text: toolResultLines.join("\n") }],
500 details: result,
501 };
502 }
503
504 const answerLines = result.answers.map((a) => {
505 const qLabel = questions.find((q) => q.id === a.id)?.label || a.id;
506 if (a.wasCustom) {
507 return `${qLabel}: user wrote: ${a.label}`;
508 }
509 return `${qLabel}: user selected: ${a.index}. ${a.label}`;
510 });
511
512 return {
513 content: [{ type: "text", text: answerLines.join("\n") }],
514 details: result,
515 };
516 },
517
518 renderCall(args, theme) {
519 const qs = (args.questions as Question[]) || [];
520 const count = qs.length;
521 const labels = qs.map((q) => q.label || q.id).join(", ");
522 let text = theme.fg("toolTitle", theme.bold("questionnaire "));
523 text += theme.fg("muted", `${count} question${count !== 1 ? "s" : ""}`);
524 if (labels) {
525 text += theme.fg("dim", ` (${truncateToWidth(labels, 40)})`);
526 }
527 return new Text(text, 0, 0);
528 },
529
530 renderResult(result, _options, theme) {
531 const details = result.details as QuestionnaireResult | undefined;
532 if (!details) {
533 const text = result.content[0];
534 return new Text(text?.type === "text" ? text.text : "", 0, 0);
535 }
536 if (details.cancelled) {
537 return new Text(theme.fg("warning", "Cancelled"), 0, 0);
538 }
539 const lines = details.answers.map((a) => {
540 if (a.id === "_note") {
541 return `${theme.fg("success", "✓ ")}${theme.fg("accent", "Note")}: ${theme.fg("muted", "(via message)")}`;
542 }
543 if (a.wasCustom) {
544 return `${theme.fg("success", "✓ ")}${theme.fg("accent", a.id)}: ${theme.fg("muted", "(via message) ")}[${a.id}]`;
545 }
546 const display = a.index ? `${a.index}. ${a.label}` : a.label;
547 return `${theme.fg("success", "✓ ")}${theme.fg("accent", a.id)}: ${display}`;
548 });
549 return new Text(lines.join("\n"), 0, 0);
550 },
551 });
552}