@pi9/ask 0.1.0

package/src/context.ts

@@ -0,0 +1,209 @@
+import type { ContextEvent } from "@earendil-works/pi-coding-agent";
+import { Check } from "typebox/value";
+
+import { AskAnsweredDetailsSchema } from "./schema.js";
+import {
+  ASK_REPLAY_CUSTOM_TYPE,
+  parseAskReplayDetails,
+  validateStoredArgs,
+} from "./replay.js";
+import type { AskAnswer, AskReplayDetails, ValidatedAskParams } from "./types.js";
+
+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 = {
+  args: ValidatedAskParams;
+  message: AssistantMessage;
+  standalone: boolean;
+};
+type ReplayRecord = {
+  message: CustomMessage;
+  details?: AskReplayDetails;
+};
+
+const ASK_SUMMARY_CUSTOM_TYPE = "ask:summary" as const;
+
+type AskSummaryPayload = {
+  type: "ask_response";
+  question: string;
+  context?: string;
+  selectionMode: "single" | "multi";
+  answer: AskAnswer;
+};
+
+/**
+ * Replace completed Ask exchanges with the small custom message that is useful
+ * to the model. The session still contains the original protocol messages; the
+ * projection must not leave a tool call without its result.
+ */
+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 || !call.standalone) 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) {
+      // A replay is a revision only when it is the one well-formed marker for
+      // this call. In particular, do not hide malformed or duplicate markers.
+      if (replays.length !== 1) continue;
+      const candidate = replays[0]!;
+      if (!candidate.details
+        || !matchesCall(candidate.details, call.args)
+        || !answerMatchesCall(candidate.details.answer, call.args)) continue;
+      replay = candidate;
+      answer = candidate.details.answer;
+    }
+
+    if (result) {
+      if (result.toolName !== "ask") continue;
+      // A valid replay is the latest answer and takes precedence regardless
+      // of whether the earlier native attempt completed successfully.
+      if (!replay) {
+        if (result.isError) continue;
+        const details = parseNativeDetails(result.details);
+        if (!details
+          || details.question !== call.args.question
+          || !answerMatchesCall(details.answer, call.args)) continue;
+        answer = details.answer;
+      }
+    } else if (!replay) {
+      continue;
+    }
+
+    if (!answer) continue;
+    const timestamp = replay ? replay.message.timestamp : result!.timestamp;
+    const summary = makeSummary(call, answer, timestamp);
+    summaries.set(call.message, summary);
+    removals.add(call.message);
+    if (result) removals.add(result);
+    if (replay) removals.add(replay.message);
+  }
+
+  const rewritten: AgentMessage[] = [];
+  for (const message of messages) {
+    const summary = summaries.get(message);
+    if (summary) {
+      rewritten.push(summary);
+      continue;
+    }
+    if (!removals.has(message)) rewritten.push(message);
+  }
+  return rewritten;
+}
+
+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 args = validateStoredArgs(block.arguments);
+      const call = args
+        ? { args, message, standalone: toolCalls.length === 1 }
+        : undefined;
+      // Keep an invalid entry as ambiguous too. A later valid call with the
+      // same ID must not make an otherwise malformed exchange transformable.
+      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 rawDetails = isRecord(message.details) && typeof message.details.toolCallId === "string"
+      ? message.details.toolCallId
+      : undefined;
+    const details = parseAskReplayDetails(message.details);
+    const toolCallId = details?.toolCallId ?? rawDetails;
+    if (toolCallId === undefined) continue;
+    const matching = records.get(toolCallId) ?? [];
+    matching.push({ message, ...(details ? { details } : {}) });
+    records.set(toolCallId, matching);
+  }
+  return records;
+}
+
+function parseNativeDetails(value: unknown): { question: string; answer: AskAnswer } | undefined {
+  return Check(AskAnsweredDetailsSchema, value) ? value : undefined;
+}
+
+function matchesCall(replay: AskReplayDetails, args: ValidatedAskParams): boolean {
+  return replay.question === args.question
+    && replay.context === args.context
+    && replay.allowMultiple === args.allowMultiple;
+}
+
+function answerMatchesCall(answer: AskAnswer, args: ValidatedAskParams): boolean {
+  if (answer.selections.length > 1 && !args.allowMultiple) return false;
+  if (answer.freeform !== undefined && !args.allowFreeform) return false;
+
+  const labels = new Set(args.options.map(option => option.label));
+  const selected = new Set<string>();
+  for (const selection of answer.selections) {
+    if (!labels.has(selection.label) || selected.has(selection.label)) return false;
+    selected.add(selection.label);
+  }
+  return true;
+}
+
+function makeSummary(call: AskCall, answer: AskAnswer, timestamp: number): AgentMessage {
+  const payload: AskSummaryPayload = {
+    type: "ask_response",
+    question: call.args.question,
+    ...(call.args.context !== undefined ? { context: call.args.context } : {}),
+    selectionMode: call.args.allowMultiple ? "multi" : "single",
+    answer: {
+      selections: answer.selections.map(selection => selectedOption(selection, call.args)),
+      ...(answer.freeform !== undefined ? { freeform: answer.freeform } : {}),
+    },
+  };
+  return {
+    role: "custom",
+    customType: ASK_SUMMARY_CUSTOM_TYPE,
+    display: false,
+    content: JSON.stringify(payload),
+    timestamp,
+  };
+}
+
+function selectedOption(selection: AskAnswer["selections"][number], args: ValidatedAskParams): AskAnswer["selections"][number] {
+  const original = args.options.find(option => option.label === selection.label);
+  const description = selection.description ?? original?.description;
+  return {
+    label: selection.label,
+    ...(description !== undefined ? { description } : {}),
+    ...(selection.comment !== undefined ? { comment: selection.comment } : {}),
+  };
+}
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+  return typeof value === "object" && value !== null && !Array.isArray(value);
+}

