@khanhicetea/pi-better-tool 0.1.0

package/src/text.ts

@@ -0,0 +1,129 @@
+/**
+ * Text normalization utilities.
+ *
+ * These are ports of the helpers pi's built-in edit tool uses (edit-diff.ts),
+ * kept byte-for-byte compatible in behavior so that suggestions produced by
+ * this package's diagnostics are guaranteed to match what the same matching
+ * engine accepts.
+ */
+
+export interface LineSpan {
+	/** Character offset of the line start (inclusive). */
+	start: number;
+	/** Character offset of the line end (exclusive, includes the trailing \n when present). */
+	end: number;
+}
+
+export function splitBom(content: string): { bom: string; text: string } {
+	return content.startsWith("") ? { bom: "", text: content.slice(1) } : { bom: "", text: content };
+}
+
+export function detectLineEnding(content: string): "\n" | "\r\n" {
+	const crlfIdx = content.indexOf("\r\n");
+	const lfIdx = content.indexOf("\n");
+	if (lfIdx === -1) return "\n";
+	if (crlfIdx === -1) return "\n";
+	return crlfIdx < lfIdx ? "\r\n" : "\n";
+}
+
+export function normalizeToLF(text: string): string {
+	return text.replace(/\r\n/g, "\n").replace(/\r/g, "\n");
+}
+
+export function restoreLineEndings(text: string, ending: "\n" | "\r\n"): string {
+	return ending === "\r\n" ? text.replace(/\n/g, "\r\n") : text;
+}
+
+/**
+ * Normalize text for fuzzy matching. Applies progressive transformations:
+ * - NFKC normalization
+ * - Strip trailing whitespace from each line
+ * - Normalize smart quotes to ASCII equivalents
+ * - Normalize Unicode dashes/hyphens to ASCII hyphen
+ * - Normalize special Unicode spaces to regular space
+ *
+ * Line-count preserving and idempotent.
+ */
+export function normalizeForFuzzyMatch(text: string): string {
+	return (
+		text
+			.normalize("NFKC")
+			// Strip trailing whitespace per line
+			.split("\n")
+			.map((line) => line.trimEnd())
+			.join("\n")
+			// Smart single quotes -> '
+			.replace(/[\u2018\u2019\u201A\u201B]/g, "'")
+			// Smart double quotes -> "
+			.replace(/[\u201C\u201D\u201E\u201F]/g, '"')
+			// Various dashes/hyphens -> -
+			// U+2010 hyphen, U+2011 non-breaking hyphen, U+2012 figure dash,
+			// U+2013 en-dash, U+2014 em-dash, U+2015 horizontal bar, U+2212 minus
+			.replace(/[\u2010\u2011\u2012\u2013\u2014\u2015\u2212]/g, "-")
+			// Special spaces -> regular space
+			// U+00A0 NBSP, U+2002-U+200A various spaces, U+202F narrow NBSP,
+			// U+205F medium math space, U+3000 ideographic space
+			.replace(/[\u00A0\u2002-\u200A\u202F\u205F\u3000]/g, " ")
+	);
+}
+
+export function splitLinesWithEndings(content: string): string[] {
+	return content.match(/[^\n]*\n|[^\n]+/g) ?? [];
+}
+
+export function getLineSpans(content: string): LineSpan[] {
+	let offset = 0;
+	return splitLinesWithEndings(content).map((line) => {
+		const span = { start: offset, end: offset + line.length };
+		offset = span.end;
+		return span;
+	});
+}
+
+/** 0-based index of the line containing `offset` (clamped to the last line). */
+export function lineAt(spans: LineSpan[], offset: number): number {
+	let lo = 0;
+	let hi = spans.length - 1;
+	let result = spans.length - 1;
+	while (lo <= hi) {
+		const mid = (lo + hi) >> 1;
+		if (spans[mid].start <= offset) {
+			result = mid;
+			lo = mid + 1;
+		} else {
+			hi = mid - 1;
+		}
+	}
+	return result;
+}
+
+/**
+ * Count occurrences of `needle` in fully fuzzy-normalized space, using the
+ * same normalization the matching engine uses for its uniqueness check.
+ * `fuzzyContent` must already be `normalizeForFuzzyMatch`-ed.
+ */
+export function countFuzzyOccurrences(fuzzyContent: string, needle: string): number {
+	if (!needle) return 0;
+	return fuzzyContent.split(normalizeForFuzzyMatch(needle)).length - 1;
+}
+
+/** All start offsets of `needle` in `haystack` (literal string search). */
+export function findAllOccurrences(haystack: string, needle: string): number[] {
+	const offsets: number[] = [];
+	if (!needle) return offsets;
+	let idx = haystack.indexOf(needle);
+	while (idx !== -1) {
+		offsets.push(idx);
+		idx = haystack.indexOf(needle, idx + needle.length);
+	}
+	return offsets;
+}
+
+/** Split into lines for similarity comparison (fuzzy-normalized, no trailing-empty artifact). */
+export function toFuzzyLines(text: string): string[] {
+	const lines = normalizeToLF(text).split("\n").map((line) => normalizeForFuzzyMatch(line));
+	if (lines.length > 0 && lines[lines.length - 1] === "") {
+		lines.pop();
+	}
+	return lines;
+}

