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