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, fallback: NonNullable<ExtensionContext["model"]>) {
192 const available = ctx.modelRegistry.getAvailable();
193 const preferred = available.find((m) => m.id === PREFERRED_EXTRACTION_MODEL_ID);
194 return preferred ?? fallback;
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 selectedPrefix: this.cyan,
276 selectedText: (s: string) => `\x1b[44m${s}\x1b[0m`,
277 description: this.gray,
278 scrollInfo: this.dim,
279 noMatch: this.dim,
280 },
281 };
282
283 this.editor = new Editor(tui, editorTheme);
284 // Disable the editor's built-in submit (which clears the editor)
285 // We'll handle Enter ourselves to preserve the text
286 this.editor.disableSubmit = true;
287 this.editor.onChange = () => {
288 this.invalidate();
289 this.tui.requestRender();
290 };
291
292 this.notesEditor = new Editor(tui, editorTheme);
293 this.notesEditor.onSubmit = (value) => {
294 this.notesText = value.trim();
295 this.notesMode = false;
296 this.invalidate();
297 this.tui.requestRender();
298 };
299 }
300
301 private saveCurrentAnswer(): void {
302 this.answers[this.currentIndex] = this.editor.getText();
303 }
304
305 private navigateTo(index: number): void {
306 if (index < 0 || index >= this.questions.length) return;
307 this.saveCurrentAnswer();
308 this.currentIndex = index;
309 this.editor.setText(this.answers[index] || "");
310 this.invalidate();
311 }
312
313 private submit(): void {
314 this.saveCurrentAnswer();
315
316 // Build the response text
317 const parts: string[] = [];
318 parts.push(`<qna>`);
319 for (let i = 0; i < this.questions.length; i++) {
320 const q = this.questions[i];
321 const a = this.answers[i]?.trim() || "(no answer)";
322 parts.push(`<q n="${i}">${q.question}</q>`);
323 parts.push(`<a n="${i}">${a}</a>`);
324 }
325 parts.push(`</qna>`);
326 if (this.notesText) {
327 parts.push(`\n<note>${this.notesText}</note>`);
328 }
329
330 this.onDone(parts.join("\n").trim());
331 }
332
333 private cancel(): void {
334 this.onDone(null);
335 }
336
337 invalidate(): void {
338 this.cachedWidth = undefined;
339 this.cachedLines = undefined;
340 }
341
342 handleInput(data: string): void {
343 // Notes mode: route to notes editor
344 if (this.notesMode) {
345 if (matchesKey(data, Key.escape)) {
346 this.notesMode = false;
347 this.notesEditor.setText(this.notesText);
348 this.invalidate();
349 this.tui.requestRender();
350 return;
351 }
352 this.notesEditor.handleInput(data);
353 this.invalidate();
354 this.tui.requestRender();
355 return;
356 }
357
358 // Handle confirmation dialog
359 if (this.showingConfirmation) {
360 if (matchesKey(data, Key.enter) || data.toLowerCase() === "y") {
361 this.submit();
362 return;
363 }
364 if (matchesKey(data, Key.escape) || matchesKey(data, Key.ctrl("c")) || data.toLowerCase() === "n") {
365 this.showingConfirmation = false;
366 this.invalidate();
367 this.tui.requestRender();
368 return;
369 }
370 return;
371 }
372
373 // Global navigation and commands
374 if (matchesKey(data, Key.escape) || matchesKey(data, Key.ctrl("c"))) {
375 this.cancel();
376 return;
377 }
378
379 // Activate notes editor (available when not typing an answer)
380 if (matchesKey(data, Key.alt("n"))) {
381 this.notesMode = true;
382 this.notesEditor.setText(this.notesText);
383 this.invalidate();
384 this.tui.requestRender();
385 return;
386 }
387
388 // Tab / Shift+Tab for navigation
389 if (matchesKey(data, Key.tab)) {
390 if (this.currentIndex < this.questions.length - 1) {
391 this.navigateTo(this.currentIndex + 1);
392 this.tui.requestRender();
393 }
394 return;
395 }
396 if (matchesKey(data, Key.shift("tab"))) {
397 if (this.currentIndex > 0) {
398 this.navigateTo(this.currentIndex - 1);
399 this.tui.requestRender();
400 }
401 return;
402 }
403
404 // Arrow up/down for question navigation when editor is empty
405 // (Editor handles its own cursor navigation when there's content)
406 if (matchesKey(data, Key.up) && this.editor.getText() === "") {
407 if (this.currentIndex > 0) {
408 this.navigateTo(this.currentIndex - 1);
409 this.tui.requestRender();
410 return;
411 }
412 }
413 if (matchesKey(data, Key.down) && this.editor.getText() === "") {
414 if (this.currentIndex < this.questions.length - 1) {
415 this.navigateTo(this.currentIndex + 1);
416 this.tui.requestRender();
417 return;
418 }
419 }
420
421 // Handle Enter ourselves (editor's submit is disabled)
422 // Plain Enter moves to next question or shows confirmation on last question
423 // Shift+Enter adds a newline (handled by editor)
424 if (matchesKey(data, Key.enter) && !matchesKey(data, Key.shift("enter"))) {
425 this.saveCurrentAnswer();
426 if (this.currentIndex < this.questions.length - 1) {
427 this.navigateTo(this.currentIndex + 1);
428 } else {
429 // On last question - show confirmation
430 this.showingConfirmation = true;
431 }
432 this.invalidate();
433 this.tui.requestRender();
434 return;
435 }
436
437 // Pass to editor
438 this.editor.handleInput(data);
439 this.invalidate();
440 this.tui.requestRender();
441 }
442
443 render(width: number): string[] {
444 if (this.cachedLines && this.cachedWidth === width) {
445 return this.cachedLines;
446 }
447
448 const lines: string[] = [];
449 const boxWidth = Math.min(width - 4, 120); // Allow wider box
450 const contentWidth = boxWidth - 4; // 2 chars padding on each side
451
452 // Helper to create horizontal lines (dim the whole thing at once)
453 const horizontalLine = (count: number) => "─".repeat(count);
454
455 // Helper to create a box line
456 const boxLine = (content: string, leftPad: number = 2): string => {
457 const paddedContent = " ".repeat(leftPad) + content;
458 const contentLen = visibleWidth(paddedContent);
459 const rightPad = Math.max(0, boxWidth - contentLen - 2);
460 return this.dim("│") + paddedContent + " ".repeat(rightPad) + this.dim("│");
461 };
462
463 const emptyBoxLine = (): string => {
464 return this.dim("│") + " ".repeat(boxWidth - 2) + this.dim("│");
465 };
466
467 const padToWidth = (line: string): string => {
468 const len = visibleWidth(line);
469 return line + " ".repeat(Math.max(0, width - len));
470 };
471
472 // Title
473 lines.push(padToWidth(this.dim(`â•${horizontalLine(boxWidth - 2)}â•®`)));
474 const title = `${this.bold(this.cyan("Questions"))} ${this.dim(`(${this.currentIndex + 1}/${this.questions.length})`)}`;
475 lines.push(padToWidth(boxLine(title)));
476 lines.push(padToWidth(this.dim(`├${horizontalLine(boxWidth - 2)}┤`)));
477
478 // Progress indicator
479 const progressParts: string[] = [];
480 for (let i = 0; i < this.questions.length; i++) {
481 const answered = (this.answers[i]?.trim() || "").length > 0;
482 const current = i === this.currentIndex;
483 if (current) {
484 progressParts.push(this.cyan("â—Ź"));
485 } else if (answered) {
486 progressParts.push(this.green("â—Ź"));
487 } else {
488 progressParts.push(this.dim("â—‹"));
489 }
490 }
491 lines.push(padToWidth(boxLine(progressParts.join(" "))));
492 lines.push(padToWidth(emptyBoxLine()));
493
494 // Current question
495 const q = this.questions[this.currentIndex];
496 const questionText = `${this.bold("Q:")} ${q.question}`;
497 const wrappedQuestion = wrapTextWithAnsi(questionText, contentWidth);
498 for (const line of wrappedQuestion) {
499 lines.push(padToWidth(boxLine(line)));
500 }
501
502 // Context if present
503 if (q.context) {
504 lines.push(padToWidth(emptyBoxLine()));
505 const contextText = this.gray(`> ${q.context}`);
506 const wrappedContext = wrapTextWithAnsi(contextText, contentWidth - 2);
507 for (const line of wrappedContext) {
508 lines.push(padToWidth(boxLine(line)));
509 }
510 }
511
512 lines.push(padToWidth(emptyBoxLine()));
513
514 // Render the editor component (multi-line input) with padding
515 // Skip the first and last lines (editor's own border lines)
516 const answerPrefix = this.bold("A: ");
517 const editorWidth = contentWidth - 4 - 3; // Extra padding + space for "A: "
518 const editorLines = this.editor.render(editorWidth);
519 for (let i = 1; i < editorLines.length - 1; i++) {
520 if (i === 1) {
521 // First content line gets the "A: " prefix
522 lines.push(padToWidth(boxLine(answerPrefix + editorLines[i])));
523 } else {
524 // Subsequent lines get padding to align with the first line
525 lines.push(padToWidth(boxLine(` ${editorLines[i]}`)));
526 }
527 }
528
529 // Notes section
530 if (this.notesMode) {
531 lines.push(padToWidth(emptyBoxLine()));
532 const notesLabel = `${this.cyan("✎")} ${this.bold("Note:")} ${this.dim("(Enter to save, Esc to discard)")}`;
533 lines.push(padToWidth(boxLine(notesLabel)));
534 const notesEditorWidth = contentWidth - 4;
535 const notesEditorLines = this.notesEditor.render(notesEditorWidth);
536 for (let i = 1; i < notesEditorLines.length - 1; i++) {
537 lines.push(padToWidth(boxLine(` ${notesEditorLines[i]}`)));
538 }
539 } else if (this.notesText) {
540 lines.push(padToWidth(emptyBoxLine()));
541 const savedNote = `${this.cyan("✎")} ${this.gray(this.notesText)}`;
542 const wrappedNote = wrapTextWithAnsi(savedNote, contentWidth);
543 for (const line of wrappedNote) {
544 lines.push(padToWidth(boxLine(line)));
545 }
546 lines.push(padToWidth(boxLine(this.dim("Alt+N to edit note"))));
547 } else {
548 lines.push(padToWidth(emptyBoxLine()));
549 lines.push(padToWidth(boxLine(this.dim("Alt+N to add a note"))));
550 }
551
552 lines.push(padToWidth(emptyBoxLine()));
553
554 // Confirmation dialog or footer with controls
555 if (this.showingConfirmation) {
556 lines.push(padToWidth(this.dim(`├${horizontalLine(boxWidth - 2)}┤`)));
557 const confirmMsg = `${this.yellow("Submit all answers?")} ${this.dim("(Enter/y to confirm, Esc/n to cancel)")}`;
558 lines.push(padToWidth(boxLine(truncateToWidth(confirmMsg, contentWidth))));
559 } else {
560 lines.push(padToWidth(this.dim(`├${horizontalLine(boxWidth - 2)}┤`)));
561 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`;
562 lines.push(padToWidth(boxLine(truncateToWidth(controls, contentWidth))));
563 }
564 lines.push(padToWidth(this.dim(`â•°${horizontalLine(boxWidth - 2)}â•Ż`)));
565
566 this.cachedWidth = width;
567 this.cachedLines = lines;
568 return lines;
569 }
570}
571
572export default function (pi: ExtensionAPI) {
573 const answerHandler = async (ctx: ExtensionContext, instruction?: string) => {
574 if (!ctx.hasUI) {
575 ctx.ui.notify("answer requires interactive mode", "error");
576 return;
577 }
578
579 if (!ctx.model) {
580 ctx.ui.notify("No model selected", "error");
581 return;
582 }
583
584 // Find the last assistant message on the current branch
585 const branch = ctx.sessionManager.getBranch();
586 let lastAssistantText: string | undefined;
587
588 for (let i = branch.length - 1; i >= 0; i--) {
589 const entry = branch[i];
590 if (entry.type === "message") {
591 const msg = entry.message;
592 if ("role" in msg && msg.role === "assistant") {
593 if (msg.stopReason !== "stop") {
594 ctx.ui.notify(`Last assistant message incomplete (${msg.stopReason})`, "error");
595 return;
596 }
597 const textParts = msg.content
598 .filter((c): c is { type: "text"; text: string } => c.type === "text")
599 .map((c) => c.text);
600 if (textParts.length > 0) {
601 lastAssistantText = textParts.join("\n");
602 break;
603 }
604 }
605 }
606 }
607
608 if (!lastAssistantText) {
609 ctx.ui.notify("No assistant messages found", "error");
610 return;
611 }
612
613 // Capture narrowed values so closures can use them without non-null assertions.
614 const sessionModel = ctx.model;
615 const assistantText = lastAssistantText;
616
617 // Run extraction with loader UI.
618 // The result distinguishes user cancellation from errors so we can
619 // show the right message after the custom UI closes.
620 type ExtractionOutcome =
621 | { kind: "ok"; result: ExtractionResult }
622 | { kind: "cancelled" }
623 | { kind: "error"; message: string };
624
625 const outcome = await ctx.ui.custom<ExtractionOutcome>((tui, theme, _kb, done) => {
626 const extractionModel = resolveExtractionModel(ctx, sessionModel);
627 const loader = new BorderedLoader(tui, theme, `Extracting questions using ${extractionModel.name}...`);
628 loader.onAbort = () => done({ kind: "cancelled" });
629
630 const tryExtract = async (model: ReturnType<typeof resolveExtractionModel>) => {
631 const auth = await ctx.modelRegistry.getApiKeyAndHeaders(model);
632 if (!auth.ok) {
633 return { kind: "model_error" as const, model };
634 }
635 const instructionBlock = instruction
636 ? `<user_instruction>\n${instruction}\n</user_instruction>\n\nExtract questions from the assistant's message based on the user's instruction above.`
637 : "Extract questions from the assistant's message for the user to fill out.";
638
639 const userMessage: UserMessage = {
640 role: "user",
641 content: [
642 {
643 type: "text",
644 text: `<last_assistant_message>\n${assistantText}\n</last_assistant_message>\n\n${instructionBlock}`,
645 },
646 ],
647 timestamp: Date.now(),
648 };
649
650 const response = await complete(
651 model,
652 {
653 systemPrompt: SYSTEM_PROMPT,
654 messages: [userMessage],
655 tools: [QUESTION_EXTRACTION_TOOL],
656 },
657 { apiKey: auth.apiKey, headers: auth.headers, signal: loader.signal },
658 );
659
660 if (response.stopReason === "aborted") {
661 return { kind: "cancelled" as const };
662 }
663
664 if (response.stopReason === "error" || response.content.length === 0) {
665 return { kind: "model_error" as const, model };
666 }
667
668 const parsed = parseToolCallResult(response);
669 if (!parsed) {
670 return {
671 kind: "parse_error" as const,
672 message: `${model.name} did not call extract_questions tool`,
673 };
674 }
675 return { kind: "ok" as const, result: parsed };
676 };
677
678 const doExtract = async () => {
679 let result = await tryExtract(extractionModel);
680
681 // If the preferred model errored and it's not already the
682 // session model, fall back to the session model.
683 if (result.kind === "model_error" && extractionModel.id !== sessionModel.id) {
684 ctx.ui.notify(`${extractionModel.name} unavailable, falling back to ${sessionModel.name}...`, "warning");
685 result = await tryExtract(sessionModel);
686 }
687
688 switch (result.kind) {
689 case "ok":
690 return done({ kind: "ok", result: result.result });
691 case "cancelled":
692 return done({ kind: "cancelled" });
693 case "model_error":
694 return done({
695 kind: "error",
696 message: `${result.model.name} returned an error with no content`,
697 });
698 case "parse_error":
699 return done({ kind: "error", message: result.message });
700 }
701 };
702
703 doExtract().catch((err) =>
704 done({
705 kind: "error",
706 message: `${err?.message ?? err}`,
707 }),
708 );
709
710 return loader;
711 });
712
713 if (outcome.kind === "cancelled") {
714 ctx.ui.notify("Cancelled", "info");
715 return;
716 }
717 if (outcome.kind === "error") {
718 ctx.ui.notify(`Extraction failed: ${outcome.message}`, "error");
719 return;
720 }
721
722 const extractionResult = outcome.result;
723
724 if (extractionResult.questions.length === 0) {
725 ctx.ui.notify("No questions found in the last message", "info");
726 return;
727 }
728
729 // Show the Q&A component
730 const answersResult = await ctx.ui.custom<string | null>((tui, _theme, _kb, done) => {
731 return new QnAComponent(extractionResult.questions, tui, done);
732 });
733
734 if (answersResult === null) {
735 ctx.ui.notify("Cancelled", "info");
736 return;
737 }
738
739 // Send the answers directly as a message and trigger a turn
740 pi.sendMessage(
741 {
742 customType: "answers",
743 content: `I answered your questions in the following way:\n\n${answersResult}`,
744 display: true,
745 },
746 { triggerTurn: true },
747 );
748 };
749
750 pi.registerCommand("answer", {
751 description:
752 "Extract questions from last assistant message into interactive Q&A. Optional: provide instructions for what to extract.",
753 handler: (args, ctx) => answerHandler(ctx, args?.trim() || undefined),
754 });
755
756 pi.registerShortcut("ctrl+.", {
757 description: "Extract and answer questions",
758 handler: answerHandler,
759 });
760}