package/src/deadline.ts

@@ -0,0 +1,78 @@
+export interface DeadlineSignal {
+  signal: AbortSignal | undefined;
+  readonly timedOut: boolean;
+  dispose(): void;
+}
+
+export function createDeadlineSignal(
+  parent: AbortSignal | undefined,
+  timeoutMs: number | undefined,
+): DeadlineSignal {
+  const hasTimeout = timeoutMs !== undefined
+    && Number.isFinite(timeoutMs)
+    && timeoutMs > 0;
+
+  if (parent === undefined && !hasTimeout) {
+    return {
+      signal: undefined,
+      timedOut: false,
+      dispose() {},
+    };
+  }
+
+  const controller = new AbortController();
+  let disposed = false;
+  let timedOut = false;
+  let timer: ReturnType<typeof setTimeout> | undefined;
+  let parentListener: (() => void) | undefined;
+
+  const cleanup = (): void => {
+    if (timer !== undefined) {
+      clearTimeout(timer);
+      timer = undefined;
+    }
+    if (parent !== undefined && parentListener !== undefined) {
+      parent.removeEventListener("abort", parentListener);
+      parentListener = undefined;
+    }
+  };
+
+  const abort = (reason?: unknown): void => {
+    if (disposed || controller.signal.aborted) return;
+    cleanup();
+    controller.abort(reason);
+  };
+
+  if (parent?.aborted) {
+    abort(parent.reason);
+    return {
+      signal: controller.signal,
+      timedOut: false,
+      dispose() {},
+    };
+  }
+
+  if (parent !== undefined) {
+    parentListener = () => abort(parent.reason);
+    parent.addEventListener("abort", parentListener, { once: true });
+  }
+
+  if (hasTimeout) {
+    timer = setTimeout(() => {
+      timedOut = true;
+      abort();
+    }, timeoutMs);
+  }
+
+  return {
+    signal: controller.signal,
+    get timedOut() {
+      return timedOut;
+    },
+    dispose() {
+      if (disposed) return;
+      disposed = true;
+      cleanup();
+    },
+  };
+}

package/src/glyphs.ts

@@ -0,0 +1,4 @@
+export const EMPTY_CIRCLE = "󰄰";
+export const CHECKED_CIRCLE = "󰄴";
+export const EMPTY_BOX = "󰄱";
+export const CHECKED_BOX = "󰄵";

package/src/index.ts