package/src/tool.ts

@@ -0,0 +1,221 @@
+/**
+ * The "better edit" tool: a drop-in override of pi's built-in edit tool.
+ *
+ * Happy-path behavior is identical to the built-in tool (same schema, same
+ * matching semantics, same result shapes so the built-in diff renderer is
+ * inherited). The difference is failure behavior: instead of a bare "Could
+ * not find edits[1]" error that forces the model to re-read the file and
+ * guess at larger context, failures include recovery context:
+ *
+ * - ambiguous oldText → occurrence line numbers + the minimal prefix/suffix
+ *   context that disambiguates each occurrence, as copy-paste-ready snippets
+ * - not-found oldText → closest matching region with a per-line comparison
+ *   and the exact file bytes to retry with
+ */
+
+import {
+	generateDiffString,
+	generateUnifiedPatch,
+	withFileMutationQueue,
+	type EditToolDetails,
+	type ExtensionAPI,
+	type ExtensionContext,
+} from "@earendil-works/pi-coding-agent";
+import { constants } from "node:fs";
+import { access as fsAccess, readFile as fsReadFile, writeFile as fsWriteFile } from "node:fs/promises";
+import { homedir } from "node:os";
+import { resolve } from "node:path";
+import { type Static, Type } from "typebox";
+import { analyzeEdits, applyAnalysis, type EditOp } from "./apply.ts";
+import { formatEditFailure } from "./diagnostics.ts";
+import { detectLineEnding, normalizeToLF, restoreLineEndings, splitBom } from "./text.ts";
+
+const replaceEditSchema = Type.Object({
+	oldText: Type.String({
+		description:
+			"Exact text for one targeted replacement. It must be unique in the original file and must not overlap with any other edits[].oldText in the same call.",
+	}),
+	newText: Type.String({ description: "Replacement text for this targeted edit." }),
+});
+
+export const betterEditSchema = Type.Object({
+	path: Type.String({ description: "Path to the file to edit (relative or absolute)" }),
+	edits: Type.Array(replaceEditSchema, {
+		description:
+			"One or more targeted replacements. Each edit is matched against the original file, not incrementally. Do not include overlapping or nested edits. If two changes touch the same block or nearby lines, merge them into one edit instead.",
+	}),
+});
+
+export type BetterEditInput = Static<typeof betterEditSchema>;
+
+function resolveToCwd(filePath: string, cwd: string): string {
+	let path = filePath;
+	if (path === "~") {
+		path = homedir();
+	} else if (path.startsWith("~/")) {
+		path = resolve(homedir(), path.slice(2));
+	}
+	return resolve(cwd, path);
+}
+
+function isSingleEditInput(value: unknown): value is { oldText: string; newText: string } {
+	if (!value || typeof value !== "object" || Array.isArray(value)) {
+		return false;
+	}
+	const edit = value as { oldText?: unknown; newText?: unknown };
+	return typeof edit.oldText === "string" && typeof edit.newText === "string";
+}
+
+/**
+ * Compatibility shim, ported from the built-in tool: some models send edits as
+ * a JSON string, or send a single edit object instead of a one-element array,
+ * or use the legacy top-level oldText/newText shape.
+ */
+export function prepareEditArguments(input: unknown): unknown {
+	if (!input || typeof input !== "object") {
+		return input;
+	}
+	const args = input as {
+		edits?: unknown;
+		oldText?: unknown;
+		newText?: unknown;
+		[path: string]: unknown;
+	};
+
+	if (typeof args.edits === "string") {
+		try {
+			const parsed: unknown = JSON.parse(args.edits);
+			if (Array.isArray(parsed)) {
+				args.edits = parsed;
+			} else if (isSingleEditInput(parsed)) {
+				args.edits = [parsed];
+			}
+		} catch {
+			// leave as-is; schema validation will report it
+		}
+	} else if (isSingleEditInput(args.edits)) {
+		args.edits = [args.edits];
+	}
+
+	if (typeof args.oldText === "string" && typeof args.newText === "string") {
+		const edits = Array.isArray(args.edits) ? [...(args.edits as EditOp[])] : [];
+		edits.push({ oldText: args.oldText, newText: args.newText });
+		const { oldText: _oldText, newText: _newText, ...rest } = args;
+		return { ...rest, edits };
+	}
+	return args;
+}
+
+export interface BetterEditSuccess {
+	content: Array<{ type: "text"; text: string }>;
+	details: EditToolDetails;
+}
+
+export async function executeBetterEdit(
+	input: BetterEditInput,
+	signal: AbortSignal | undefined,
+	ctx: Pick<ExtensionContext, "cwd">,
+): Promise<BetterEditSuccess> {
+	const edits: EditOp[] = input.edits ?? [];
+	if (!Array.isArray(edits) || edits.length === 0) {
+		throw new Error("Edit tool input is invalid. edits must contain at least one replacement.");
+	}
+	const absolutePath = resolveToCwd(input.path, ctx.cwd);
+
+	return withFileMutationQueue(absolutePath, async () => {
+		// Do not reject from an abort event listener here: that would release the
+		// mutation queue while an in-flight filesystem operation may still finish.
+		const throwIfAborted = () => {
+			if (signal?.aborted) throw new Error("Operation aborted");
+		};
+		throwIfAborted();
+
+		try {
+			await fsAccess(absolutePath, constants.R_OK | constants.W_OK);
+		} catch (error) {
+			throwIfAborted();
+			const errorMessage =
+				error instanceof Error && "code" in error
+					? `Error code: ${(error as NodeJS.ErrnoException).code}`
+					: String(error);
+			throw new Error(`Could not edit file: ${input.path}. ${errorMessage}.`);
+		}
+		throwIfAborted();
+
+		const buffer = await fsReadFile(absolutePath);
+		const rawContent = buffer.toString("utf-8");
+		throwIfAborted();
+
+		const { bom, text: content } = splitBom(rawContent);
+		const originalEnding = detectLineEnding(content);
+		const normalizedContent = normalizeToLF(content);
+
+		const result = analyzeEdits(normalizedContent, edits);
+		if (!result.ok) {
+			throw new Error(
+				formatEditFailure({
+					path: input.path,
+					normalizedContent,
+					edits,
+					failure: result.failure,
+				}),
+			);
+		}
+
+		const { baseContent, newContent } = applyAnalysis(normalizedContent, result.analysis);
+		if (baseContent === newContent) {
+			throw new Error(
+				formatEditFailure({
+					path: input.path,
+					normalizedContent,
+					edits,
+					failure: { kind: "no-change" },
+				}),
+			);
+		}
+		throwIfAborted();
+
+		const finalContent = bom + restoreLineEndings(newContent, originalEnding);
+		await fsWriteFile(absolutePath, finalContent, "utf-8");
+		throwIfAborted();
+
+		const diffResult = generateDiffString(baseContent, newContent);
+		const patch = generateUnifiedPatch(input.path, baseContent, newContent);
+		return {
+			content: [
+				{
+					type: "text",
+					text: `Successfully replaced ${edits.length} block(s) in ${input.path}.`,
+				},
+			],
+			details: {
+				diff: diffResult.diff,
+				patch,
+				firstChangedLine: diffResult.firstChangedLine,
+			} satisfies EditToolDetails,
+		};
+	});
+}
+
+export function registerBetterEditTool(pi: ExtensionAPI): void {
+	pi.registerTool({
+		name: "edit",
+		label: "edit",
+		description:
+			"Edit a single file using exact text replacement. Every edits[].oldText must match a unique, non-overlapping region of the original file. If two changes affect the same block or nearby lines, merge them into one edit instead of emitting overlapping edits. Do not include large unchanged regions just to connect distant changes. On failure the error includes recovery context (closest matching region or per-occurrence disambiguation snippets) so you can retry immediately without re-reading the file.",
+		promptSnippet:
+			"Make precise file edits with exact text replacement; failures return recovery context (closest match or disambiguation snippets)",
+		promptGuidelines: [
+			"Use edit for precise changes (edits[].oldText must match exactly)",
+			"When changing multiple separate locations in one file, use one edit call with multiple entries in edits[] instead of multiple edit calls",
+			"Each edits[].oldText is matched against the original file, not after earlier edits are applied. Do not emit overlapping or nested edits. Merge nearby changes into one edit.",
+			"Keep edits[].oldText as small as possible while still being unique in the file. Do not pad with large unchanged regions.",
+			"When an edit call fails, the error message already contains recovery context: the closest matching region with exact file bytes (not-found) or each occurrence with a ready-to-use disambiguated oldText (ambiguous match). Retry the edit using that text directly instead of re-reading the file.",
+		],
+		parameters: betterEditSchema,
+		prepareArguments: (args: unknown): BetterEditInput => prepareEditArguments(args) as BetterEditInput,
+		async execute(_toolCallId, input, signal, _onUpdate, ctx) {
+			return executeBetterEdit(input, signal, ctx);
+		},
+	});
+}