@pi9/ask 0.1.0 → 0.3.0

package/src/index.ts

@@ -1,20 +1,25 @@
 import type { ExtensionAPI, SessionEntry, Theme } from "@earendil-works/pi-coding-agent";
-import { Text } from "@earendil-works/pi-tui";
+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,16 +28,11 @@ 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" : "dim";
+  const questionColor = state.status === "answered" ? "text" : "muted";
   const title = `${theme.fg("toolTitle", "ask")} ${theme.fg(questionColor, args.question)}`;
   if (state.status) return title;
 
@@ -41,7 +41,7 @@ function renderAskCall(args: AskParams, theme: Theme, state: AskRendererState):
   const timeout = args.timeout !== undefined && args.timeout > 0
     ? ` · timeout:${formatTimeout(args.timeout)}`
     : "";
-  return `${title}\n${theme.fg("dim", `╰ ${mode}options:${optionCount}${timeout}`)}`;
+  return `${title}\n${theme.fg("muted", `╰ ${mode}options:${optionCount}${timeout}`)}`;
 }
 
 function formatTimeout(timeoutMs: number): string {
@@ -50,26 +50,55 @@ function formatTimeout(timeoutMs: number): string {
   return `${Number.isInteger(seconds) ? seconds : Number(seconds.toFixed(2))}s`;
 }
 
-function renderAnsweredOptions(args: AskParams, answer: AskAnswer, theme: Theme): string {
-  const selections = new Map(answer.selections.map(selection => [selection.label, selection]));
-  const checkedGlyph = args.allowMultiple === true ? CHECKED_BOX : CHECKED_CIRCLE;
-  const emptyGlyph = args.allowMultiple === true ? EMPTY_BOX : EMPTY_CIRCLE;
-  const lines = args.options.map(option => {
-    const label = option.label.trim();
-    const selection = selections.get(label);
-    const comment = selection?.comment ? ` (${selection.comment})` : "";
-    const text = `${label}${comment}`;
-    return selection
-      ? `${theme.fg("success", checkedGlyph)} ${theme.fg("text", text)}`
-      : theme.fg("dim", `${emptyGlyph} ${text}`);
-  });
-  if (answer.freeform) lines.push(`${theme.fg("success", checkedGlyph)} ${theme.fg("text", answer.freeform)}`);
-  return lines.map((line, index) => `${index === 0 ? `${theme.fg("dim", "╰")} ` : "  "}${line}`).join("\n");
+class AnsweredOptions implements Component {
+  constructor(
+    private readonly args: AskParams,
+    private readonly answer: AskAnswer,
+    private readonly theme: Theme,
+  ) {}
+
+  invalidate(): void {}
+
+  render(width: number): string[] {
+    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, index) => {
+      const label = option.label.trim();
+      const selection = selections.get(index);
+      const comment = selection?.comment ? ` (${selection.comment})` : "";
+      return {
+        marker: this.theme.fg(selection ? "success" : "muted", selection ? checkedGlyph : emptyGlyph),
+        text: this.theme.fg(selection ? "text" : "muted", `${label}${comment}`),
+      };
+    });
+    if (this.answer.freeform) {
+      rows.push({
+        marker: this.theme.fg("success", checkedGlyph),
+        text: this.theme.fg("text", this.answer.freeform),
+      });
+    }
+
+    return rows.flatMap((row, index) => wrapAnsweredRow(
+      `${index === 0 ? `${this.theme.fg("muted", "╰")} ` : "  "}${row.marker} `,
+      row.text,
+      width,
+    ));
+  }
+}
+
+function wrapAnsweredRow(prefix: string, text: string, width: number): string[] {
+  const safeWidth = Math.max(1, Number.isFinite(width) ? Math.floor(width) : 1);
+  const prefixWidth = visibleWidth(prefix);
+  if (prefixWidth >= safeWidth) return wrapTextWithAnsi(`${prefix}${text}`, safeWidth);
+
+  const wrapped = wrapTextWithAnsi(text, safeWidth - prefixWidth);
+  const continuation = " ".repeat(prefixWidth);
+  return wrapped.map((line, index) => `${index === 0 ? prefix : continuation}${line}`);
 }
 
 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>();
 
@@ -100,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();
@@ -137,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" };
     }
   });
 
@@ -177,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 {
@@ -220,7 +222,7 @@ export default function askExtension(pi: ExtensionAPI) {
         ?? (result.details?.status === "answered" ? result.details.answer : undefined);
       context.state.status = answer === undefined ? "settled" : "answered";
       context.state.callComponent?.setText(renderAskCall(context.args, theme, context.state));
-      if (answer) return new Text(renderAnsweredOptions(context.args, answer, theme), 0, 0);
+      if (answer) return new AnsweredOptions(context.args, answer, theme);
       const text = result.content.find((item) => item.type === "text")?.text ?? "Ask completed.";
       const color = result.details?.status === "cancelled" || result.details?.status === "unanswered"
         ? "muted"

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 };
 }