@@ -0,0 +1,231 @@
+import type { ExtensionAPI, SessionEntry, Theme } from "@earendil-works/pi-coding-agent";
+import { Text } from "@earendil-works/pi-tui";
+
+import { rewriteAskContext } from "./context.js";
+import { resolveTimeoutMs } from "./config.js";
+import { createDeadlineSignal } from "./deadline.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";
+
+type AskRendererState = {
+  callComponent?: Text;
+  status?: "answered" | "settled";
+  invalidate?: () => void;
+};
+
+type AskReplayState =
+  | { status: "idle" }
+  | { status: "prompting" }
+  | {
+      status: "dispatched";
+      details: ReturnType<typeof buildAskReplayMessage>["details"];
+      clearEditor: boolean;
+    };
+
+function renderAskCall(args: AskParams, theme: Theme, state: AskRendererState): string {
+  const questionColor = state.status === "answered" ? "text" : "dim";
+  const title = `${theme.fg("toolTitle", "ask")} ${theme.fg(questionColor, args.question)}`;
+  if (state.status) return title;
+
+  const optionCount = args.options.length;
+  const mode = args.allowMultiple === true ? "multi · " : "";
+  const timeout = args.timeout !== undefined && args.timeout > 0
+    ? ` · timeout:${formatTimeout(args.timeout)}`
+    : "";
+  return `${title}\n${theme.fg("dim", `╰ ${mode}options:${optionCount}${timeout}`)}`;
+}
+
+function formatTimeout(timeoutMs: number): string {
+  if (timeoutMs < 1000) return `${timeoutMs}ms`;
+  const seconds = timeoutMs / 1000;
+  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");
+}
+
+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>();
+
+  const applyRevision = (toolCallId: string, answer: AskAnswer) => {
+    revisedAnswers.set(toolCallId, answer);
+    rendererStates.get(toolCallId)?.invalidate?.();
+  };
+  const restoreRevisions = (entries: readonly SessionEntry[]) => {
+    revisedAnswers.clear();
+    for (const entry of entries) {
+      if (entry.type !== "custom_message" || entry.customType !== ASK_REPLAY_CUSTOM_TYPE) continue;
+      const details = parseAskReplayDetails(entry.details);
+      if (details) revisedAnswers.set(details.toolCallId, details.answer);
+    }
+    for (const state of rendererStates.values()) state.invalidate?.();
+  };
+
+  const reconcileAskTool = (hasUI: boolean) => {
+    if (hasUI) return;
+    const activeTools = pi.getActiveTools();
+    if (!activeTools.includes("ask")) return;
+    pi.setActiveTools(activeTools.filter(name => name !== "ask"));
+  };
+
+  pi.on("session_start", (_event, ctx) => {
+    reconcileAskTool(ctx.hasUI);
+    restoreRevisions(ctx.sessionManager.getBranch());
+  });
+  pi.on("before_agent_start", (_event, ctx) => reconcileAskTool(ctx.hasUI));
+  pi.on("context", (event) => ({ messages: rewriteAskContext(event.messages) }));
+  pi.on("agent_settled", (_event, ctx) => {
+    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();
+    restoreRevisions(entries);
+    const byId = new Map(entries.map(entry => [entry.id, entry]));
+    if (event.summaryEntry) byId.set(event.summaryEntry.id, event.summaryEntry);
+    const resolution = resolveAskReplayTarget(event, id => byId.get(id));
+    if (resolution.status !== "resolved") {
+      if (resolution.reason === "mixed-tools" || resolution.reason === "multiple-tool-calls" || resolution.reason === "invalid-arguments") {
+        ctx.ui.notify("This Ask cannot be re-answered because its original tool call is mixed or invalid.", "warning");
+      }
+      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),
+    );
+    try {
+      const answer = await launchQuestionnaire(ctx, resolution.params, deadline.signal);
+      if (!answer) return;
+      const message = buildAskReplayMessage(resolution.toolCallId, resolution.params, answer);
+      pi.sendMessage(message, { triggerTurn: true, deliverAs: "followUp" });
+      applyRevision(message.details.toolCallId, message.details.answer);
+      replayState = { status: "dispatched", details: message.details, clearEditor: guardEditor };
+    } finally {
+      deadline.dispose();
+      if (replayState.status !== "dispatched") {
+        replayState = { status: "idle" };
+        if (guardEditor) setTimeout(() => ctx.ui.setEditorText(""), 0);
+      }
+    }
+  });
+
+  pi.registerTool<typeof AskParamsSchema, AskToolDetails, AskRendererState>({
+    name: "ask",
+    label: "Ask",
+    description: "Ask the user one focused question with selectable options. Blocks until answered, cancelled, or timed out.",
+    promptSnippet: "Ask the user a focused question with selectable options when input is required",
+    promptGuidelines: [
+      "Use ask only when you can offer a short list of useful options; put open-ended questions in your normal response instead.",
+      "An ask_response is a completed ask whose tool call was removed from the context; treat its answer as final and do not re-ask.",
+    ],
+    parameters: AskParamsSchema,
+    executionMode: "sequential",
+
+    async execute(_toolCallId, rawParams, signal, _onUpdate, ctx) {
+      const params = validateAskParams(rawParams as AskParams);
+      if (!ctx.hasUI) return buildUiUnavailableResponse(params.question);
+
+      const deadline = createDeadlineSignal(
+        signal,
+        resolveTimeoutMs(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);
+            pi.events.emit("ask:unanswered", result.details);
+            return result;
+          }
+          const result = buildCancelledResponse(params.question);
+          pi.events.emit("ask:cancelled", result.details);
+          return result;
+        }
+
+        const result = buildAnsweredResponse(params.question, answer);
+        pi.events.emit("ask:answered", result.details);
+        return result;
+      } finally {
+        deadline.dispose();
+      }
+    },
+
+    renderCall(args, theme, context) {
+      const text = context.lastComponent instanceof Text ? context.lastComponent : new Text("", 0, 0);
+      context.state.callComponent = text;
+      text.setText(renderAskCall(args, theme, context.state));
+      return text;
+    },
+    renderResult(result, _options, theme, context) {
+      context.state.invalidate = context.invalidate;
+      rendererStates.set(context.toolCallId, context.state);
+      const answer = revisedAnswers.get(context.toolCallId)
+        ?? (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);
+      const text = result.content.find((item) => item.type === "text")?.text ?? "Ask completed.";
+      const color = result.details?.status === "cancelled" || result.details?.status === "unanswered"
+        ? "muted"
+        : "text";
+      return new Text(theme.fg(color, text), 0, 0);
+    },
+  });
+}

