@oh-my-pi/pi-coding-agent 13.5.1 → 13.5.3

package/src/tools/ask.ts

@@ -14,9 +14,11 @@
  *   - Use recommended: <index> to mark the default option; "(Recommended)" suffix is added automatically
  *   - Questions may time out and auto-select the recommended option (configurable, disabled in plan mode)
  */
+
 import type { AgentTool, AgentToolContext, AgentToolResult, AgentToolUpdateCallback } from "@oh-my-pi/pi-agent-core";
 import type { Component } from "@oh-my-pi/pi-tui";
 import { TERMINAL, Text } from "@oh-my-pi/pi-tui";
+import { ptree, untilAborted } from "@oh-my-pi/pi-utils";
 import { type Static, Type } from "@sinclair/typebox";
 import { renderPromptTemplate } from "../config/prompt-templates";
 import type { RenderResultOptions } from "../extensibility/custom-tools/types";
@@ -25,6 +27,7 @@ import askDescription from "../prompts/tools/ask.md" with { type: "text" };
 import { renderStatusLine } from "../tui";
 import type { ToolSession } from ".";
 import { formatErrorMessage, formatMeta, formatTitle } from "./render-utils";
+import { ToolAbortError } from "./tool-errors";
 
 // =============================================================================
 // Types
@@ -110,14 +113,9 @@ interface UIContext {
 	select(
 		prompt: string,
 		options: string[],
-		options_?: { initialIndex?: number; timeout?: number; outline?: boolean },
+		options_?: { initialIndex?: number; signal?: AbortSignal; outline?: boolean },
 	): Promise<string | undefined>;
-	input(prompt: string): Promise<string | undefined>;
-}
-
-interface AskQuestionOptions {
-	/** Timeout in milliseconds, null/undefined to disable */
-	timeout?: number | null;
+	input(prompt: string, options_?: { signal?: AbortSignal }): Promise<string | undefined>;
 }
 
 async function askSingleQuestion(
@@ -126,9 +124,8 @@ async function askSingleQuestion(
 	optionLabels: string[],
 	multi: boolean,
 	recommended?: number,
-	options?: AskQuestionOptions,
+	signal?: AbortSignal,
 ): Promise<SelectionResult> {
-	const timeout = options?.timeout ?? undefined;
 	const doneLabel = getDoneOptionLabel();
 	let selectedOptions: string[] = [];
 	let customInput: string | undefined;
@@ -152,22 +149,27 @@ async function askSingleQuestion(
 			opts.push(OTHER_OPTION);
 
 			const prefix = selected.size > 0 ? `(${selected.size} selected) ` : "";
-			const selectionStart = Date.now();
-			const choice = await ui.select(`${prefix}${question}`, opts, {
-				initialIndex: cursorIndex,
-				timeout: timeout ?? undefined,
-				outline: true,
-			});
-			const elapsed = Date.now() - selectionStart;
-			const timedOut = timeout != null && elapsed >= timeout;
+			const choice = signal
+				? await untilAborted(signal, () =>
+						ui.select(`${prefix}${question}`, opts, {
+							initialIndex: cursorIndex,
+							signal,
+							outline: true,
+						}),
+					)
+				: await ui.select(`${prefix}${question}`, opts, {
+						initialIndex: cursorIndex,
+						signal,
+						outline: true,
+					});
 
 			if (choice === undefined || choice === doneLabel) break;
 
 			if (choice === OTHER_OPTION) {
-				if (!timedOut) {
-					const input = await ui.input("Enter your response:");
-					if (input) customInput = input;
-				}
+				const input = signal
+					? await untilAborted(signal, () => ui.input("Enter your response:", { signal }))
+					: await ui.input("Enter your response:", { signal });
+				if (input) customInput = input;
 				break;
 			}
 
@@ -192,21 +194,28 @@ async function askSingleQuestion(
 					selected.add(opt);
 				}
 			}
-
-			if (timedOut) {
-				break;
-			}
 		}
 		selectedOptions = Array.from(selected);
 	} else {
 		const displayLabels = addRecommendedSuffix(optionLabels, recommended);
-		const choice = await ui.select(question, [...displayLabels, OTHER_OPTION], {
-			timeout: timeout ?? undefined,
-			initialIndex: recommended,
-			outline: true,
-		});
+		const choice = signal
+			? await untilAborted(signal, () =>
+					ui.select(question, [...displayLabels, OTHER_OPTION], {
+						initialIndex: recommended,
+						signal,
+						outline: true,
+					}),
+				)
+			: await ui.select(question, [...displayLabels, OTHER_OPTION], {
+					initialIndex: recommended,
+					signal,
+					outline: true,
+				});
+
 		if (choice === OTHER_OPTION) {
-			const input = await ui.input("Enter your response:");
+			const input = signal
+				? await untilAborted(signal, () => ui.input("Enter your response:", { signal }))
+				: await ui.input("Enter your response:", { signal });
 			if (input) customInput = input;
 		} else if (choice) {
 			selectedOptions = [stripRecommendedSuffix(choice)];
@@ -265,7 +274,7 @@ export class AskTool implements AgentTool<typeof askSchema, AskToolDetails> {
 	async execute(
 		_toolCallId: string,
 		params: AskParams,
-		_signal?: AbortSignal,
+		signal?: AbortSignal,
 		_onUpdate?: AgentToolUpdateCallback<AskToolDetails>,
 		context?: AgentToolContext,
 	): Promise<AgentToolResult<AskToolDetails>> {
@@ -277,7 +286,11 @@ export class AskTool implements AgentTool<typeof askSchema, AskToolDetails> {
 			};
 		}
 
-		const { ui } = context;
+		const extensionUi = context.ui;
+		const ui: UIContext = {
+			select: (prompt, options, dialogOptions) => extensionUi.select(prompt, options, dialogOptions),
+			input: (prompt, dialogOptions) => extensionUi.input(prompt, undefined, dialogOptions),
+		};
 
 		// Determine timeout based on settings and plan mode
 		const planModeEnabled = this.session.getPlanModeState?.()?.enabled ?? false;
@@ -296,18 +309,41 @@ export class AskTool implements AgentTool<typeof askSchema, AskToolDetails> {
 			};
 		}
 
+		const askQuestion = async (q: AskParams["questions"][number]) => {
+			const optionLabels = q.options.map(o => o.label);
+			const timeoutSignal = timeout == null ? undefined : AbortSignal.timeout(timeout);
+			const questionSignal = ptree.combineSignals(signal, timeoutSignal);
+			try {
+				const { selectedOptions, customInput } = await askSingleQuestion(
+					ui,
+					q.question,
+					optionLabels,
+					q.multi ?? false,
+					q.recommended,
+					questionSignal,
+				);
+				return { optionLabels, selectedOptions, customInput, timedOut: false };
+			} catch (error) {
+				if (error instanceof Error && error.name === "AbortError") {
+					if (signal?.aborted) {
+						throw new ToolAbortError("Ask input was cancelled");
+					}
+					if (timeoutSignal?.aborted) {
+						return { optionLabels, selectedOptions: [], customInput: undefined, timedOut: true };
+					}
+				}
+				throw error;
+			}
+		};
+
 		if (params.questions.length === 1) {
 			const [q] = params.questions;
-			const optionLabels = q.options.map(o => o.label);
-			const { selectedOptions, customInput } = await askSingleQuestion(
-				ui,
-				q.question,
-				optionLabels,
-				q.multi ?? false,
-				q.recommended,
-				{ timeout },
-			);
+			const { optionLabels, selectedOptions, customInput, timedOut } = await askQuestion(q);
 
+			if (!timedOut && selectedOptions.length === 0 && !customInput) {
+				context.abort();
+				throw new ToolAbortError("Ask tool was cancelled by the user");
+			}
 			const details: AskToolDetails = {
 				question: q.question,
 				options: optionLabels,
@@ -333,16 +369,12 @@ export class AskTool implements AgentTool<typeof askSchema, AskToolDetails> {
 		const results: QuestionResult[] = [];
 
 		for (const q of params.questions) {
-			const optionLabels = q.options.map(o => o.label);
-			const { selectedOptions, customInput } = await askSingleQuestion(
-				ui,
-				q.question,
-				optionLabels,
-				q.multi ?? false,
-				q.recommended,
-				{ timeout },
-			);
+			const { optionLabels, selectedOptions, customInput, timedOut } = await askQuestion(q);
 
+			if (!timedOut && selectedOptions.length === 0 && !customInput) {
+				context.abort();
+				throw new ToolAbortError("Ask tool was cancelled by the user");
+			}
 			results.push({
 				id: q.id,
 				question: q.question,

package/src/tools/bash.ts

@@ -98,7 +98,11 @@ export class BashTool implements AgentTool<BashToolSchema, BashToolDetails> {
 	constructor(private readonly session: ToolSession) {
 		this.#asyncEnabled = this.session.settings.get("async.enabled");
 		this.parameters = this.#asyncEnabled ? bashSchemaWithAsync : bashSchemaBase;
-		this.description = renderPromptTemplate(bashDescription, { asyncEnabled: this.#asyncEnabled });
+		this.description = renderPromptTemplate(bashDescription, {
+			asyncEnabled: this.#asyncEnabled,
+			hasAstGrep: this.session.settings.get("astGrep.enabled"),
+			hasAstEdit: this.session.settings.get("astEdit.enabled"),
+		});
 	}
 
 	#formatResultOutput(result: BashResult | BashInteractiveResult, headLines?: number, tailLines?: number): string {

package/src/tools/checkpoint.ts

@@ -0,0 +1,128 @@
+import type { AgentTool, AgentToolContext, AgentToolResult, AgentToolUpdateCallback } from "@oh-my-pi/pi-agent-core";
+import { type Static, Type } from "@sinclair/typebox";
+import { renderPromptTemplate } from "../config/prompt-templates";
+import checkpointDescription from "../prompts/tools/checkpoint.md" with { type: "text" };
+import rewindDescription from "../prompts/tools/rewind.md" with { type: "text" };
+import type { ToolSession } from ".";
+import type { OutputMeta } from "./output-meta";
+import { ToolError } from "./tool-errors";
+import { toolResult } from "./tool-result";
+
+export interface CheckpointState {
+	/** Number of in-memory messages at checkpoint (AFTER checkpoint tool result is appended) */
+	checkpointMessageCount: number;
+	/** Session entry ID at checkpoint (for session tree branching) */
+	checkpointEntryId: string | null;
+	/** Timestamp */
+	startedAt: string;
+}
+
+const checkpointSchema = Type.Object({
+	goal: Type.String({ description: "What you are investigating and why" }),
+});
+
+type CheckpointParams = Static<typeof checkpointSchema>;
+
+const rewindSchema = Type.Object({
+	report: Type.String({ description: "Concise investigation findings to retain after rewind" }),
+});
+
+type RewindParams = Static<typeof rewindSchema>;
+
+export interface CheckpointToolDetails {
+	goal: string;
+	startedAt: string;
+	meta?: OutputMeta;
+}
+
+export interface RewindToolDetails {
+	report: string;
+	rewound: boolean;
+	meta?: OutputMeta;
+}
+
+function isTopLevelSession(session: ToolSession): boolean {
+	const depth = session.taskDepth;
+	return depth === undefined || depth === 0;
+}
+
+export class CheckpointTool implements AgentTool<typeof checkpointSchema, CheckpointToolDetails> {
+	readonly name = "checkpoint";
+	readonly label = "Checkpoint";
+	readonly description: string;
+	readonly parameters = checkpointSchema;
+	readonly strict = true;
+
+	constructor(private readonly session: ToolSession) {
+		this.description = renderPromptTemplate(checkpointDescription);
+	}
+
+	static createIf(session: ToolSession): CheckpointTool | null {
+		if (!isTopLevelSession(session)) return null;
+		return new CheckpointTool(session);
+	}
+
+	async execute(
+		_toolCallId: string,
+		params: CheckpointParams,
+		_signal?: AbortSignal,
+		_onUpdate?: AgentToolUpdateCallback<CheckpointToolDetails>,
+		_context?: AgentToolContext,
+	): Promise<AgentToolResult<CheckpointToolDetails>> {
+		if (!isTopLevelSession(this.session)) {
+			throw new ToolError("Checkpoint not available in subagents.");
+		}
+		if (this.session.getCheckpointState?.()) {
+			throw new ToolError("Checkpoint already active.");
+		}
+		const startedAt = new Date().toISOString();
+		return toolResult<CheckpointToolDetails>({ goal: params.goal, startedAt })
+			.text(
+				[
+					"Checkpoint created.",
+					`Goal: ${params.goal}`,
+					"Run your investigation, then call rewind with a concise report.",
+				].join("\n"),
+			)
+			.done();
+	}
+}
+
+export class RewindTool implements AgentTool<typeof rewindSchema, RewindToolDetails> {
+	readonly name = "rewind";
+	readonly label = "Rewind";
+	readonly description: string;
+	readonly parameters = rewindSchema;
+	readonly strict = true;
+
+	constructor(private readonly session: ToolSession) {
+		this.description = renderPromptTemplate(rewindDescription);
+	}
+
+	static createIf(session: ToolSession): RewindTool | null {
+		if (!isTopLevelSession(session)) return null;
+		return new RewindTool(session);
+	}
+
+	async execute(
+		_toolCallId: string,
+		params: RewindParams,
+		_signal?: AbortSignal,
+		_onUpdate?: AgentToolUpdateCallback<RewindToolDetails>,
+		_context?: AgentToolContext,
+	): Promise<AgentToolResult<RewindToolDetails>> {
+		if (!isTopLevelSession(this.session)) {
+			throw new ToolError("Checkpoint not available in subagents.");
+		}
+		if (!this.session.getCheckpointState?.()) {
+			throw new ToolError("No active checkpoint.");
+		}
+		const report = params.report.trim();
+		if (report.length === 0) {
+			throw new ToolError("Report cannot be empty.");
+		}
+		return toolResult<RewindToolDetails>({ report, rewound: true })
+			.text(["Rewind requested.", "Report captured for context replacement."].join("\n"))
+			.done();
+	}
+}

package/src/tools/index.ts

@@ -22,6 +22,7 @@ import { BashTool } from "./bash";
 import { BrowserTool } from "./browser";
 import { CalculatorTool } from "./calculator";
 import { CancelJobTool } from "./cancel-job";
+import { type CheckpointState, CheckpointTool, RewindTool } from "./checkpoint";
 import { ExitPlanModeTool } from "./exit-plan-mode";
 import { FetchTool } from "./fetch";
 import { FindTool } from "./find";
@@ -30,6 +31,7 @@ import { NotebookTool } from "./notebook";
 import { wrapToolWithMetaNotice } from "./output-meta";
 import { PythonTool } from "./python";
 import { ReadTool } from "./read";
+import { RenderMermaidTool } from "./render-mermaid";
 import { ResolveTool } from "./resolve";
 import { reportFindingTool } from "./review";
 import { loadSshTool } from "./ssh";
@@ -54,6 +56,7 @@ export * from "./bash";
 export * from "./browser";
 export * from "./calculator";
 export * from "./cancel-job";
+export * from "./checkpoint";
 export * from "./exit-plan-mode";
 export * from "./fetch";
 export * from "./find";
@@ -63,6 +66,7 @@ export * from "./notebook";
 export * from "./pending-action";
 export * from "./python";
 export * from "./read";
+export * from "./render-mermaid";
 export * from "./resolve";
 export * from "./review";
 export * from "./ssh";
@@ -143,6 +147,10 @@ export interface ToolSession {
 	setTodoPhases?: (phases: TodoPhase[]) => void;
 	/** Pending action store for preview/apply workflows */
 	pendingActionStore?: import("./pending-action").PendingActionStore;
+	/** Get active checkpoint state if any. */
+	getCheckpointState?: () => CheckpointState | undefined;
+	/** Set or clear active checkpoint state. */
+	setCheckpointState?: (state: CheckpointState | null) => void;
 }
 
 type ToolFactory = (session: ToolSession) => Tool | null | Promise<Tool | null>;
@@ -150,6 +158,7 @@ type ToolFactory = (session: ToolSession) => Tool | null | Promise<Tool | null>;
 export const BUILTIN_TOOLS: Record<string, ToolFactory> = {
 	ast_grep: s => new AstGrepTool(s),
 	ast_edit: s => new AstEditTool(s),
+	render_mermaid: s => new RenderMermaidTool(s),
 	ask: AskTool.createIf,
 	bash: s => new BashTool(s),
 	python: s => new PythonTool(s),
@@ -162,6 +171,8 @@ export const BUILTIN_TOOLS: Record<string, ToolFactory> = {
 	notebook: s => new NotebookTool(s),
 	read: s => new ReadTool(s),
 	browser: s => new BrowserTool(s),
+	checkpoint: CheckpointTool.createIf,
+	rewind: RewindTool.createIf,
 	task: TaskTool.create,
 	cancel_job: CancelJobTool.createIf,
 	await: AwaitTool.createIf,
@@ -271,6 +282,24 @@ export async function createTools(session: ToolSession, toolNames?: string[]): P
 	) {
 		requestedTools.push("bash");
 	}
+
+	// Auto-include AST counterparts when their text-based sibling is present
+	if (requestedTools) {
+		if (
+			requestedTools.includes("grep") &&
+			!requestedTools.includes("ast_grep") &&
+			session.settings.get("astGrep.enabled")
+		) {
+			requestedTools.push("ast_grep");
+		}
+		if (
+			requestedTools.includes("edit") &&
+			!requestedTools.includes("ast_edit") &&
+			session.settings.get("astEdit.enabled")
+		) {
+			requestedTools.push("ast_edit");
+		}
+	}
 	const allTools: Record<string, ToolFactory> = { ...BUILTIN_TOOLS, ...HIDDEN_TOOLS };
 	const isToolAllowed = (name: string) => {
 		if (name === "lsp") return enableLsp;
@@ -281,12 +310,14 @@ export async function createTools(session: ToolSession, toolNames?: string[]): P
 		if (name === "grep") return session.settings.get("grep.enabled");
 		if (name === "ast_grep") return session.settings.get("astGrep.enabled");
 		if (name === "ast_edit") return session.settings.get("astEdit.enabled");
+		if (name === "render_mermaid") return session.settings.get("renderMermaid.enabled");
 		if (name === "notebook") return session.settings.get("notebook.enabled");
 		if (name === "fetch") return session.settings.get("fetch.enabled");
 		if (name === "web_search") return session.settings.get("web_search.enabled");
 		if (name === "lsp") return session.settings.get("lsp.enabled");
 		if (name === "calc") return session.settings.get("calc.enabled");
 		if (name === "browser") return session.settings.get("browser.enabled");
+		if (name === "checkpoint" || name === "rewind") return session.settings.get("checkpoint.enabled");
 		if (name === "task") {
 			const maxDepth = session.settings.get("task.maxRecursionDepth") ?? 2;
 			const currentDepth = session.taskDepth ?? 0;

package/src/tools/render-mermaid.ts

@@ -0,0 +1,67 @@
+import type { AgentTool, AgentToolContext, AgentToolResult, AgentToolUpdateCallback } from "@oh-my-pi/pi-agent-core";
+import { type MermaidAsciiRenderOptions, renderMermaidAscii } from "@oh-my-pi/pi-utils";
+import { type Static, Type } from "@sinclair/typebox";
+import { renderPromptTemplate } from "../config/prompt-templates";
+import renderMermaidDescription from "../prompts/tools/render-mermaid.md" with { type: "text" };
+import type { ToolSession } from "./index";
+
+const renderMermaidSchema = Type.Object({
+	mermaid: Type.String({ description: "Mermaid graph source text" }),
+	config: Type.Optional(
+		Type.Object({
+			useAscii: Type.Optional(Type.Boolean()),
+			paddingX: Type.Optional(Type.Number()),
+			paddingY: Type.Optional(Type.Number()),
+			boxBorderPadding: Type.Optional(Type.Number()),
+		}),
+	),
+});
+
+type RenderMermaidParams = Static<typeof renderMermaidSchema>;
+
+function sanitizeRenderConfig(config: MermaidAsciiRenderOptions | undefined): MermaidAsciiRenderOptions | undefined {
+	if (!config) return undefined;
+	return {
+		useAscii: config.useAscii,
+		boxBorderPadding:
+			config.boxBorderPadding === undefined ? undefined : Math.max(0, Math.floor(config.boxBorderPadding)),
+		paddingX: config.paddingX === undefined ? undefined : Math.max(0, Math.floor(config.paddingX)),
+		paddingY: config.paddingY === undefined ? undefined : Math.max(0, Math.floor(config.paddingY)),
+	};
+}
+export interface RenderMermaidToolDetails {
+	artifactId?: string;
+}
+
+export class RenderMermaidTool implements AgentTool<typeof renderMermaidSchema, RenderMermaidToolDetails> {
+	readonly name = "render_mermaid";
+	readonly label = "RenderMermaid";
+	readonly description: string;
+	readonly parameters = renderMermaidSchema;
+	readonly strict = true;
+
+	constructor(private readonly session: ToolSession) {
+		this.description = renderPromptTemplate(renderMermaidDescription);
+	}
+
+	async execute(
+		_toolCallId: string,
+		params: RenderMermaidParams,
+		_signal?: AbortSignal,
+		_onUpdate?: AgentToolUpdateCallback<RenderMermaidToolDetails>,
+		_context?: AgentToolContext,
+	): Promise<AgentToolResult<RenderMermaidToolDetails>> {
+		const ascii = renderMermaidAscii(params.mermaid, sanitizeRenderConfig(params.config));
+		const { path: artifactPath, id: artifactId } =
+			(await this.session.allocateOutputArtifact?.("render_mermaid")) ?? {};
+		if (artifactPath) {
+			await Bun.write(artifactPath, ascii);
+		}
+
+		const artifactLine = artifactId ? `\n\nSaved artifact: artifact://${artifactId}` : "";
+		return {
+			content: [{ type: "text", text: `${ascii}${artifactLine}` }],
+			details: { artifactId },
+		};
+	}
+}

package/src/utils/prompt-format.ts

@@ -86,9 +86,8 @@ export function formatPromptContent(content: string, options: PromptFormatOption
 
 	for (let i = 0; i < lines.length; i++) {
 		let line = lines[i].trimEnd();
-		const trimmed = line.trimStart();
-
-		if (CODE_FENCE.test(trimmed)) {
+		let trimmedStart = line.trimStart();
+		if (CODE_FENCE.test(trimmedStart)) {
 			inCodeBlock = !inCodeBlock;
 			result.push(line);
 			continue;
@@ -102,30 +101,29 @@ export function formatPromptContent(content: string, options: PromptFormatOption
 		if (replaceAsciiSymbols) {
 			line = replaceCommonAsciiSymbols(line);
 		}
+		trimmedStart = line.trimStart();
+		const trimmed = line.trim();
 
-		const isOpeningXml = OPENING_XML.test(trimmed) && !trimmed.endsWith("/>");
-		if (isOpeningXml && line.length === trimmed.length) {
-			const match = OPENING_XML.exec(trimmed);
+		const isOpeningXml = OPENING_XML.test(trimmedStart) && !trimmedStart.endsWith("/>");
+		if (isOpeningXml && line.length === trimmedStart.length) {
+			const match = OPENING_XML.exec(trimmedStart);
 			if (match) topLevelTags.push(match[1]);
 		}
 
-		const closingMatch = CLOSING_XML.exec(trimmed);
+		const closingMatch = CLOSING_XML.exec(trimmedStart);
 		if (closingMatch) {
 			const tagName = closingMatch[1];
 			if (topLevelTags.length > 0 && topLevelTags[topLevelTags.length - 1] === tagName) {
-				line = trimmed;
 				topLevelTags.pop();
-			} else {
-				line = line.trimEnd();
 			}
-		} else if (isPreRender && trimmed.startsWith("{{")) {
-			line = trimmed;
-		} else if (TABLE_SEP.test(trimmed)) {
-			line = compactTableSep(trimmed);
-		} else if (TABLE_ROW.test(trimmed)) {
-			line = compactTableRow(trimmed);
-		} else {
-			line = line.trimEnd();
+		} else if (isPreRender && trimmedStart.startsWith("{{")) {
+			/* keep indentation as-is in pre-render for Handlebars markers */
+		} else if (TABLE_SEP.test(trimmedStart)) {
+			const leadingWhitespace = line.slice(0, line.length - trimmedStart.length);
+			line = `${leadingWhitespace}${compactTableSep(trimmedStart)}`;
+		} else if (TABLE_ROW.test(trimmedStart)) {
+			const leadingWhitespace = line.slice(0, line.length - trimmedStart.length);
+			line = `${leadingWhitespace}${compactTableRow(trimmedStart)}`;
 		}
 
 		if (shouldBoldRfc2119) {