package/src/session.ts

@@ -0,0 +1,254 @@
+import type {
+  ContextEvent,
+  SessionEntry,
+  SessionTreeEvent,
+} from "@earendil-works/pi-coding-agent";
+import { Check } from "typebox/value";
+
+import {
+  AskAnsweredDetailsSchema,
+  AskParamsSchema,
+  answerMatchesAsk,
+  normalizeAsk,
+  parseAskReplayDetails,
+  type Ask,
+  type AskAnswer,
+  type AskParams,
+  type AskReplayDetails,
+} from "./domain.js";
+
+export const ASK_REPLAY_CUSTOM_TYPE = "ask:reanswer" as const;
+const ASK_SUMMARY_CUSTOM_TYPE = "ask:summary" as const;
+
+type AgentMessage = ContextEvent["messages"][number];
+type AssistantMessage = Extract<AgentMessage, { role: "assistant" }>;
+type ToolResultMessage = Extract<AgentMessage, { role: "toolResult" }>;
+type CustomMessage = Extract<AgentMessage, { role: "custom" }>;
+
+type AskCall = {
+  ask: Ask;
+  message: AssistantMessage;
+};
+
+type ReplayRecord = {
+  message: CustomMessage;
+  details?: AskReplayDetails;
+};
+
+type AskSummaryPayload = {
+  type: "ask_response";
+  question: string;
+  context?: string;
+  selectionMode: "single" | "multi";
+  answer: {
+    selections: Array<{ label: string; description?: string; comment?: string }>;
+    freeform?: string;
+  };
+};
+
+export type AskReplayMessage = {
+  customType: typeof ASK_REPLAY_CUSTOM_TYPE;
+  content: "";
+  display: false;
+  details: AskReplayDetails;
+};
+
+export type AskReplayResolution =
+  | { status: "resolved"; toolCallId: string; ask: Ask }
+  | {
+      status: "not-replayable";
+      reason: "no-entry" | "not-assistant" | "not-ask" | "multiple-tool-calls" | "mixed-tools" | "invalid-arguments";
+    };
+
+export function buildAskReplayMessage(toolCallId: string, answer: AskAnswer): AskReplayMessage {
+  return {
+    customType: ASK_REPLAY_CUSTOM_TYPE,
+    content: "",
+    display: false,
+    details: { toolCallId, answer },
+  };
+}
+
+/** Resolve only the selected leaf, or its immediate parent when Pi created a branch summary. */
+export function resolveAskReplayTarget(
+  event: Pick<SessionTreeEvent, "newLeafId" | "summaryEntry">,
+  getEntry: (id: string) => SessionEntry | undefined | null,
+): AskReplayResolution {
+  if (!event.newLeafId) return rejected("no-entry");
+
+  let entry = getEntry(event.newLeafId);
+  const summary = event.summaryEntry?.id === event.newLeafId
+    ? event.summaryEntry
+    : entry?.type === "branch_summary" ? entry : undefined;
+  if (summary) entry = summary.parentId ? getEntry(summary.parentId) : undefined;
+  if (!entry) return rejected("no-entry");
+
+  let selectedToolCallId: string | undefined;
+  if (entry.type === "message" && entry.message.role === "toolResult" && entry.message.toolName === "ask") {
+    selectedToolCallId = entry.message.toolCallId;
+    entry = entry.parentId ? getEntry(entry.parentId) : undefined;
+  }
+  if (!entry) return rejected("no-entry");
+  if (entry.type !== "message" || entry.message.role !== "assistant") return rejected("not-assistant");
+
+  const calls = entry.message.content.filter(item => item.type === "toolCall");
+  const askCalls = calls.filter(call => call.name === "ask");
+  if (calls.length !== 1) {
+    if (askCalls.length > 0 && askCalls.length < calls.length) return rejected("mixed-tools");
+    if (askCalls.length > 1) return rejected("multiple-tool-calls");
+    return rejected("not-ask");
+  }
+
+  const call = calls[0]!;
+  if (call.name !== "ask" || (selectedToolCallId !== undefined && call.id !== selectedToolCallId)) {
+    return rejected("not-ask");
+  }
+
+  const ask = parseStoredAsk(call.arguments);
+  if (!ask) return rejected("invalid-arguments");
+  return { status: "resolved", toolCallId: call.id, ask };
+}
+
+export function parseStoredAsk(value: unknown): Ask | undefined {
+  if (!Check(AskParamsSchema, value)) return undefined;
+  try {
+    return normalizeAsk(value as AskParams);
+  } catch {
+    return undefined;
+  }
+}
+
+/** Replace completed standalone Ask exchanges with their decision-focused answer. */
+export function rewriteAskContext(messages: readonly AgentMessage[]): AgentMessage[] {
+  const calls = collectCalls(messages);
+  const nativeResults = collectNativeResults(messages);
+  const replayRecords = collectReplayRecords(messages);
+
+  const summaries = new Map<AgentMessage, AgentMessage>();
+  const removals = new Set<AgentMessage>();
+  for (const [toolCallId, call] of calls) {
+    if (!call) continue;
+
+    const results = nativeResults.get(toolCallId) ?? [];
+    if (results.length > 1) continue;
+    const result = results[0];
+
+    const replays = replayRecords.get(toolCallId) ?? [];
+    let answer: AskAnswer | undefined;
+    let replay: ReplayRecord | undefined;
+    if (replays.length > 0) {
+      if (replays.length !== 1) continue;
+      const candidate = replays[0]!;
+      if (!candidate.details || !answerMatchesAsk(candidate.details.answer, call.ask)) continue;
+      replay = candidate;
+      answer = candidate.details.answer;
+    }
+
+    if (result) {
+      if (result.toolName !== "ask") continue;
+      if (!replay) {
+        if (result.isError) continue;
+        const details = parseNativeDetails(result.details);
+        if (!details || !answerMatchesAsk(details.answer, call.ask)) continue;
+        answer = details.answer;
+      }
+    } else if (!replay) {
+      continue;
+    }
+
+    if (!answer) continue;
+    const timestamp = replay ? replay.message.timestamp : result!.timestamp;
+    summaries.set(call.message, makeSummary(call, answer, timestamp));
+    removals.add(call.message);
+    if (result) removals.add(result);
+    if (replay) removals.add(replay.message);
+  }
+
+  return messages.flatMap((message) => {
+    const summary = summaries.get(message);
+    if (summary) return [summary];
+    return removals.has(message) ? [] : [message];
+  });
+}
+
+function collectCalls(messages: readonly AgentMessage[]): Map<string, AskCall | undefined> {
+  const calls = new Map<string, AskCall | undefined>();
+  for (const message of messages) {
+    if (message.role !== "assistant") continue;
+    const toolCalls = message.content.filter(block => block.type === "toolCall");
+    for (const block of toolCalls) {
+      if (block.name !== "ask") continue;
+      const ask = toolCalls.length === 1 ? parseStoredAsk(block.arguments) : undefined;
+      const call = ask ? { ask, message } : undefined;
+      calls.set(block.id, calls.has(block.id) ? undefined : call);
+    }
+  }
+  return calls;
+}
+
+function collectNativeResults(messages: readonly AgentMessage[]): Map<string, ToolResultMessage[]> {
+  const results = new Map<string, ToolResultMessage[]>();
+  for (const message of messages) {
+    if (message.role !== "toolResult") continue;
+    const matching = results.get(message.toolCallId) ?? [];
+    matching.push(message);
+    results.set(message.toolCallId, matching);
+  }
+  return results;
+}
+
+function collectReplayRecords(messages: readonly AgentMessage[]): Map<string, ReplayRecord[]> {
+  const records = new Map<string, ReplayRecord[]>();
+  for (const message of messages) {
+    if (message.role !== "custom" || message.customType !== ASK_REPLAY_CUSTOM_TYPE) continue;
+    const rawToolCallId = isRecord(message.details) && typeof message.details.toolCallId === "string"
+      ? message.details.toolCallId
+      : undefined;
+    const details = parseAskReplayDetails(message.details);
+    const toolCallId = details?.toolCallId ?? rawToolCallId;
+    if (toolCallId === undefined) continue;
+    const matching = records.get(toolCallId) ?? [];
+    matching.push({ message, ...(details ? { details } : {}) });
+    records.set(toolCallId, matching);
+  }
+  return records;
+}
+
+function parseNativeDetails(value: unknown): { status: "answered"; answer: AskAnswer } | undefined {
+  return Check(AskAnsweredDetailsSchema, value) ? value : undefined;
+}
+
+function makeSummary(call: AskCall, answer: AskAnswer, timestamp: number): AgentMessage {
+  const payload: AskSummaryPayload = {
+    type: "ask_response",
+    question: call.ask.question,
+    ...(call.ask.context !== undefined ? { context: call.ask.context } : {}),
+    selectionMode: call.ask.allowMultiple ? "multi" : "single",
+    answer: {
+      selections: answer.selections.map((selection) => {
+        const option = call.ask.options[selection.option]!;
+        return {
+          label: option.label,
+          ...(option.description !== undefined ? { description: option.description } : {}),
+          ...(selection.comment !== undefined ? { comment: selection.comment } : {}),
+        };
+      }),
+      ...(answer.freeform !== undefined ? { freeform: answer.freeform } : {}),
+    },
+  };
+  return {
+    role: "custom",
+    customType: ASK_SUMMARY_CUSTOM_TYPE,
+    display: false,
+    content: JSON.stringify(payload),
+    timestamp,
+  };
+}
+
+function rejected(reason: Exclude<AskReplayResolution, { status: "resolved" }>["reason"]): AskReplayResolution {
+  return { status: "not-replayable", reason };
+}
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+  return typeof value === "object" && value !== null && !Array.isArray(value);
+}