package/src/preview.ts

@@ -0,0 +1,55 @@
+import { getMarkdownTheme } from "@earendil-works/pi-coding-agent";
+import { Markdown, truncateToWidth, visibleWidth } from "@earendil-works/pi-tui";
+
+/** Terminal width at which an authored preview gets its own pane. */
+export const PREVIEW_WIDE_THRESHOLD = 88;
+
+export interface PreviewPaneLayout {
+  leftWidth: number;
+  rightWidth: number;
+}
+
+/**
+ * Choose the option/preview pane sizes without allowing either pane to become
+ * unusably narrow. The split is intentionally opt-in at a comfortable width.
+ */
+export function getPreviewPaneLayout(width: number): PreviewPaneLayout | undefined {
+  const renderWidth = safeWidth(width);
+  if (renderWidth < PREVIEW_WIDE_THRESHOLD) return undefined;
+
+  const available = renderWidth - 1; // one column for the separator
+  const leftWidth = Math.max(32, Math.floor(available * 0.4));
+  const rightWidth = available - leftWidth;
+  if (rightWidth < 40) return undefined;
+  return { leftWidth, rightWidth };
+}
+
+/** Render authored option text using Pi's normal markdown theme. */
+export function renderPreviewMarkdown(preview: string, width: number): string[] {
+  if (!preview.trim()) return [];
+  const lines = new Markdown(preview, 0, 0, getMarkdownTheme()).render(safeWidth(width));
+  return lines.map(line => fitAndPad(line, safeWidth(width)));
+}
+
+/** Compose one visible option row with its corresponding preview row. */
+export function composePreviewRow(
+  left: string,
+  preview: string,
+  layout: PreviewPaneLayout,
+  separator: string,
+): string {
+  return [
+    fitAndPad(left, layout.leftWidth),
+    fitAndPad(separator, 1),
+    fitAndPad(preview, layout.rightWidth),
+  ].join("");
+}
+
+function safeWidth(width: number): number {
+  return Math.max(1, Number.isFinite(width) ? Math.floor(width) : 1);
+}
+
+function fitAndPad(line: string, width: number): string {
+  const clipped = visibleWidth(line) > width ? truncateToWidth(line, width, "") : line;
+  return clipped + " ".repeat(Math.max(0, width - visibleWidth(clipped)));
+}

package/src/questionnaire.ts

@@ -0,0 +1,37 @@
+import type { ExtensionUIContext } from "@earendil-works/pi-coding-agent";
+
+import { AskComponent } from "./component.js";
+import type { AskAnswer, ValidatedAskParams } from "./types.js";
+
+interface QuestionnaireLaunchContext {
+  ui: Pick<ExtensionUIContext, "custom">;
+}
+
+export async function launchQuestionnaire(
+  ctx: QuestionnaireLaunchContext,
+  params: ValidatedAskParams,
+  signal?: AbortSignal,
+): Promise<AskAnswer | null> {
+  let abortListener: (() => void) | undefined;
+  try {
+    const answer = await ctx.ui.custom<AskAnswer | null>((tui, theme, keybindings, done) => {
+      const component = new AskComponent({
+        ...params,
+        tui,
+        theme,
+        keybindings,
+        onSubmit: done,
+        onCancel: () => done(null),
+      });
+
+      abortListener = () => component.cancel();
+      if (signal?.aborted) abortListener();
+      else if (signal) signal.addEventListener("abort", abortListener, { once: true });
+
+      return component;
+    });
+    return answer ?? null;
+  } finally {
+    if (abortListener) signal?.removeEventListener("abort", abortListener);
+  }
+}