@pi9/ask 0.2.0 → 0.3.0

package/README.md

@@ -5,6 +5,7 @@ Ask adds interactive questions to [Pi](https://github.com/earendil-works/pi-mono
 ## Feature overview
 
 - **Answer revisions** — Reopen an Ask entry from `/tree` and continue with a new answer.
+- **Tool-call rewriting** — Replace answered standalone Ask exchanges in model context with compact, decision-focused responses while preserving the original session history.
 - **Lean tool definition** — A highly optimized description and schema minimize prompt overhead while staying compatible across providers.
 - **Rich previews** — Inspect Markdown previews before choosing.
 - **Multi-select** — Choose any number of options, then submit them together.
@@ -61,6 +62,16 @@ Comments let you qualify a selection without writing a separate message. When a
 
 ## Advanced behavior
 
+### Tool-call rewriting
+
+Ask rewrites answered standalone exchanges before they are sent back to the model, replacing the Ask call and its answer record with a compact `ask_response` message. The replacement keeps the question, optional background, single- or multi-select mode, selected option labels and descriptions, comments, and any freeform response. It leaves out unselected options and all Markdown previews.
+
+The main purpose is to keep proposed alternatives from being mistaken for decisions. Once you choose an answer, the model sees only what you selected and the context attached to that selection—not every option it originally offered. As a secondary benefit, dropping unused answers, previews, and the surrounding tool-call protocol reduces the amount of context carried into later turns.
+
+The answer can come from either a successful native tool result or a valid revision submitted from `/tree`. A revision becomes the authoritative answer, replacing the earlier result in model context even if that result is absent or unsuccessful.
+
+This rewrite affects only the model-facing context. The original messages remain in the stored session, so the conversation still renders normally and the question can still be reopened from `/tree`. Ask rewrites only unambiguous calls and answers that validate against each other; otherwise, it leaves the messages untouched rather than risk producing a misleading response or an unmatched tool call.
+
 ### Responsive previews
 
 Previews appear beside the options in terminals at least 88 columns wide and stack below them in narrower layouts. They support Markdown and update as you move between options. While you edit a comment or freeform response, the preview is hidden to leave more room for writing.

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@pi9/ask",
-  "version": "0.2.0",
+  "version": "0.3.0",
   "description": "Preview-rich, replayable Pi questions with multi-select and comments.",
   "license": "MIT",
   "author": "Chase Cummings <chaseecummings@gmail.com>",

package/src/component.ts

@@ -33,11 +33,11 @@ import {
   type FocusRange,
   type ViewportOverflow,
 } from "./viewport.js";
-import type { AskAnswer, ValidatedAskParams } from "./types.js";
+import type { Ask, AskAnswer } from "./domain.js";
 
 const FRAME_WIDE_PREVIEW = true; // Set to false to compare the same layout without its outer frame.
 
-type AskComponentOptions = ValidatedAskParams & {
+type AskComponentOptions = Ask & {
   tui: TUI;
   theme: Theme;
   keybindings: KeybindingsManager;
@@ -377,17 +377,17 @@ export class AskComponent implements Component, Focusable {
 
       if (row.kind === "option") {
         const source = row.option;
-        const checked = this.questionnaireState.checked.has(source.label);
+        const checked = this.questionnaireState.checked.has(row.index);
         const badge = this.questionnaireState.config.allowMultiple && checked
           ? this.config.theme.fg("success", " [selected]")
           : "";
-        const comment = this.questionnaireState.comments.has(source.label)
+        const comment = this.questionnaireState.comments.has(row.index)
           ? this.config.theme.fg("warning", " ✎")
           : "";
         addPrefixed(marker, `${source.label}${badge}${comment}`, selected ? "accent" : "text");
         if (source.description) addPrefixed("  ", source.description, "muted");
         if (selected) {
-          const commentText = this.questionnaireState.comments.get(source.label);
+          const commentText = this.questionnaireState.comments.get(row.index);
           if (commentText) addPrefixed("     ", `✎ ${commentText}`, "dim");
         }
       } else {

package/src/deadline.ts

@@ -1,16 +1,41 @@
+export const MAX_TIMEOUT_MS = 2_147_483_647;
+
+type AskEnvironment = Readonly<{
+  PI9_ASK_TIMEOUT_MS?: string;
+}>;
+
 export interface DeadlineSignal {
   signal: AbortSignal | undefined;
   readonly timedOut: boolean;
   dispose(): void;
 }
 
+export function resolveTimeoutMs(
+  perCallTimeout: number | undefined,
+  env: AskEnvironment,
+): number | undefined {
+  if (perCallTimeout !== undefined) {
+    return Number.isInteger(perCallTimeout) && perCallTimeout > 0 && perCallTimeout <= MAX_TIMEOUT_MS
+      ? perCallTimeout
+      : undefined;
+  }
+
+  const envTimeout = env.PI9_ASK_TIMEOUT_MS;
+  if (envTimeout === undefined || !/^\d+$/.test(envTimeout)) return undefined;
+
+  const timeout = Number(envTimeout);
+  return Number.isInteger(timeout) && timeout > 0 && timeout <= MAX_TIMEOUT_MS
+    ? timeout
+    : undefined;
+}
+
 export function createDeadlineSignal(
   parent: AbortSignal | undefined,
-  timeoutMs: number | undefined,
+  perCallTimeout: number | undefined,
+  env: AskEnvironment = {},
 ): DeadlineSignal {
-  const hasTimeout = timeoutMs !== undefined
-    && Number.isFinite(timeoutMs)
-    && timeoutMs > 0;
+  const timeoutMs = resolveTimeoutMs(perCallTimeout, env);
+  const hasTimeout = timeoutMs !== undefined;
 
   if (parent === undefined && !hasTimeout) {
     return {

package/src/domain.ts

@@ -0,0 +1,144 @@
+import { Type, type Static } from "typebox";
+import { Check } from "typebox/value";
+
+import { MAX_TIMEOUT_MS } from "./deadline.js";
+
+export const AskOptionSchema = Type.Object({
+  label: Type.String({ minLength: 1, description: "Short option title; keep distinct across options." }),
+  description: Type.Optional(Type.String({ description: "Brief explanation when the label alone is ambiguous." })),
+  preview: Type.Optional(Type.String({ description: "Markdown shown to the user, e.g. code, plans, or comparisons; not returned in the answer." })),
+}, { additionalProperties: false });
+
+export const AskParamsSchema = Type.Object({
+  question: Type.String({ minLength: 1 }),
+  context: Type.Optional(Type.String({ description: "Optional background shown above the question, e.g. what prompted the ask." })),
+  options: Type.Array(AskOptionSchema, { minItems: 1, description: "Suggested answers." }),
+  allowMultiple: Type.Optional(Type.Boolean({ default: false, description: "Allow selecting multiple options. Defaults to false." })),
+  allowFreeform: Type.Optional(Type.Boolean({ default: true, description: "Allow a typed response. Defaults to true." })),
+  timeout: Type.Optional(Type.Integer({ minimum: 0, maximum: MAX_TIMEOUT_MS, description: "Timeout in ms for answers that would go stale; the call returns unanswered on expiry. Zero disables any default." })),
+}, { additionalProperties: false });
+
+export const AskSelectionSchema = Type.Object({
+  option: Type.Integer({ minimum: 0 }),
+  comment: Type.Optional(Type.String({ minLength: 1 })),
+}, { additionalProperties: false });
+
+export const AskAnswerSchema = Type.Object({
+  selections: Type.Array(AskSelectionSchema),
+  freeform: Type.Optional(Type.String({ minLength: 1 })),
+}, { additionalProperties: false });
+
+export const AskReplayDetailsSchema = Type.Object({
+  toolCallId: Type.String({ minLength: 1 }),
+  answer: AskAnswerSchema,
+}, { additionalProperties: false });
+
+export const AskAnsweredDetailsSchema = Type.Object({
+  status: Type.Literal("answered"),
+  answer: AskAnswerSchema,
+}, { additionalProperties: false });
+
+export type AskOption = Static<typeof AskOptionSchema>;
+export type AskParams = Static<typeof AskParamsSchema>;
+export type AskSelection = Static<typeof AskSelectionSchema>;
+export type AskAnswer = Static<typeof AskAnswerSchema>;
+export type AskReplayDetails = Static<typeof AskReplayDetailsSchema>;
+export type AskAnsweredDetails = Static<typeof AskAnsweredDetailsSchema>;
+
+export type Ask = Omit<AskParams, "allowMultiple" | "allowFreeform"> & {
+  allowMultiple: boolean;
+  allowFreeform: boolean;
+};
+
+export type AskToolDetails =
+  | AskAnsweredDetails
+  | { status: "unanswered" | "cancelled" | "ui_unavailable" };
+
+export type AskResponse = {
+  content: Array<{ type: "text"; text: string }>;
+  details: AskToolDetails;
+};
+
+export function normalizeAsk(params: AskParams): Ask {
+  const question = params.question.trim();
+  if (!question) throw new Error("Ask question must not be empty.");
+
+  const context = trimOptional(params.context);
+  const options = params.options.map(normalizeOption);
+  if (options.length === 0) throw new Error("Ask needs at least one option.");
+
+  const labels = new Set<string>();
+  for (const option of options) {
+    if (labels.has(option.label)) throw new Error(`Ask options contain duplicate label: ${option.label}.`);
+    labels.add(option.label);
+  }
+
+  return {
+    question,
+    ...(context === undefined ? {} : { context }),
+    options,
+    allowMultiple: params.allowMultiple ?? false,
+    allowFreeform: params.allowFreeform ?? true,
+    ...(params.timeout === undefined ? {} : { timeout: params.timeout }),
+  };
+}
+
+export function parseAskReplayDetails(value: unknown): AskReplayDetails | undefined {
+  return Check(AskReplayDetailsSchema, value) ? value : undefined;
+}
+
+export function answerMatchesAsk(answer: AskAnswer, ask: Ask): boolean {
+  if (answer.selections.length > 1 && !ask.allowMultiple) return false;
+  if (answer.freeform !== undefined && !ask.allowFreeform) return false;
+
+  const selected = new Set<number>();
+  for (const selection of answer.selections) {
+    if (selection.option >= ask.options.length || selected.has(selection.option)) return false;
+    selected.add(selection.option);
+  }
+  return true;
+}
+
+export function formatAskAnswer(ask: Ask, answer: AskAnswer): string {
+  const lines = answer.selections.map((selection) => {
+    const option = ask.options[selection.option]!;
+    const description = option.description ? ` — ${option.description}` : "";
+    const comment = selection.comment ? ` (${selection.comment})` : "";
+    return `Selected: ${option.label}${description}${comment}`;
+  });
+  if (answer.freeform) lines.push(`Freeform: ${answer.freeform}`);
+  return lines.join("\n") || "No answer provided.";
+}
+
+export function buildAskResponse(ask: Ask, details: AskToolDetails): AskResponse {
+  if (details.status === "answered") {
+    return {
+      content: [{ type: "text", text: formatAskAnswer(ask, details.answer) }],
+      details,
+    };
+  }
+
+  const text = {
+    unanswered: "The question timed out without an answer.",
+    cancelled: "User cancelled the question.",
+    ui_unavailable: "Interactive UI is unavailable.",
+  }[details.status];
+  return { content: [{ type: "text", text }], details };
+}
+
+function normalizeOption(option: AskOption): AskOption {
+  const label = option.label.trim();
+  if (!label) throw new Error("Ask option label must not be empty.");
+  const description = trimOptional(option.description);
+  const preview = option.preview?.trim() ? option.preview : undefined;
+  return {
+    label,
+    ...(description === undefined ? {} : { description }),
+    ...(preview === undefined ? {} : { preview }),
+  };
+}
+
+function trimOptional(value: string | undefined): string | undefined {
+  const trimmed = value?.trim();
+  return trimmed || undefined;
+}

package/src/index.ts

@@ -1,20 +1,25 @@
 import type { ExtensionAPI, SessionEntry, Theme } from "@earendil-works/pi-coding-agent";
 import { Text, visibleWidth, wrapTextWithAnsi, type Component } from "@earendil-works/pi-tui";
 
-import { rewriteAskContext } from "./context.js";
-import { resolveTimeoutMs } from "./config.js";
 import { createDeadlineSignal } from "./deadline.js";
+import {
+  AskParamsSchema,
+  buildAskResponse,
+  normalizeAsk,
+  parseAskReplayDetails,
+  type AskAnswer,
+  type AskParams,
+  type AskToolDetails,
+} from "./domain.js";
 import { CHECKED_BOX, CHECKED_CIRCLE, EMPTY_BOX, EMPTY_CIRCLE } from "./glyphs.js";
 import { launchQuestionnaire } from "./questionnaire.js";
-import { renderAskReanswerMessage } from "./replay-renderer.js";
-import { ASK_REPLAY_CUSTOM_TYPE, buildAskReplayMessage, parseAskReplayDetails, resolveAskReplayTarget } from "./replay.js";
-import { buildAnsweredResponse, buildCancelledResponse, buildUiUnavailableResponse, buildUnansweredResponse } from "./response.js";
 import { askWithRpc } from "./rpc.js";
-import { AskParamsSchema } from "./schema.js";
-import type { AskAnswer, AskParams, AskToolDetails } from "./types.js";
-import { validateAskParams } from "./validation.js";
-
-const TREE_EDITOR_GUARD = "\u200B";
+import {
+  ASK_REPLAY_CUSTOM_TYPE,
+  buildAskReplayMessage,
+  resolveAskReplayTarget,
+  rewriteAskContext,
+} from "./session.js";
 
 type AskRendererState = {
   callComponent?: Text;
@@ -23,13 +28,8 @@ type AskRendererState = {
 };
 
 type AskReplayState =
-  | { status: "idle" }
-  | { status: "prompting" }
-  | {
-      status: "dispatched";
-      details: ReturnType<typeof buildAskReplayMessage>["details"];
-      clearEditor: boolean;
-    };
+  | { status: "idle" | "prompting" }
+  | { status: "dispatched"; details: ReturnType<typeof buildAskReplayMessage>["details"] };
 
 function renderAskCall(args: AskParams, theme: Theme, state: AskRendererState): string {
   const questionColor = state.status === "answered" ? "text" : "muted";
@@ -60,12 +60,12 @@ class AnsweredOptions implements Component {
   invalidate(): void {}
 
   render(width: number): string[] {
-    const selections = new Map(this.answer.selections.map(selection => [selection.label, selection]));
+    const selections = new Map(this.answer.selections.map(selection => [selection.option, selection]));
     const checkedGlyph = this.args.allowMultiple === true ? CHECKED_BOX : CHECKED_CIRCLE;
     const emptyGlyph = this.args.allowMultiple === true ? EMPTY_BOX : EMPTY_CIRCLE;
-    const rows = this.args.options.map(option => {
+    const rows = this.args.options.map((option, index) => {
       const label = option.label.trim();
-      const selection = selections.get(label);
+      const selection = selections.get(index);
       const comment = selection?.comment ? ` (${selection.comment})` : "";
       return {
         marker: this.theme.fg(selection ? "success" : "muted", selection ? checkedGlyph : emptyGlyph),
@@ -99,7 +99,6 @@ function wrapAnsweredRow(prefix: string, text: string, width: number): string[]
 
 export default function askExtension(pi: ExtensionAPI) {
   let replayState: AskReplayState = { status: "idle" };
-  let replayTreeSelection = false;
   const revisedAnswers = new Map<string, AskAnswer>();
   const rendererStates = new Map<string, AskRendererState>();
 
@@ -130,29 +129,17 @@ export default function askExtension(pi: ExtensionAPI) {
   });
   pi.on("before_agent_start", (_event, ctx) => reconcileAskTool(ctx.hasUI));
   pi.on("context", (event) => ({ messages: rewriteAskContext(event.messages) }));
-  pi.on("agent_settled", (_event, ctx) => {
+  pi.on("agent_settled", () => {
     if (replayState.status !== "dispatched") return;
-    if (replayState.clearEditor) {
-      replayState = { ...replayState, clearEditor: false };
-      setTimeout(() => ctx.ui.setEditorText(""), 0);
-    }
     pi.events.emit("ask:reanswered", replayState.details);
     replayState = { status: "idle" };
   });
   pi.on("session_shutdown", () => {
     replayState = { status: "idle" };
-    replayTreeSelection = false;
     revisedAnswers.clear();
     rendererStates.clear();
   });
-  pi.registerMessageRenderer(ASK_REPLAY_CUSTOM_TYPE, renderAskReanswerMessage);
-  pi.on("session_before_tree", (event, ctx) => {
-    const target = ctx.sessionManager.getEntry(event.preparation.targetId);
-    replayTreeSelection = target?.type === "custom_message" && target.customType === ASK_REPLAY_CUSTOM_TYPE;
-  });
   pi.on("session_tree", async (event, ctx) => {
-    const suppressEditorRestore = replayTreeSelection;
-    replayTreeSelection = false;
     if (ctx.mode !== "tui" || replayState.status !== "idle") return;
 
     const entries = ctx.sessionManager.getBranch();
@@ -167,30 +154,18 @@ export default function askExtension(pi: ExtensionAPI) {
       return;
     }
 
-    // Pi restores selected custom-message content after this hook returns. Keep
-    // the editor temporarily non-empty until the replay turn settles so the
-    // selected marker text cannot replace the guard.
-    const guardEditor = suppressEditorRestore && !ctx.ui.getEditorText().trim();
-    if (guardEditor) ctx.ui.setEditorText(TREE_EDITOR_GUARD);
-
     replayState = { status: "prompting" };
-    const deadline = createDeadlineSignal(
-      undefined,
-      resolveTimeoutMs(resolution.params.timeout, process.env),
-    );
+    const deadline = createDeadlineSignal(undefined, resolution.ask.timeout, process.env);
     try {
-      const answer = await launchQuestionnaire(ctx, resolution.params, deadline.signal);
+      const answer = await launchQuestionnaire(ctx, resolution.ask, deadline.signal);
       if (!answer) return;
-      const message = buildAskReplayMessage(resolution.toolCallId, resolution.params, answer);
+      const message = buildAskReplayMessage(resolution.toolCallId, answer);
       pi.sendMessage(message, { triggerTurn: true, deliverAs: "followUp" });
       applyRevision(message.details.toolCallId, message.details.answer);
-      replayState = { status: "dispatched", details: message.details, clearEditor: guardEditor };
+      replayState = { status: "dispatched", details: message.details };
     } finally {
       deadline.dispose();
-      if (replayState.status !== "dispatched") {
-        replayState = { status: "idle" };
-        if (guardEditor) setTimeout(() => ctx.ui.setEditorText(""), 0);
-      }
+      if (replayState.status !== "dispatched") replayState = { status: "idle" };
     }
   });
 
@@ -207,29 +182,26 @@ export default function askExtension(pi: ExtensionAPI) {
     executionMode: "sequential",
 
     async execute(_toolCallId, rawParams, signal, _onUpdate, ctx) {
-      const params = validateAskParams(rawParams as AskParams);
-      if (!ctx.hasUI) return buildUiUnavailableResponse(params.question);
+      const params = normalizeAsk(rawParams as AskParams);
+      if (!ctx.hasUI) return buildAskResponse(params, { status: "ui_unavailable" });
 
-      const deadline = createDeadlineSignal(
-        signal,
-        resolveTimeoutMs(params.timeout, process.env),
-      );
+      const deadline = createDeadlineSignal(signal, params.timeout, process.env);
       try {
         const answer = ctx.mode === "tui"
           ? await launchQuestionnaire(ctx, params, deadline.signal)
           : await askWithRpc(ctx.ui, params, deadline.signal);
         if (answer === null) {
           if (deadline.timedOut) {
-            const result = buildUnansweredResponse(params.question);
+            const result = buildAskResponse(params, { status: "unanswered" });
             pi.events.emit("ask:unanswered", result.details);
             return result;
           }
-          const result = buildCancelledResponse(params.question);
+          const result = buildAskResponse(params, { status: "cancelled" });
           pi.events.emit("ask:cancelled", result.details);
           return result;
         }
 
-        const result = buildAnsweredResponse(params.question, answer);
+        const result = buildAskResponse(params, { status: "answered", answer });
         pi.events.emit("ask:answered", result.details);
         return result;
       } finally {

package/src/questionnaire.ts

@@ -1,7 +1,7 @@
 import type { ExtensionUIContext } from "@earendil-works/pi-coding-agent";
 
 import { AskComponent } from "./component.js";
-import type { AskAnswer, ValidatedAskParams } from "./types.js";
+import type { Ask, AskAnswer } from "./domain.js";
 
 interface QuestionnaireLaunchContext {
   ui: Pick<ExtensionUIContext, "custom">;
@@ -9,7 +9,7 @@ interface QuestionnaireLaunchContext {
 
 export async function launchQuestionnaire(
   ctx: QuestionnaireLaunchContext,
-  params: ValidatedAskParams,
+  params: Ask,
   signal?: AbortSignal,
 ): Promise<AskAnswer | null> {
   let abortListener: (() => void) | undefined;

package/src/rpc.ts

@@ -1,101 +1,93 @@
-import type { AskAnswer, AskOption, ValidatedAskParams } from "./types.js";
+import type { Ask, AskAnswer, AskOption, AskSelection } from "./domain.js";
 
 export interface AskDialogUI {
   select(title: string, options: string[], dialogOptions?: { signal?: AbortSignal }): Promise<string | undefined>;
   input(title: string, placeholder?: string, dialogOptions?: { signal?: AbortSignal }): Promise<string | undefined>;
 }
 
-type AskSelection = AskAnswer["selections"][number];
-
 const FREEFORM_CHOICE = "Type a response…";
 
 export async function askWithRpc(
   ui: AskDialogUI,
-  params: ValidatedAskParams,
+  ask: Ask,
   signal?: AbortSignal,
 ): Promise<AskAnswer | null> {
-  const prompt = params.context ? `${params.context}\n\n${params.question}` : params.question;
+  const prompt = ask.context ? `${ask.context}\n\n${ask.question}` : ask.question;
   if (signal?.aborted) return null;
 
-  return params.allowMultiple
-    ? runMultiSelect(ui, prompt, params.options, params.allowFreeform, signal)
-    : runSingleSelect(ui, prompt, params.options, params.allowFreeform, signal);
+  return ask.allowMultiple
+    ? runMultiSelect(ui, prompt, ask, signal)
+    : runSingleSelect(ui, prompt, ask, signal);
 }
 
 async function runSingleSelect(
   ui: AskDialogUI,
   prompt: string,
-  options: AskOption[],
-  allowFreeform: boolean,
+  ask: Ask,
   signal?: AbortSignal,
 ): Promise<AskAnswer | null> {
-  const choices = options.map((option, index) => `${index + 1}. ${formatOption(option)}`);
-  if (allowFreeform) choices.push(`${options.length + 1}. ${FREEFORM_CHOICE}`);
+  const choices = ask.options.map((option, index) => `${index + 1}. ${formatOption(option)}`);
+  if (ask.allowFreeform) choices.push(`${ask.options.length + 1}. ${FREEFORM_CHOICE}`);
 
   const selected = await selectDialog(ui, prompt, choices, signal);
   if (selected == null || signal?.aborted) return null;
 
   const selectedIndex = choices.indexOf(selected);
   if (selectedIndex === -1) return null;
-  if (selectedIndex === options.length) {
+  if (selectedIndex === ask.options.length) {
     const freeform = await readInput(ui, prompt, signal);
     return freeform === null ? null : makeAnswer([], freeform);
   }
 
-  const selectedOption = options[selectedIndex];
-  if (!selectedOption) return null;
-  const selections = await collectComments(ui, prompt, [selectedOption], signal);
+  const selections = await collectComments(ui, prompt, ask.options, [selectedIndex], signal);
   return selections === null ? null : makeAnswer(selections);
 }
 
 async function runMultiSelect(
   ui: AskDialogUI,
   prompt: string,
-  options: AskOption[],
-  allowFreeform: boolean,
+  ask: Ask,
   signal?: AbortSignal,
 ): Promise<AskAnswer | null> {
-  const numberedOptions = options
+  const numberedOptions = ask.options
     .map((option, index) => `${index + 1}. ${formatOption(option)}`)
     .join("\n");
   const selectionPrompt = `${prompt}\n\n${numberedOptions}\n\nEnter option numbers separated by commas:`;
   const rawSelection = await readInput(ui, selectionPrompt, signal);
   if (rawSelection === null) return null;
 
-  const selectedOptions = parseSelections(rawSelection, options);
-  const freeform = allowFreeform
+  const selected = parseSelections(rawSelection, ask.options.length);
+  const freeform = ask.allowFreeform
     ? await readInput(ui, `${prompt}\n\nAdditional response (optional):`, signal)
     : undefined;
   if (freeform === null) return null;
 
-  const selections = await collectComments(ui, prompt, selectedOptions, signal);
+  const selections = await collectComments(ui, prompt, ask.options, selected, signal);
   return selections === null ? null : makeAnswer(selections, freeform);
 }
 
-function parseSelections(raw: string, options: AskOption[]): AskOption[] {
-  const indices = new Set<number>();
+function parseSelections(raw: string, optionCount: number): number[] {
+  const selected = new Set<number>();
   for (const token of raw.split(",")) {
     const value = token.trim();
     if (!/^\d+$/.test(value)) continue;
 
     const index = Number(value) - 1;
-    if (Number.isSafeInteger(index) && index >= 0 && index < options.length) indices.add(index);
+    if (Number.isSafeInteger(index) && index >= 0 && index < optionCount) selected.add(index);
   }
-  return options.filter((_option, index) => indices.has(index));
+  return [...selected].sort((a, b) => a - b);
 }
 
 async function collectComments(
   ui: AskDialogUI,
   prompt: string,
   options: AskOption[],
+  selected: number[],
   signal?: AbortSignal,
 ): Promise<AskSelection[] | null> {
-  const selections: AskSelection[] = options.map(toSelection);
-  for (let index = 0; index < options.length; index++) {
-    const option = options[index];
-    const selection = selections[index];
-    if (!option || !selection) continue;
-
+  const selections: AskSelection[] = selected.map(option => ({ option }));
+  for (const selection of selections) {
+    const option = options[selection.option]!;
     const comment = await readInput(ui, `${prompt}\n\nComment for \"${option.label}\" (optional):`, signal);
     if (comment === null) return null;
     if (comment) selection.comment = comment;
@@ -110,12 +102,6 @@ async function readInput(ui: AskDialogUI, title: string, signal?: AbortSignal):
   return value.trim();
 }
 
-function toSelection(option: AskOption): AskSelection {
-  const selection: AskSelection = { label: option.label };
-  if (option.description !== undefined) selection.description = option.description;
-  return selection;
-}
-
 function makeAnswer(selections: AskSelection[], freeform?: string): AskAnswer {
   return freeform ? { selections, freeform } : { selections };
 }