1// SPDX-FileCopyrightText: Amolith <amolith@secluded.site>
2// SPDX-FileCopyrightText: Armin Ronacher <armin.ronacher@active-4.com>
3//
4// SPDX-License-Identifier: Apache-2.0
5
6/**
7 * Q&A extraction hook - extracts questions from assistant responses
8 *
9 * Custom interactive TUI for answering questions.
10 *
11 * Demonstrates the "prompt generator" pattern with custom TUI:
12 * 1. /answer command gets the last assistant message
13 * 2. Shows a spinner while extracting questions as structured JSON
14 * 3. Presents an interactive TUI to navigate and answer questions
15 * 4. Submits the compiled answers when done
16 */
17
18import { complete, type Tool, type UserMessage } from "@mariozechner/pi-ai";
19import type { ExtensionAPI, ExtensionContext } from "@mariozechner/pi-coding-agent";
20import { BorderedLoader } from "@mariozechner/pi-coding-agent";
21import {
22 type Component,
23 Editor,
24 type EditorTheme,
25 Key,
26 matchesKey,
27 truncateToWidth,
28 type TUI,
29 visibleWidth,
30 wrapTextWithAnsi,
31} from "@mariozechner/pi-tui";
32import { Type } from "@sinclair/typebox";
33
34/** Escape characters that would break pseudo-XML tag boundaries. */
35function escapeXml(s: string): string {
36 return s.replace(/&/g, "&").replace(/</g, "<").replace(/>/g, ">");
37}
38
39// Structured output format for question extraction
40interface ExtractedQuestion {
41 question: string;
42 context?: string;
43}
44
45interface ExtractionResult {
46 questions: ExtractedQuestion[];
47}
48
49const SYSTEM_PROMPT = `You extract questions from assistant messages so a user can answer them without reading the full message. Call the extract_questions tool with your results.
50
51Each extracted question must be a standalone object with a "question" string and an optional "context" string. The user will read ONLY the extracted questions, not the original message, so both the question and context must be fully self-contained.
52
53Rules:
54- Each question must make sense on its own without referring back to the message (no "Issue 1", "the above", "as mentioned")
55- Context should include all relevant details, options, and reasoning needed to answer — it replaces reading the original message
56- Merge duplicate or overlapping questions into one (e.g. an analysis section and a plan section about the same topic become one question)
57- Capture both explicit questions ("Which approach?") and implicit ones (proposals that need approval, recommendations that need confirmation)
58- Keep questions in logical order
59- If no questions are found, call the tool with an empty questions array
60
61<examples>
62<example>
63<description>Simple explicit question</description>
64<input>I can set up the database with either PostgreSQL or MySQL. The existing infrastructure uses PostgreSQL for two other services. Which database should I use?</input>
65<output>
66{"question": "Which database should be used for the new service?", "context": "Options are PostgreSQL or MySQL. The existing infrastructure already uses PostgreSQL for two other services."}
67</output>
68</example>
69
70<example>
71<description>Implicit proposal needing approval</description>
72<input>I looked at the failing test. The flakiness comes from a race condition in the cleanup goroutine. I think the simplest fix is to add a sync.WaitGroup so the test waits for cleanup to finish before asserting. Want me to go ahead with that?</input>
73<output>
74{"question": "Should the flaky test be fixed by adding a sync.WaitGroup to wait for the cleanup goroutine before asserting?", "context": "The test failure is caused by a race condition in the cleanup goroutine. The WaitGroup approach ensures cleanup completes before assertions run."}
75</output>
76</example>
77
78<example>
79<description>Long message with analysis and plan sections that overlap — questions must be merged, deduplicated, and fully standalone</description>
80<input>
81## Issue 1: Non-deterministic map iteration for overrides → rawArgs
82
83The problem is in RunE at line ~75: for k, vals := range overrides gives random order each run.
84
85Three good approaches:
861. Sort the keys before iterating. Simplest, most common Go pattern.
872. Use []restic.Flag or []struct{key, vals} instead of a map. More invasive but preserves construction order.
883. Skip the map→rawArgs→parsePassthrough round-trip entirely. Instead of serializing overrides to rawArgs and then having runCommand parse them back, just pass the overrides map directly to runCommand and merge there. This eliminates the round-trip and the ordering problem disappears.
89
90Recommend approach 3 as best: it kills the root cause.
91
92## Issue 2: Double-pointers (**screens.Snapshot etc.)
93
94The problem is buildCommandScreens needs to both return screens AND communicate which specific screen instances were created back to runInteractive. Double-pointers are gnarly.
95
96Simpler approach: return a struct holding typed screen pointers. The ResolveFunc closure captures the struct.
97
98## Issue 3: Picker height overflow
99
100The session already subtracts chrome height via adjustedSizeMsg() before forwarding WindowSizeMsg to screens. The legacy pickerModel does its own subtraction because it runs standalone. Assessment: this finding is a false positive.
101
102## Issue 2.1: huh form boilerplate
103
104The buildForm/Init/Update boilerplate is repetitive across Overwrite, Target, Preset, Snapshot. Code is simple and clear in each screen. Recommendation: skip — borderline DRY.
105
106## Issue 2.2: drain test helpers
107
108The drain* test helpers are copy-pasted with only the type changed. Perfect case for Go generics: func drain[S ui.Screen](s S, cmd tea.Cmd) (S, tea.Cmd)
109
110## Issue 3.1: Snapshot theme handler clears user input
111
112When BackgroundColorMsg arrives during phaseSnapshotManual, buildManualForm() resets s.entered = "". Fix: preserve s.entered before rebuilding. Also needs rebuildSelectForm() call during phaseSnapshotSelecting.
113
114## Issue 4: WindowSizeMsg missed during phaseFileLoading
115
116If the screen is loading when the initial WindowSizeMsg arrives, it only goes to the spinner and the size is lost. Fix: store the last WindowSizeMsg on the FilePicker struct, apply when picker is built.
117
118## Issue 5.1: Same as Issue 1
119
120Same root cause — map iteration in RunE. Fix for Issue 1 handles this.
121
122## Issue 6.1: Duplicated constant
123
124Both internal/config and internal/restic define the same commandSuffix constant.
125
126## Issue 6.2: Inconsistent empty-string semantics
127
128rc.Environ treats "" as present (checks map key existence), but os.Getenv checks != "" so an empty env var counts as absent.
129
130## Plan
131
1321. Remove legacy restore prompts (promptRestore, promptSnapshotID, promptFileSelection, printRestoreSummary, standalone pickerModel)
1332. Eliminate map→rawArgs round-trip (Issues 1 and 5.1)
1343. Replace double-pointers with return struct (Issue 2)
1354. Extract generic drain test helper (Issue 2.2)
1365. Fix snapshot theme handler (Issue 3.1)
1376. Store WindowSizeMsg during loading (Issue 4)
1387. Share commandSuffix constant (Issue 6.1)
1398. Fix inconsistent empty-string check (Issue 6.2)
140
141NOT planning to do Issue 2.1 (huh form boilerplate extraction).
142</input>
143<output>
144{"question": "How should the non-deterministic map iteration in RunE be fixed, where 'for k, vals := range overrides' produces random argument order each run?", "context": "Three approaches: (1) Sort map keys before iterating — simplest, small code change. (2) Replace map[string][]string with a slice of key-value pairs to preserve construction order. (3) Eliminate the map→rawArgs→parsePassthrough round-trip entirely by passing overrides directly to runCommand — this also fixes the identical issue in rawArgs parsing. Recommended: approach 3."}
145{"question": "Should the double-pointer output parameters (**screens.Snapshot, etc.) in buildCommandScreens be replaced with a returned struct?", "context": "Currently uses double-pointers so runInteractive can read results from specific screen instances. A commandScreens struct holding typed screen pointers would be returned instead, and the ResolveFunc closure captures it. Idiomatic Go — return a struct instead of output parameters."}
146{"question": "The picker height overflow appears to be a false positive because the session already subtracts chrome height via adjustedSizeMsg() before forwarding WindowSizeMsg to screens. Safe to skip?", "context": "session.go subtracts 5 chrome lines (breadcrumb + title + help bar) before forwarding. The legacy pickerModel does its own subtraction because it runs standalone with tea.NewProgram. The screen adapter doesn't need additional subtraction."}
147{"question": "Should the repeated huh form boilerplate (buildForm/Init/Update) across Overwrite, Target, Preset, and Snapshot screens be extracted into a shared helper?", "context": "Each screen repeats same form setup pattern with slight behavior variations. Code is currently simple and self-contained in each screen. Recommended: skip — borderline DRY."}
148{"question": "Should the copy-pasted type-specific drain test helpers be replaced with a single generic drain[S ui.Screen] function?", "context": "Multiple test files have identical drain helpers differing only in type parameter (e.g. drainSnapshot, drainFilePicker). A single generic function would replace all of them."}
149{"question": "Should the snapshot theme handler be fixed to preserve user input when the terminal theme changes?", "context": "When BackgroundColorMsg arrives during phaseSnapshotManual, buildManualForm() resets s.entered to empty, wiping user input. Fix: preserve s.entered before rebuild. Also needs rebuildSelectForm() call during phaseSnapshotSelecting."}
150{"question": "Should FilePicker store the last WindowSizeMsg to prevent size loss during the loading phase?", "context": "If the screen is loading when the initial WindowSizeMsg arrives, the message only goes to the spinner and size is lost. When buildPicker runs later, picker doesn't know terminal size. Fix: add lastSize field, apply when picker built."}
151{"question": "Should the duplicated commandSuffix constant be extracted to a shared package?", "context": "Both internal/config and internal/restic define the same constant."}
152{"question": "Should os.Getenv() != '' checks be replaced with os.LookupEnv so empty environment variables count as present?", "context": "rc.Environ treats empty string as present (checks map key existence), but os.Getenv checks != '' so an empty env var counts as absent. Using os.LookupEnv would make behavior consistent — explicitly setting RESTIC_REPOSITORY='' would count as configured."}
153{"question": "Should the legacy interactive restore prompts (promptRestore, promptSnapshotID, promptFileSelection, printRestoreSummary, standalone pickerModel) be removed from root.go?", "context": "The TUI session now handles the full restore flow through screens, making these standalone prompt functions unused."}
154</output>
155</example>
156</examples>`;
157
158/**
159 * Tool definition for structured question extraction. Models produce
160 * structured tool-call arguments far more reliably than free-form JSON
161 * in a text response.
162 */
163const QUESTION_EXTRACTION_TOOL: Tool = {
164 name: "extract_questions",
165 description:
166 "Extract questions from the assistant's message. Each question is a self-contained object the user will read instead of the original message.",
167 parameters: Type.Object({
168 questions: Type.Array(
169 Type.Object({
170 question: Type.String({
171 description: "A complete, standalone question. Must make sense without reading the original message.",
172 }),
173 context: Type.Optional(
174 Type.String({
175 description:
176 "Self-contained context with all details, options, and reasoning needed to answer the question.",
177 }),
178 ),
179 }),
180 {
181 description:
182 "Array of question objects. Merge overlapping questions. Each item MUST be an object with 'question' and optional 'context' strings, NOT a plain string.",
183 },
184 ),
185 }),
186};
187
188// Preferred model for extraction — lightweight is fine for structured JSON output.
189// Falls back to the session model if this one isn't available.
190const PREFERRED_EXTRACTION_MODEL_ID = "nemotron-3-super-120b-a12b";
191
192/**
193 * Resolve the extraction model: prefer a lightweight model from the registry,
194 * fall back to the current session model.
195 */
196function resolveExtractionModel(ctx: ExtensionContext, fallback: NonNullable<ExtensionContext["model"]>) {
197 const available = ctx.modelRegistry.getAvailable();
198 const preferred = available.find((m) => m.id === PREFERRED_EXTRACTION_MODEL_ID);
199 return preferred ?? fallback;
200}
201
202/**
203 * Parse a tool call response into an ExtractionResult.
204 */
205function parseToolCallResult(response: Awaited<ReturnType<typeof complete>>): ExtractionResult | null {
206 const toolCall = response.content.find((c) => c.type === "toolCall" && c.name === "extract_questions");
207
208 if (!toolCall || toolCall.type !== "toolCall") {
209 console.error("Model did not call extract_questions:", response.content);
210 return null;
211 }
212
213 const args = toolCall.arguments as Record<string, unknown>;
214 if (!Array.isArray(args.questions)) {
215 console.error("[answer] expected questions array, got:", typeof args.questions, args);
216 return null;
217 }
218
219 // Validate each question item — model may return plain strings instead
220 // of objects if it ignores the schema, so we handle both gracefully.
221 const questions: ExtractedQuestion[] = [];
222 for (const item of args.questions) {
223 if (typeof item === "string") {
224 // Model returned bare string instead of object — use as question text
225 questions.push({ question: item });
226 continue;
227 }
228 if (typeof item !== "object" || item === null) continue;
229 const obj = item as Record<string, unknown>;
230 if (typeof obj.question !== "string") {
231 console.error("[answer] skipping item with no 'question' string:", Object.keys(obj));
232 continue;
233 }
234 questions.push({
235 question: obj.question,
236 context: typeof obj.context === "string" ? obj.context : undefined,
237 });
238 }
239
240 return { questions };
241}
242
243/**
244 * Interactive Q&A component for answering extracted questions
245 */
246class QnAComponent implements Component {
247 private questions: ExtractedQuestion[];
248 private answers: string[];
249 private currentIndex: number = 0;
250 private editor: Editor;
251 private tui: TUI;
252 private onDone: (result: string | null) => void;
253 private showingConfirmation: boolean = false;
254 private notesMode: boolean = false;
255 private notesText: string = "";
256 private notesEditor: Editor;
257
258 // Cache
259 private cachedWidth?: number;
260 private cachedLines?: string[];
261
262 // Colors - using proper reset sequences
263 private dim = (s: string) => `\x1b[2m${s}\x1b[0m`;
264 private bold = (s: string) => `\x1b[1m${s}\x1b[0m`;
265 private cyan = (s: string) => `\x1b[36m${s}\x1b[0m`;
266 private green = (s: string) => `\x1b[32m${s}\x1b[0m`;
267 private yellow = (s: string) => `\x1b[33m${s}\x1b[0m`;
268 private gray = (s: string) => `\x1b[90m${s}\x1b[0m`;
269
270 constructor(questions: ExtractedQuestion[], tui: TUI, onDone: (result: string | null) => void) {
271 this.questions = questions;
272 this.answers = questions.map(() => "");
273 this.tui = tui;
274 this.onDone = onDone;
275
276 // Create a minimal theme for the editor
277 const editorTheme: EditorTheme = {
278 borderColor: this.dim,
279 selectList: {
280 selectedPrefix: this.cyan,
281 selectedText: (s: string) => `\x1b[44m${s}\x1b[0m`,
282 description: this.gray,
283 scrollInfo: this.dim,
284 noMatch: this.dim,
285 },
286 };
287
288 this.editor = new Editor(tui, editorTheme);
289 // Disable the editor's built-in submit (which clears the editor)
290 // We'll handle Enter ourselves to preserve the text
291 this.editor.disableSubmit = true;
292 this.editor.onChange = () => {
293 this.invalidate();
294 this.tui.requestRender();
295 };
296
297 this.notesEditor = new Editor(tui, editorTheme);
298 this.notesEditor.onSubmit = (value) => {
299 this.notesText = value.trim();
300 this.notesMode = false;
301 this.invalidate();
302 this.tui.requestRender();
303 };
304 }
305
306 private saveCurrentAnswer(): void {
307 this.answers[this.currentIndex] = this.editor.getText();
308 }
309
310 private navigateTo(index: number): void {
311 if (index < 0 || index >= this.questions.length) return;
312 this.saveCurrentAnswer();
313 this.currentIndex = index;
314 this.editor.setText(this.answers[index] || "");
315 this.invalidate();
316 }
317
318 private submit(): void {
319 this.saveCurrentAnswer();
320
321 // Build the response text
322 const parts: string[] = [];
323 parts.push(`<qna>`);
324 for (let i = 0; i < this.questions.length; i++) {
325 const q = this.questions[i];
326 const a = this.answers[i]?.trim() || "(no answer)";
327 parts.push(`<q n="${i}">${escapeXml(q.question)}</q>`);
328 parts.push(`<a n="${i}">${escapeXml(a)}</a>`);
329 }
330 parts.push(`</qna>`);
331 if (this.notesText) {
332 parts.push(`\n<note>${escapeXml(this.notesText)}</note>`);
333 }
334
335 this.onDone(parts.join("\n").trim());
336 }
337
338 private cancel(): void {
339 this.onDone(null);
340 }
341
342 invalidate(): void {
343 this.cachedWidth = undefined;
344 this.cachedLines = undefined;
345 }
346
347 handleInput(data: string): void {
348 // Notes mode: route to notes editor
349 if (this.notesMode) {
350 if (matchesKey(data, Key.escape)) {
351 this.notesMode = false;
352 this.notesEditor.setText(this.notesText);
353 this.invalidate();
354 this.tui.requestRender();
355 return;
356 }
357 this.notesEditor.handleInput(data);
358 this.invalidate();
359 this.tui.requestRender();
360 return;
361 }
362
363 // Handle confirmation dialog
364 if (this.showingConfirmation) {
365 if (matchesKey(data, Key.enter) || data.toLowerCase() === "y") {
366 this.submit();
367 return;
368 }
369 if (matchesKey(data, Key.escape) || matchesKey(data, Key.ctrl("c")) || data.toLowerCase() === "n") {
370 this.showingConfirmation = false;
371 this.invalidate();
372 this.tui.requestRender();
373 return;
374 }
375 return;
376 }
377
378 // Global navigation and commands
379 if (matchesKey(data, Key.escape) || matchesKey(data, Key.ctrl("c"))) {
380 this.cancel();
381 return;
382 }
383
384 // Activate notes editor (available when not typing an answer)
385 if (matchesKey(data, Key.alt("n"))) {
386 this.notesMode = true;
387 this.notesEditor.setText(this.notesText);
388 this.invalidate();
389 this.tui.requestRender();
390 return;
391 }
392
393 // Tab / Shift+Tab for navigation
394 if (matchesKey(data, Key.tab)) {
395 if (this.currentIndex < this.questions.length - 1) {
396 this.navigateTo(this.currentIndex + 1);
397 this.tui.requestRender();
398 }
399 return;
400 }
401 if (matchesKey(data, Key.shift("tab"))) {
402 if (this.currentIndex > 0) {
403 this.navigateTo(this.currentIndex - 1);
404 this.tui.requestRender();
405 }
406 return;
407 }
408
409 // Arrow up/down for question navigation when editor is empty
410 // (Editor handles its own cursor navigation when there's content)
411 if (matchesKey(data, Key.up) && this.editor.getText() === "") {
412 if (this.currentIndex > 0) {
413 this.navigateTo(this.currentIndex - 1);
414 this.tui.requestRender();
415 return;
416 }
417 }
418 if (matchesKey(data, Key.down) && this.editor.getText() === "") {
419 if (this.currentIndex < this.questions.length - 1) {
420 this.navigateTo(this.currentIndex + 1);
421 this.tui.requestRender();
422 return;
423 }
424 }
425
426 // Handle Enter ourselves (editor's submit is disabled)
427 // Plain Enter moves to next question or shows confirmation on last question
428 // Shift+Enter adds a newline (handled by editor)
429 if (matchesKey(data, Key.enter) && !matchesKey(data, Key.shift("enter"))) {
430 this.saveCurrentAnswer();
431 if (this.currentIndex < this.questions.length - 1) {
432 this.navigateTo(this.currentIndex + 1);
433 } else {
434 // On last question - show confirmation
435 this.showingConfirmation = true;
436 }
437 this.invalidate();
438 this.tui.requestRender();
439 return;
440 }
441
442 // Pass to editor
443 this.editor.handleInput(data);
444 this.invalidate();
445 this.tui.requestRender();
446 }
447
448 render(width: number): string[] {
449 if (this.cachedLines && this.cachedWidth === width) {
450 return this.cachedLines;
451 }
452
453 const lines: string[] = [];
454 const boxWidth = Math.min(width - 4, 120); // Allow wider box
455 const contentWidth = boxWidth - 4; // 2 chars padding on each side
456
457 // Helper to create horizontal lines (dim the whole thing at once)
458 const horizontalLine = (count: number) => "─".repeat(count);
459
460 // Helper to create a box line
461 const boxLine = (content: string, leftPad: number = 2): string => {
462 const paddedContent = " ".repeat(leftPad) + content;
463 const contentLen = visibleWidth(paddedContent);
464 const rightPad = Math.max(0, boxWidth - contentLen - 2);
465 return this.dim("│") + paddedContent + " ".repeat(rightPad) + this.dim("│");
466 };
467
468 const emptyBoxLine = (): string => {
469 return this.dim("│") + " ".repeat(boxWidth - 2) + this.dim("│");
470 };
471
472 const padToWidth = (line: string): string => {
473 const len = visibleWidth(line);
474 return line + " ".repeat(Math.max(0, width - len));
475 };
476
477 // Title
478 lines.push(padToWidth(this.dim(`â•${horizontalLine(boxWidth - 2)}â•®`)));
479 const title = `${this.bold(this.cyan("Questions"))} ${this.dim(`(${this.currentIndex + 1}/${this.questions.length})`)}`;
480 lines.push(padToWidth(boxLine(title)));
481 lines.push(padToWidth(this.dim(`├${horizontalLine(boxWidth - 2)}┤`)));
482
483 // Progress indicator
484 const progressParts: string[] = [];
485 for (let i = 0; i < this.questions.length; i++) {
486 const answered = (this.answers[i]?.trim() || "").length > 0;
487 const current = i === this.currentIndex;
488 if (current) {
489 progressParts.push(this.cyan("â—Ź"));
490 } else if (answered) {
491 progressParts.push(this.green("â—Ź"));
492 } else {
493 progressParts.push(this.dim("â—‹"));
494 }
495 }
496 lines.push(padToWidth(boxLine(progressParts.join(" "))));
497 lines.push(padToWidth(emptyBoxLine()));
498
499 // Current question
500 const q = this.questions[this.currentIndex];
501 const questionText = `${this.bold("Q:")} ${q.question}`;
502 const wrappedQuestion = wrapTextWithAnsi(questionText, contentWidth);
503 for (const line of wrappedQuestion) {
504 lines.push(padToWidth(boxLine(line)));
505 }
506
507 // Context if present
508 if (q.context) {
509 lines.push(padToWidth(emptyBoxLine()));
510 const contextText = this.gray(`> ${q.context}`);
511 const wrappedContext = wrapTextWithAnsi(contextText, contentWidth - 2);
512 for (const line of wrappedContext) {
513 lines.push(padToWidth(boxLine(line)));
514 }
515 }
516
517 lines.push(padToWidth(emptyBoxLine()));
518
519 // Render the editor component (multi-line input) with padding
520 // Skip the first and last lines (editor's own border lines)
521 const answerPrefix = this.bold("A: ");
522 const editorWidth = contentWidth - 4 - 3; // Extra padding + space for "A: "
523 const editorLines = this.editor.render(editorWidth);
524 for (let i = 1; i < editorLines.length - 1; i++) {
525 if (i === 1) {
526 // First content line gets the "A: " prefix
527 lines.push(padToWidth(boxLine(answerPrefix + editorLines[i])));
528 } else {
529 // Subsequent lines get padding to align with the first line
530 lines.push(padToWidth(boxLine(` ${editorLines[i]}`)));
531 }
532 }
533
534 // Notes section
535 if (this.notesMode) {
536 lines.push(padToWidth(emptyBoxLine()));
537 const notesLabel = `${this.cyan("✎")} ${this.bold("Note:")} ${this.dim("(Enter to save, Esc to discard)")}`;
538 lines.push(padToWidth(boxLine(notesLabel)));
539 const notesEditorWidth = contentWidth - 4;
540 const notesEditorLines = this.notesEditor.render(notesEditorWidth);
541 for (let i = 1; i < notesEditorLines.length - 1; i++) {
542 lines.push(padToWidth(boxLine(` ${notesEditorLines[i]}`)));
543 }
544 } else if (this.notesText) {
545 lines.push(padToWidth(emptyBoxLine()));
546 const savedNote = `${this.cyan("✎")} ${this.gray(this.notesText)}`;
547 const wrappedNote = wrapTextWithAnsi(savedNote, contentWidth);
548 for (const line of wrappedNote) {
549 lines.push(padToWidth(boxLine(line)));
550 }
551 lines.push(padToWidth(boxLine(this.dim("Alt+N to edit note"))));
552 } else {
553 lines.push(padToWidth(emptyBoxLine()));
554 lines.push(padToWidth(boxLine(this.dim("Alt+N to add a note"))));
555 }
556
557 lines.push(padToWidth(emptyBoxLine()));
558
559 // Confirmation dialog or footer with controls
560 if (this.showingConfirmation) {
561 lines.push(padToWidth(this.dim(`├${horizontalLine(boxWidth - 2)}┤`)));
562 const confirmMsg = `${this.yellow("Submit all answers?")} ${this.dim("(Enter/y to confirm, Esc/n to cancel)")}`;
563 lines.push(padToWidth(boxLine(truncateToWidth(confirmMsg, contentWidth))));
564 } else {
565 lines.push(padToWidth(this.dim(`├${horizontalLine(boxWidth - 2)}┤`)));
566 const controls = `${this.dim("Tab/Enter")} next · ${this.dim("Shift+Tab")} prev · ${this.dim("Shift+Enter")} newline · ${this.dim("Alt+N")} note · ${this.dim("Esc")} cancel`;
567 lines.push(padToWidth(boxLine(truncateToWidth(controls, contentWidth))));
568 }
569 lines.push(padToWidth(this.dim(`â•°${horizontalLine(boxWidth - 2)}â•Ż`)));
570
571 this.cachedWidth = width;
572 this.cachedLines = lines;
573 return lines;
574 }
575}
576
577export default function (pi: ExtensionAPI) {
578 const answerHandler = async (ctx: ExtensionContext, instruction?: string) => {
579 if (!ctx.hasUI) {
580 ctx.ui.notify("answer requires interactive mode", "error");
581 return;
582 }
583
584 if (!ctx.model) {
585 ctx.ui.notify("No model selected", "error");
586 return;
587 }
588
589 // Find the last assistant message on the current branch
590 const branch = ctx.sessionManager.getBranch();
591 let lastAssistantText: string | undefined;
592
593 for (let i = branch.length - 1; i >= 0; i--) {
594 const entry = branch[i];
595 if (entry.type === "message") {
596 const msg = entry.message;
597 if ("role" in msg && msg.role === "assistant") {
598 if (msg.stopReason !== "stop") {
599 ctx.ui.notify(`Last assistant message incomplete (${msg.stopReason})`, "error");
600 return;
601 }
602 const textParts = msg.content
603 .filter((c): c is { type: "text"; text: string } => c.type === "text")
604 .map((c) => c.text);
605 if (textParts.length > 0) {
606 lastAssistantText = textParts.join("\n");
607 break;
608 }
609 }
610 }
611 }
612
613 if (!lastAssistantText) {
614 ctx.ui.notify("No assistant messages found", "error");
615 return;
616 }
617
618 // Capture narrowed values so closures can use them without non-null assertions.
619 const sessionModel = ctx.model;
620 const assistantText = lastAssistantText;
621
622 // Run extraction with loader UI.
623 // The result distinguishes user cancellation from errors so we can
624 // show the right message after the custom UI closes.
625 type ExtractionOutcome =
626 | { kind: "ok"; result: ExtractionResult }
627 | { kind: "cancelled" }
628 | { kind: "error"; message: string };
629
630 const outcome = await ctx.ui.custom<ExtractionOutcome>((tui, theme, _kb, done) => {
631 const extractionModel = resolveExtractionModel(ctx, sessionModel);
632 const loader = new BorderedLoader(tui, theme, `Extracting questions using ${extractionModel.name}...`);
633
634 // Guard against double-completion: loader.onAbort fires on user
635 // cancel, but the in-flight promise may also resolve/reject after
636 // the abort. Only the first call to finish() takes effect.
637 let finished = false;
638 const finish = (result: ExtractionOutcome) => {
639 if (finished) return;
640 finished = true;
641 done(result);
642 };
643
644 loader.onAbort = () => finish({ kind: "cancelled" });
645
646 const tryExtract = async (model: ReturnType<typeof resolveExtractionModel>) => {
647 const auth = await ctx.modelRegistry.getApiKeyAndHeaders(model);
648 if (!auth.ok) {
649 return { kind: "model_error" as const, model };
650 }
651 const instructionBlock = instruction
652 ? `<user_instruction>\n${instruction}\n</user_instruction>\n\nExtract questions from the assistant's message based on the user's instruction above.`
653 : "Extract questions from the assistant's message for the user to fill out.";
654
655 const userMessage: UserMessage = {
656 role: "user",
657 content: [
658 {
659 type: "text",
660 text: `<last_assistant_message>\n${assistantText}\n</last_assistant_message>\n\n${instructionBlock}`,
661 },
662 ],
663 timestamp: Date.now(),
664 };
665
666 const response = await complete(
667 model,
668 {
669 systemPrompt: SYSTEM_PROMPT,
670 messages: [userMessage],
671 tools: [QUESTION_EXTRACTION_TOOL],
672 },
673 { apiKey: auth.apiKey, headers: auth.headers, signal: loader.signal },
674 );
675
676 if (response.stopReason === "aborted") {
677 return { kind: "cancelled" as const };
678 }
679
680 if (response.stopReason === "error" || response.content.length === 0) {
681 return { kind: "model_error" as const, model };
682 }
683
684 const parsed = parseToolCallResult(response);
685 if (!parsed) {
686 return {
687 kind: "parse_error" as const,
688 message: `${model.name} did not call extract_questions tool`,
689 };
690 }
691 return { kind: "ok" as const, result: parsed };
692 };
693
694 const doExtract = async () => {
695 let result = await tryExtract(extractionModel);
696
697 // If the preferred model errored and it's not already the
698 // session model, fall back to the session model.
699 if (result.kind === "model_error" && extractionModel.id !== sessionModel.id) {
700 ctx.ui.notify(`${extractionModel.name} unavailable, falling back to ${sessionModel.name}...`, "warning");
701 result = await tryExtract(sessionModel);
702 }
703
704 switch (result.kind) {
705 case "ok":
706 return finish({ kind: "ok", result: result.result });
707 case "cancelled":
708 return finish({ kind: "cancelled" });
709 case "model_error":
710 return finish({
711 kind: "error",
712 message: `${result.model.name} returned an error with no content`,
713 });
714 case "parse_error":
715 return finish({ kind: "error", message: result.message });
716 }
717 };
718
719 doExtract().catch((err) =>
720 finish({
721 kind: "error",
722 message: `${err?.message ?? err}`,
723 }),
724 );
725
726 return loader;
727 });
728
729 if (outcome.kind === "cancelled") {
730 ctx.ui.notify("Cancelled", "info");
731 return;
732 }
733 if (outcome.kind === "error") {
734 ctx.ui.notify(`Extraction failed: ${outcome.message}`, "error");
735 return;
736 }
737
738 const extractionResult = outcome.result;
739
740 if (extractionResult.questions.length === 0) {
741 ctx.ui.notify("No questions found in the last message", "info");
742 return;
743 }
744
745 // Show the Q&A component
746 const answersResult = await ctx.ui.custom<string | null>((tui, _theme, _kb, done) => {
747 return new QnAComponent(extractionResult.questions, tui, done);
748 });
749
750 if (answersResult === null) {
751 ctx.ui.notify("Cancelled", "info");
752 return;
753 }
754
755 // Send the answers directly as a message and trigger a turn
756 pi.sendMessage(
757 {
758 customType: "answers",
759 content: `I answered your questions in the following way:\n\n${answersResult}`,
760 display: true,
761 },
762 { triggerTurn: true },
763 );
764 };
765
766 pi.registerCommand("answer", {
767 description:
768 "Extract questions from last assistant message into interactive Q&A. Optional: provide instructions for what to extract.",
769 handler: (args, ctx) => answerHandler(ctx, args?.trim() || undefined),
770 });
771
772 pi.registerShortcut("ctrl+.", {
773 description: "Extract and answer questions",
774 handler: answerHandler,
775 });
776}