@@ -0,0 +1,759 @@
+// SPDX-FileCopyrightText: Amolith <amolith@secluded.site>
+// SPDX-FileCopyrightText: Armin Ronacher <armin.ronacher@active-4.com>
+//
+// SPDX-License-Identifier: Apache-2.0
+
+/**
+ * Q&A extraction hook - extracts questions from assistant responses
+ *
+ * Custom interactive TUI for answering questions.
+ *
+ * Demonstrates the "prompt generator" pattern with custom TUI:
+ * 1. /answer command gets the last assistant message
+ * 2. Shows a spinner while extracting questions as structured JSON
+ * 3. Presents an interactive TUI to navigate and answer questions
+ * 4. Submits the compiled answers when done
+ */
+
+import { complete, type Tool, type UserMessage } from "@mariozechner/pi-ai";
+import type { ExtensionAPI, ExtensionContext } from "@mariozechner/pi-coding-agent";
+import { BorderedLoader } from "@mariozechner/pi-coding-agent";
+import {
+ type Component,
+ Editor,
+ type EditorTheme,
+ Key,
+ matchesKey,
+ truncateToWidth,
+ type TUI,
+ visibleWidth,
+ wrapTextWithAnsi,
+} from "@mariozechner/pi-tui";
+import { Type } from "@sinclair/typebox";
+
+// Structured output format for question extraction
+interface ExtractedQuestion {
+ question: string;
+ context?: string;
+}
+
+interface ExtractionResult {
+ questions: ExtractedQuestion[];
+}
+
+const 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.
+
+Each 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.
+
+Rules:
+- Each question must make sense on its own without referring back to the message (no "Issue 1", "the above", "as mentioned")
+- Context should include all relevant details, options, and reasoning needed to answer — it replaces reading the original message
+- Merge duplicate or overlapping questions into one (e.g. an analysis section and a plan section about the same topic become one question)
+- Capture both explicit questions ("Which approach?") and implicit ones (proposals that need approval, recommendations that need confirmation)
+- Keep questions in logical order
+- If no questions are found, call the tool with an empty questions array
+
+<examples>
+<example>
+<description>Simple explicit question</description>
+<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>
+<output>
+{"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."}
+</output>
+</example>
+
+<example>
+<description>Implicit proposal needing approval</description>
+<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>
+<output>
+{"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."}
+</output>
+</example>
+
+<example>
+<description>Long message with analysis and plan sections that overlap — questions must be merged, deduplicated, and fully standalone</description>
+<input>
+## Issue 1: Non-deterministic map iteration for overrides → rawArgs
+
+The problem is in RunE at line ~75: for k, vals := range overrides gives random order each run.
+
+Three good approaches:
+1. Sort the keys before iterating. Simplest, most common Go pattern.
+2. Use []restic.Flag or []struct{key, vals} instead of a map. More invasive but preserves construction order.
+3. 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.
+
+Recommend approach 3 as best: it kills the root cause.
+
+## Issue 2: Double-pointers (**screens.Snapshot etc.)
+
+The problem is buildCommandScreens needs to both return screens AND communicate which specific screen instances were created back to runInteractive. Double-pointers are gnarly.
+
+Simpler approach: return a struct holding typed screen pointers. The ResolveFunc closure captures the struct.
+
+## Issue 3: Picker height overflow
+
+The 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.
+
+## Issue 2.1: huh form boilerplate
+
+The buildForm/Init/Update boilerplate is repetitive across Overwrite, Target, Preset, Snapshot. Code is simple and clear in each screen. Recommendation: skip — borderline DRY.
+
+## Issue 2.2: drain test helpers
+
+The 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)
+
+## Issue 3.1: Snapshot theme handler clears user input
+
+When BackgroundColorMsg arrives during phaseSnapshotManual, buildManualForm() resets s.entered = "". Fix: preserve s.entered before rebuilding. Also needs rebuildSelectForm() call during phaseSnapshotSelecting.
+
+## Issue 4: WindowSizeMsg missed during phaseFileLoading
+
+If 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.
+
+## Issue 5.1: Same as Issue 1
+
+Same root cause — map iteration in RunE. Fix for Issue 1 handles this.
+
+## Issue 6.1: Duplicated constant
+
+Both internal/config and internal/restic define the same commandSuffix constant.
+
+## Issue 6.2: Inconsistent empty-string semantics
+
+rc.Environ treats "" as present (checks map key existence), but os.Getenv checks != "" so an empty env var counts as absent.
+
+## Plan
+
+1. Remove legacy restore prompts (promptRestore, promptSnapshotID, promptFileSelection, printRestoreSummary, standalone pickerModel)
+2. Eliminate map→rawArgs round-trip (Issues 1 and 5.1)
+3. Replace double-pointers with return struct (Issue 2)
+4. Extract generic drain test helper (Issue 2.2)
+5. Fix snapshot theme handler (Issue 3.1)
+6. Store WindowSizeMsg during loading (Issue 4)
+7. Share commandSuffix constant (Issue 6.1)
+8. Fix inconsistent empty-string check (Issue 6.2)
+
+NOT planning to do Issue 2.1 (huh form boilerplate extraction).
+</input>
+<output>
+{"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."}
+{"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."}
+{"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."}
+{"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."}
+{"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."}
+{"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."}
+{"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."}
+{"question": "Should the duplicated commandSuffix constant be extracted to a shared package?", "context": "Both internal/config and internal/restic define the same constant."}
+{"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."}
+{"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."}
+</output>
+</example>
+</examples>`;
+
+/**
+ * Tool definition for structured question extraction. Models produce
+ * structured tool-call arguments far more reliably than free-form JSON
+ * in a text response.
+ */
+const QUESTION_EXTRACTION_TOOL: Tool = {
+ name: "extract_questions",
+ description:
+ "Extract questions from the assistant's message. Each question is a self-contained object the user will read instead of the original message.",
+ parameters: Type.Object({
+ questions: Type.Array(
+ Type.Object({
+ question: Type.String({
+ description: "A complete, standalone question. Must make sense without reading the original message.",
+ }),
+ context: Type.Optional(
+ Type.String({
+ description:
+ "Self-contained context with all details, options, and reasoning needed to answer the question.",
+ }),
+ ),
+ }),
+ {
+ description:
+ "Array of question objects. Merge overlapping questions. Each item MUST be an object with 'question' and optional 'context' strings, NOT a plain string.",
+ },
+ ),
+ }),
+};
+
+// Preferred model for extraction — lightweight is fine for structured JSON output.
+// Falls back to the session model if this one isn't available.
+const PREFERRED_EXTRACTION_MODEL_ID = "nemotron-3-super-120b-a12b";
+
+/**
+ * Resolve the extraction model: prefer a lightweight model from the registry,
+ * fall back to the current session model.
+ */
+function resolveExtractionModel(ctx: ExtensionContext) {
+ const available = ctx.modelRegistry.getAvailable();
+ const preferred = available.find((m) => m.id === PREFERRED_EXTRACTION_MODEL_ID);
+ return preferred ?? ctx.model!;
+}
+
+/**
+ * Parse a tool call response into an ExtractionResult.
+ */
+function parseToolCallResult(response: Awaited<ReturnType<typeof complete>>): ExtractionResult | null {
+ const toolCall = response.content.find((c) => c.type === "toolCall" && c.name === "extract_questions");
+
+ if (!toolCall || toolCall.type !== "toolCall") {
+ console.error("Model did not call extract_questions:", response.content);
+ return null;
+ }
+
+ const args = toolCall.arguments as Record<string, unknown>;
+ if (!Array.isArray(args.questions)) {
+ console.error("[answer] expected questions array, got:", typeof args.questions, args);
+ return null;
+ }
+
+ // Validate each question item — model may return plain strings instead
+ // of objects if it ignores the schema, so we handle both gracefully.
+ const questions: ExtractedQuestion[] = [];
+ for (const item of args.questions) {
+ if (typeof item === "string") {
+ // Model returned bare string instead of object — use as question text
+ questions.push({ question: item });
+ continue;
+ }
+ if (typeof item !== "object" || item === null) continue;
+ const obj = item as Record<string, unknown>;
+ if (typeof obj.question !== "string") {
+ console.error("[answer] skipping item with no 'question' string:", Object.keys(obj));
+ continue;
+ }
+ questions.push({
+ question: obj.question,
+ context: typeof obj.context === "string" ? obj.context : undefined,
+ });
+ }
+
+ return { questions };
+}
+
+/**
+ * Interactive Q&A component for answering extracted questions
+ */
+class QnAComponent implements Component {
+ private questions: ExtractedQuestion[];
+ private answers: string[];
+ private currentIndex: number = 0;
+ private editor: Editor;
+ private tui: TUI;
+ private onDone: (result: string | null) => void;
+ private showingConfirmation: boolean = false;
+ private notesMode: boolean = false;
+ private notesText: string = "";
+ private notesEditor: Editor;
+
+ // Cache
+ private cachedWidth?: number;
+ private cachedLines?: string[];
+
+ // Colors - using proper reset sequences
+ private dim = (s: string) => `\x1b[2m${s}\x1b[0m`;
+ private bold = (s: string) => `\x1b[1m${s}\x1b[0m`;
+ private cyan = (s: string) => `\x1b[36m${s}\x1b[0m`;
+ private green = (s: string) => `\x1b[32m${s}\x1b[0m`;
+ private yellow = (s: string) => `\x1b[33m${s}\x1b[0m`;
+ private gray = (s: string) => `\x1b[90m${s}\x1b[0m`;
+
+ constructor(questions: ExtractedQuestion[], tui: TUI, onDone: (result: string | null) => void) {
+ this.questions = questions;
+ this.answers = questions.map(() => "");
+ this.tui = tui;
+ this.onDone = onDone;
+
+ // Create a minimal theme for the editor
+ const editorTheme: EditorTheme = {
+ borderColor: this.dim,
+ selectList: {
+ selectedBg: (s: string) => `\x1b[44m${s}\x1b[0m`,
+ matchHighlight: this.cyan,
+ itemSecondary: this.gray,
+ },
+ };
+
+ this.editor = new Editor(tui, editorTheme);
+ // Disable the editor's built-in submit (which clears the editor)
+ // We'll handle Enter ourselves to preserve the text
+ this.editor.disableSubmit = true;
+ this.editor.onChange = () => {
+ this.invalidate();
+ this.tui.requestRender();
+ };
+
+ this.notesEditor = new Editor(tui, editorTheme);
+ this.notesEditor.onSubmit = (value) => {
+ this.notesText = value.trim();
+ this.notesMode = false;
+ this.invalidate();
+ this.tui.requestRender();
+ };
+ }
+
+ private allQuestionsAnswered(): boolean {
+ this.saveCurrentAnswer();
+ return this.answers.every((a) => (a?.trim() || "").length > 0);
+ }
+
+ private saveCurrentAnswer(): void {
+ this.answers[this.currentIndex] = this.editor.getText();
+ }
+
+ private navigateTo(index: number): void {
+ if (index < 0 || index >= this.questions.length) return;
+ this.saveCurrentAnswer();
+ this.currentIndex = index;
+ this.editor.setText(this.answers[index] || "");
+ this.invalidate();
+ }
+
+ private submit(): void {
+ this.saveCurrentAnswer();
+
+ // Build the response text
+ const parts: string[] = [];
+ parts.push(`<qna>`);
+ for (let i = 0; i < this.questions.length; i++) {
+ const q = this.questions[i];
+ const a = this.answers[i]?.trim() || "(no answer)";
+ parts.push(`<q n="${i}">${q.question}</q>`);
+ parts.push(`<a n="${i}">${a}</a>`);
+ }
+ parts.push(`</qna>`);
+ if (this.notesText) {
+ parts.push(`\n<note>${this.notesText}</note>`);
+ }
+
+ this.onDone(parts.join("\n").trim());
+ }
+
+ private cancel(): void {
+ this.onDone(null);
+ }
+
+ invalidate(): void {
+ this.cachedWidth = undefined;
+ this.cachedLines = undefined;
+ }
+
+ handleInput(data: string): void {
+ // Notes mode: route to notes editor
+ if (this.notesMode) {
+ if (matchesKey(data, Key.escape)) {
+ this.notesMode = false;
+ this.notesEditor.setText(this.notesText);
+ this.invalidate();
+ this.tui.requestRender();
+ return;
+ }
+ this.notesEditor.handleInput(data);
+ this.invalidate();
+ this.tui.requestRender();
+ return;
+ }
+
+ // Handle confirmation dialog
+ if (this.showingConfirmation) {
+ if (matchesKey(data, Key.enter) || data.toLowerCase() === "y") {
+ this.submit();
+ return;
+ }
+ if (matchesKey(data, Key.escape) || matchesKey(data, Key.ctrl("c")) || data.toLowerCase() === "n") {
+ this.showingConfirmation = false;
+ this.invalidate();
+ this.tui.requestRender();
+ return;
+ }
+ return;
+ }
+
+ // Global navigation and commands
+ if (matchesKey(data, Key.escape) || matchesKey(data, Key.ctrl("c"))) {
+ this.cancel();
+ return;
+ }
+
+ // Activate notes editor (available when not typing an answer)
+ if (matchesKey(data, Key.alt("n"))) {
+ this.notesMode = true;
+ this.notesEditor.setText(this.notesText);
+ this.invalidate();
+ this.tui.requestRender();
+ return;
+ }
+
+ // Tab / Shift+Tab for navigation
+ if (matchesKey(data, Key.tab)) {
+ if (this.currentIndex < this.questions.length - 1) {
+ this.navigateTo(this.currentIndex + 1);
+ this.tui.requestRender();
+ }
+ return;
+ }
+ if (matchesKey(data, Key.shift("tab"))) {
+ if (this.currentIndex > 0) {
+ this.navigateTo(this.currentIndex - 1);
+ this.tui.requestRender();
+ }
+ return;
+ }
+
+ // Arrow up/down for question navigation when editor is empty
+ // (Editor handles its own cursor navigation when there's content)
+ if (matchesKey(data, Key.up) && this.editor.getText() === "") {
+ if (this.currentIndex > 0) {
+ this.navigateTo(this.currentIndex - 1);
+ this.tui.requestRender();
+ return;
+ }
+ }
+ if (matchesKey(data, Key.down) && this.editor.getText() === "") {
+ if (this.currentIndex < this.questions.length - 1) {
+ this.navigateTo(this.currentIndex + 1);
+ this.tui.requestRender();
+ return;
+ }
+ }
+
+ // Handle Enter ourselves (editor's submit is disabled)
+ // Plain Enter moves to next question or shows confirmation on last question
+ // Shift+Enter adds a newline (handled by editor)
+ if (matchesKey(data, Key.enter) && !matchesKey(data, Key.shift("enter"))) {
+ this.saveCurrentAnswer();
+ if (this.currentIndex < this.questions.length - 1) {
+ this.navigateTo(this.currentIndex + 1);
+ } else {
+ // On last question - show confirmation
+ this.showingConfirmation = true;
+ }
+ this.invalidate();
+ this.tui.requestRender();
+ return;
+ }
+
+ // Pass to editor
+ this.editor.handleInput(data);
+ this.invalidate();
+ this.tui.requestRender();
+ }
+
+ render(width: number): string[] {
+ if (this.cachedLines && this.cachedWidth === width) {
+ return this.cachedLines;
+ }
+
+ const lines: string[] = [];
+ const boxWidth = Math.min(width - 4, 120); // Allow wider box
+ const contentWidth = boxWidth - 4; // 2 chars padding on each side
+
+ // Helper to create horizontal lines (dim the whole thing at once)
+ const horizontalLine = (count: number) => "─".repeat(count);
+
+ // Helper to create a box line
+ const boxLine = (content: string, leftPad: number = 2): string => {
+ const paddedContent = " ".repeat(leftPad) + content;
+ const contentLen = visibleWidth(paddedContent);
+ const rightPad = Math.max(0, boxWidth - contentLen - 2);
+ return this.dim("│") + paddedContent + " ".repeat(rightPad) + this.dim("│");
+ };
+
+ const emptyBoxLine = (): string => {
+ return this.dim("│") + " ".repeat(boxWidth - 2) + this.dim("│");
+ };
+
+ const padToWidth = (line: string): string => {
+ const len = visibleWidth(line);
+ return line + " ".repeat(Math.max(0, width - len));
+ };
+
+ // Title
+ lines.push(padToWidth(this.dim("â•" + horizontalLine(boxWidth - 2) + "â•®")));
+ const title = `${this.bold(this.cyan("Questions"))} ${this.dim(`(${this.currentIndex + 1}/${this.questions.length})`)}`;
+ lines.push(padToWidth(boxLine(title)));
+ lines.push(padToWidth(this.dim("├" + horizontalLine(boxWidth - 2) + "┤")));
+
+ // Progress indicator
+ const progressParts: string[] = [];
+ for (let i = 0; i < this.questions.length; i++) {
+ const answered = (this.answers[i]?.trim() || "").length > 0;
+ const current = i === this.currentIndex;
+ if (current) {
+ progressParts.push(this.cyan("â—Ź"));
+ } else if (answered) {
+ progressParts.push(this.green("â—Ź"));
+ } else {
+ progressParts.push(this.dim("â—‹"));
+ }
+ }
+ lines.push(padToWidth(boxLine(progressParts.join(" "))));
+ lines.push(padToWidth(emptyBoxLine()));
+
+ // Current question
+ const q = this.questions[this.currentIndex];
+ const questionText = `${this.bold("Q:")} ${q.question}`;
+ const wrappedQuestion = wrapTextWithAnsi(questionText, contentWidth);
+ for (const line of wrappedQuestion) {
+ lines.push(padToWidth(boxLine(line)));
+ }
+
+ // Context if present
+ if (q.context) {
+ lines.push(padToWidth(emptyBoxLine()));
+ const contextText = this.gray(`> ${q.context}`);
+ const wrappedContext = wrapTextWithAnsi(contextText, contentWidth - 2);
+ for (const line of wrappedContext) {
+ lines.push(padToWidth(boxLine(line)));
+ }
+ }
+
+ lines.push(padToWidth(emptyBoxLine()));
+
+ // Render the editor component (multi-line input) with padding
+ // Skip the first and last lines (editor's own border lines)
+ const answerPrefix = this.bold("A: ");
+ const editorWidth = contentWidth - 4 - 3; // Extra padding + space for "A: "
+ const editorLines = this.editor.render(editorWidth);
+ for (let i = 1; i < editorLines.length - 1; i++) {
+ if (i === 1) {
+ // First content line gets the "A: " prefix
+ lines.push(padToWidth(boxLine(answerPrefix + editorLines[i])));
+ } else {
+ // Subsequent lines get padding to align with the first line
+ lines.push(padToWidth(boxLine(" " + editorLines[i])));
+ }
+ }
+
+ // Notes section
+ if (this.notesMode) {
+ lines.push(padToWidth(emptyBoxLine()));
+ const notesLabel = `${this.cyan("✎")} ${this.bold("Note:")} ${this.dim("(Enter to save, Esc to discard)")}`;
+ lines.push(padToWidth(boxLine(notesLabel)));
+ const notesEditorWidth = contentWidth - 4;
+ const notesEditorLines = this.notesEditor.render(notesEditorWidth);
+ for (let i = 1; i < notesEditorLines.length - 1; i++) {
+ lines.push(padToWidth(boxLine(" " + notesEditorLines[i])));
+ }
+ } else if (this.notesText) {
+ lines.push(padToWidth(emptyBoxLine()));
+ const savedNote = `${this.cyan("✎")} ${this.gray(this.notesText)}`;
+ const wrappedNote = wrapTextWithAnsi(savedNote, contentWidth);
+ for (const line of wrappedNote) {
+ lines.push(padToWidth(boxLine(line)));
+ }
+ lines.push(padToWidth(boxLine(this.dim("Alt+N to edit note"))));
+ } else {
+ lines.push(padToWidth(emptyBoxLine()));
+ lines.push(padToWidth(boxLine(this.dim("Alt+N to add a note"))));
+ }
+
+ lines.push(padToWidth(emptyBoxLine()));
+
+ // Confirmation dialog or footer with controls
+ if (this.showingConfirmation) {
+ lines.push(padToWidth(this.dim("├" + horizontalLine(boxWidth - 2) + "┤")));
+ const confirmMsg = `${this.yellow("Submit all answers?")} ${this.dim("(Enter/y to confirm, Esc/n to cancel)")}`;
+ lines.push(padToWidth(boxLine(truncateToWidth(confirmMsg, contentWidth))));
+ } else {
+ lines.push(padToWidth(this.dim("├" + horizontalLine(boxWidth - 2) + "┤")));
+ 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`;
+ lines.push(padToWidth(boxLine(truncateToWidth(controls, contentWidth))));
+ }
+ lines.push(padToWidth(this.dim("â•°" + horizontalLine(boxWidth - 2) + "â•Ż")));
+
+ this.cachedWidth = width;
+ this.cachedLines = lines;
+ return lines;
+ }
+}
+
+export default function (pi: ExtensionAPI) {
+ const answerHandler = async (ctx: ExtensionContext, instruction?: string) => {
+ if (!ctx.hasUI) {
+ ctx.ui.notify("answer requires interactive mode", "error");
+ return;
+ }
+
+ if (!ctx.model) {
+ ctx.ui.notify("No model selected", "error");
+ return;
+ }
+
+ // Find the last assistant message on the current branch
+ const branch = ctx.sessionManager.getBranch();
+ let lastAssistantText: string | undefined;
+
+ for (let i = branch.length - 1; i >= 0; i--) {
+ const entry = branch[i];
+ if (entry.type === "message") {
+ const msg = entry.message;
+ if ("role" in msg && msg.role === "assistant") {
+ if (msg.stopReason !== "stop") {
+ ctx.ui.notify(`Last assistant message incomplete (${msg.stopReason})`, "error");
+ return;
+ }
+ const textParts = msg.content
+ .filter((c): c is { type: "text"; text: string } => c.type === "text")
+ .map((c) => c.text);
+ if (textParts.length > 0) {
+ lastAssistantText = textParts.join("\n");
+ break;
+ }
+ }
+ }
+ }
+
+ if (!lastAssistantText) {
+ ctx.ui.notify("No assistant messages found", "error");
+ return;
+ }
+
+ // Run extraction with loader UI.
+ // The result distinguishes user cancellation from errors so we can
+ // show the right message after the custom UI closes.
+ type ExtractionOutcome =
+ | { kind: "ok"; result: ExtractionResult }
+ | { kind: "cancelled" }
+ | { kind: "error"; message: string };
+
+ const outcome = await ctx.ui.custom<ExtractionOutcome>((tui, theme, _kb, done) => {
+ const extractionModel = resolveExtractionModel(ctx);
+ const loader = new BorderedLoader(tui, theme, `Extracting questions using ${extractionModel.name}...`);
+ loader.onAbort = () => done({ kind: "cancelled" });
+
+ const tryExtract = async (model: ReturnType<typeof resolveExtractionModel>) => {
+ const auth = await ctx.modelRegistry.getApiKeyAndHeaders(model);
+ if (!auth.ok) {
+ return { kind: "model_error" as const, model };
+ }
+ const instructionBlock = instruction
+ ? `<user_instruction>\n${instruction}\n</user_instruction>\n\nExtract questions from the assistant's message based on the user's instruction above.`
+ : "Extract questions from the assistant's message for the user to fill out.";
+
+ const userMessage: UserMessage = {
+ role: "user",
+ content: [
+ {
+ type: "text",
+ text: `<last_assistant_message>\n${lastAssistantText!}\n</last_assistant_message>\n\n${instructionBlock}`,
+ },
+ ],
+ timestamp: Date.now(),
+ };
+
+ const response = await complete(
+ model,
+ {
+ systemPrompt: SYSTEM_PROMPT,
+ messages: [userMessage],
+ tools: [QUESTION_EXTRACTION_TOOL],
+ },
+ { apiKey: auth.apiKey, headers: auth.headers, signal: loader.signal },
+ );
+
+ if (response.stopReason === "aborted") {
+ return { kind: "cancelled" as const };
+ }
+
+ if (response.stopReason === "error" || response.content.length === 0) {
+ return { kind: "model_error" as const, model };
+ }
+
+ const parsed = parseToolCallResult(response);
+ if (!parsed) {
+ return {
+ kind: "parse_error" as const,
+ message: `${model.name} did not call extract_questions tool`,
+ };
+ }
+ return { kind: "ok" as const, result: parsed };
+ };
+
+ const doExtract = async () => {
+ let result = await tryExtract(extractionModel);
+
+ // If the preferred model errored and it's not already the
+ // session model, fall back to the session model.
+ if (result.kind === "model_error" && extractionModel.id !== ctx.model!.id) {
+ ctx.ui.notify(`${extractionModel.name} unavailable, falling back to ${ctx.model!.name}...`, "warning");
+ result = await tryExtract(ctx.model!);
+ }
+
+ switch (result.kind) {
+ case "ok":
+ return done({ kind: "ok", result: result.result });
+ case "cancelled":
+ return done({ kind: "cancelled" });
+ case "model_error":
+ return done({
+ kind: "error",
+ message: `${result.model.name} returned an error with no content`,
+ });
+ case "parse_error":
+ return done({ kind: "error", message: result.message });
+ }
+ };
+
+ doExtract().catch((err) =>
+ done({
+ kind: "error",
+ message: `${err?.message ?? err}`,
+ }),
+ );
+
+ return loader;
+ });
+
+ if (outcome.kind === "cancelled") {
+ ctx.ui.notify("Cancelled", "info");
+ return;
+ }
+ if (outcome.kind === "error") {
+ ctx.ui.notify(`Extraction failed: ${outcome.message}`, "error");
+ return;
+ }
+
+ const extractionResult = outcome.result;
+
+ if (extractionResult.questions.length === 0) {
+ ctx.ui.notify("No questions found in the last message", "info");
+ return;
+ }
+
+ // Show the Q&A component
+ const answersResult = await ctx.ui.custom<string | null>((tui, _theme, _kb, done) => {
+ return new QnAComponent(extractionResult.questions, tui, done);
+ });
+
+ if (answersResult === null) {
+ ctx.ui.notify("Cancelled", "info");
+ return;
+ }
+
+ // Send the answers directly as a message and trigger a turn
+ pi.sendMessage(
+ {
+ customType: "answers",
+ content: "I answered your questions in the following way:\n\n" + answersResult,
+ display: true,
+ },
+ { triggerTurn: true },
+ );
+ };
+
+ pi.registerCommand("answer", {
+ description:
+ "Extract questions from last assistant message into interactive Q&A. Optional: provide instructions for what to extract.",
+ handler: (args, ctx) => answerHandler(ctx, args?.trim() || undefined),
+ });
+
+ pi.registerShortcut("ctrl+.", {
+ description: "Extract and answer questions",
+ handler: answerHandler,
+ });
+}