@wierdbytes/pi-facelift 0.4.0 → 0.4.1

package/index.ts

@@ -25,7 +25,7 @@
 
 import * as childProcess from "node:child_process";
 import { existsSync, readFileSync } from "node:fs";
-import { basename, dirname, extname, join, relative } from "node:path";
+import { basename, dirname, extname, join, relative, resolve as resolvePath } from "node:path";
 
 import type { ImageContent, TextContent } from "@earendil-works/pi-ai";
 import type {
@@ -1782,6 +1782,52 @@ export default function piFaceliftExtension(pi: PiFaceliftApi, deps?: PiFacelift
 		},
 	});
 
+	// ===================================================================
+	// Per-path mutation lock
+	// -------------------------------------------------------------------
+	// Both `write` and `edit` snapshot the file *before* delegating to the
+	// built-in tool's execute(). When two mutations to the same path are
+	// queued concurrently the built-in tool serializes the actual write,
+	// but our pre-execute snapshots run in parallel — so the second call
+	// would diff against the file state from *before* the first call, not
+	// the state the first call left behind. Result: stale, misleading
+	// diffs in the rendered output.
+	//
+	// We fix it here (not in the built-in tool) by chaining
+	// snapshot→execute pairs per resolved path. The chain is purely
+	// in-process and only covers mutations going through our wrapper;
+	// external writers (other processes, the user) are still racy, but
+	// that's the same correctness contract the built-in tool already has.
+	// ===================================================================
+
+	const pathMutationTails = new Map<string, Promise<void>>();
+	function serializePerPath<T>(path: string, fn: () => Promise<T>): Promise<T> {
+		if (!path) return fn();
+		const key = (() => {
+			try {
+				return resolvePath(path);
+			} catch {
+				return path;
+			}
+		})();
+		const prev = pathMutationTails.get(key) ?? Promise.resolve();
+		// `prev.then(fn, fn)` — run fn after prev settles, regardless of
+		// whether prev rejected. We must not let an earlier failure poison
+		// later snapshots on the same path.
+		const result = prev.then(fn, fn);
+		const tail = result.then(
+			() => undefined,
+			() => undefined,
+		);
+		pathMutationTails.set(key, tail);
+		tail.then(() => {
+			// Drop the entry only if no newer caller has chained onto us;
+			// otherwise we'd unblock callers waiting in the chain.
+			if (pathMutationTails.get(key) === tail) pathMutationTails.delete(key);
+		});
+		return result;
+	}
+
 	// ===================================================================
 	// write — split diff (or syntax-highlighted preview for new files)
 	// ===================================================================
@@ -1802,14 +1848,19 @@ export default function piFaceliftExtension(pi: PiFaceliftApi, deps?: PiFacelift
 				ctx: ExtensionContext,
 			) {
 				const fp = params.path ?? "";
-				let oldContent: string | null = null;
-				try {
-					if (fp && existsSync(fp)) oldContent = readFileSync(fp, "utf-8");
-				} catch {
-					oldContent = null;
-				}
-
-				const result = (await origWrite.execute(tid, params, sig, upd, ctx)) as ToolResultLike;
+				// Snapshot + execute must run atomically per-path so concurrent
+				// writes to the same file don't both diff against the same
+				// stale pre-state.
+				const { oldContent, result } = await serializePerPath(fp, async () => {
+					let snapshot: string | null = null;
+					try {
+						if (fp && existsSync(fp)) snapshot = readFileSync(fp, "utf-8");
+					} catch {
+						snapshot = null;
+					}
+					const r = (await origWrite.execute(tid, params, sig, upd, ctx)) as ToolResultLike;
+					return { oldContent: snapshot, result: r };
+				});
 				const newContent = params.content ?? "";
 				const language = diffLang(fp);
 
@@ -1875,6 +1926,9 @@ export default function piFaceliftExtension(pi: PiFaceliftApi, deps?: PiFacelift
 
 				if (d?._type === "writeNew") {
 					const lab = `${FG_GREEN}✓ new file${RST} ${FG_DIM}(${d.lineCount} lines)${RST}`;
+					// Diff/preview output is always rendered in full — no compact
+					// mode, no `… N more lines` footer — so `ctx.expanded` is
+					// irrelevant to the rendered body and stays out of the key.
 					const key = `writeNew:${themeCacheKey(theme)}:${d.filePath}:${d.lineCount}:${w}:${status}`;
 					if (ctx.state._wk !== key) {
 						ctx.state._wk = key;
@@ -1883,11 +1937,7 @@ export default function piFaceliftExtension(pi: PiFaceliftApi, deps?: PiFacelift
 							hlDiffBlock(d.content, d.language)
 								.then((hlLines: string[]) => {
 									if (ctx.state._wk !== key) return;
-									const maxShow = ctx.expanded ? hlLines.length : MAX_PREVIEW_LINES;
-									const preview = hlLines.slice(0, maxShow).join("\n");
-									const rem = hlLines.length - maxShow;
-									const body =
-										rem > 0 ? `${preview}\n${FG_DIM}… ${rem} more lines${RST}` : preview;
+									const body = hlLines.join("\n");
 									ctx.state._wt = frameResultWithBottomLabel(body, lab, status, t, w);
 									ctx.invalidate();
 								})
@@ -1903,11 +1953,14 @@ export default function piFaceliftExtension(pi: PiFaceliftApi, deps?: PiFacelift
 
 				if (d?._type === "writeDiff") {
 					const lab = `${summarizeDiff(d.diff.added, d.diff.removed)} ${FG_DIM}(${d.diff.lines.length} diff lines)${RST}`;
+					// Diff output is always rendered in full — no compact mode,
+					// no `… N more lines` footer — so `ctx.expanded` is
+					// irrelevant to the rendered body and stays out of the key.
 					const key = `writeDiff:${themeCacheKey(theme)}:${d.filePath}:${d.diff.added}:${d.diff.removed}:${d.diff.lines.length}:${w}:${status}`;
 					if (ctx.state._wk !== key) {
 						ctx.state._wk = key;
 						ctx.state._wt = frameResultWithBottomLabel("", lab, status, t, w);
-						const maxLines = ctx.expanded ? d.diff.lines.length : MAX_PREVIEW_LINES;
+						const maxLines = d.diff.lines.length;
 						const innerW = Math.max(40, w - 1);
 						const layout = decideDiffLayout([d.diff], innerW, maxLines);
 						// Reserve 1 col for the outer frame rail.
@@ -1994,14 +2047,21 @@ export default function piFaceliftExtension(pi: PiFaceliftApi, deps?: PiFacelift
 				// `parseDiff(op.oldText, op.newText)` numbers every diff from
 				// line 1 and the rendered gutter ends up off by the line
 				// offset of the edit inside the file.
-				let preEditFile: string | null = null;
-				try {
-					if (fp && existsSync(fp)) preEditFile = readFileSync(fp, "utf-8");
-				} catch {
-					preEditFile = null;
-				}
-
-				const result = (await origEdit.execute(tid, params, sig, upd, ctx)) as ToolResultLike;
+				//
+				// Pair the snapshot with the underlying execute() inside a
+				// per-path lock so concurrent edits to the same file don't
+				// both snapshot the same stale pre-state and render diffs
+				// against the wrong baseline.
+				const { preEditFile, result } = await serializePerPath(fp, async () => {
+					let snapshot: string | null = null;
+					try {
+						if (fp && existsSync(fp)) snapshot = readFileSync(fp, "utf-8");
+					} catch {
+						snapshot = null;
+					}
+					const r = (await origEdit.execute(tid, params, sig, upd, ctx)) as ToolResultLike;
+					return { preEditFile: snapshot, result: r };
+				});
 				const ops = getEditOperations(params);
 				if (ops.length === 0) return result;
 
@@ -2082,35 +2142,39 @@ export default function piFaceliftExtension(pi: PiFaceliftApi, deps?: PiFacelift
 					editCount === 1 ? `1 edit ${summary}` : `${editCount} edits ${summary}`;
 				const lab = `${editsLabel} ${FG_DIM}(${diffLineTotal} diff lines)${RST}`;
 
-				const key = `editDiff:${themeCacheKey(theme)}:${d.filePath}:${editCount}:${d.totalAdded}:${d.totalRemoved}:${diffLineTotal}:${w}:${status}:${ctx.expanded ? 1 : 0}`;
+				// Every edit block is rendered in full — no compact mode,
+				// no `… N more edit blocks` footer, no per-edit line cap —
+				// so `ctx.expanded` doesn't affect the body and stays out of
+				// the key.
+				const key = `editDiff:${themeCacheKey(theme)}:${d.filePath}:${editCount}:${d.totalAdded}:${d.totalRemoved}:${diffLineTotal}:${w}:${status}`;
 				if (ctx.state._ek !== key) {
 					ctx.state._ek = key;
 					ctx.state._et = frameResultWithBottomLabel("", lab, status, t, w);
 
 					const innerW = Math.max(40, w - 1);
-					const perEditCap = ctx.expanded
-						? Number.POSITIVE_INFINITY
-						: Math.max(8, Math.floor(MAX_PREVIEW_LINES / Math.max(1, Math.min(editCount, 3))));
-					const shownEdits = ctx.expanded ? d.edits : d.edits.slice(0, 3);
-					const remaining = d.edits.length - shownEdits.length;
 
 					// Pick one layout for every edit in this call ("consistent" mode,
 					// the default) so we never end up with `Edit 1 split, Edit 2
 					// unified`. The user can override via .pi/settings.json's
-					// `diffLayout` ("split"/"unified"/"per-edit").
-					const probeCap = Number.isFinite(perEditCap) ? perEditCap : MAX_PREVIEW_LINES;
+					// `diffLayout` ("split"/"unified"/"per-edit"). We probe layout
+					// with the longest edit so a single tall diff doesn't force
+					// every other edit into unified mode unnecessarily.
+					const probeCap = Math.max(
+						1,
+						...d.edits.map((e) => e.diff.lines.length),
+					);
 					const layout = decideDiffLayout(
-						shownEdits.map((e) => e.diff),
+						d.edits.map((e) => e.diff),
 						innerW,
 						probeCap,
 					);
 
 					Promise.all(
-						shownEdits.map((edit, idx) =>
+						d.edits.map((edit, idx) =>
 							renderSplit(
 								edit.diff,
 								d.language,
-								Number.isFinite(perEditCap) ? perEditCap : edit.diff.lines.length,
+								edit.diff.lines.length,
 								dc,
 								innerW,
 								{ frameless: true, layout },
@@ -2129,12 +2193,7 @@ export default function piFaceliftExtension(pi: PiFaceliftApi, deps?: PiFacelift
 					)
 						.then((sections) => {
 							if (ctx.state._ek !== key) return;
-							const joined = sections.join("\n\n");
-							const suffix =
-								remaining > 0
-									? `\n\n${FG_DIM}… ${remaining} more edit block${remaining === 1 ? "" : "s"}${RST}`
-									: "";
-							ctx.state._et = frameResultWithBottomLabel(`${joined}${suffix}`, lab, status, t, w);
+							ctx.state._et = frameResultWithBottomLabel(sections.join("\n\n"), lab, status, t, w);
 							ctx.invalidate();
 						})
 						.catch((err: unknown) => {

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@wierdbytes/pi-facelift",
-  "version": "0.4.0",
+  "version": "0.4.1",
   "description": "Cosmetic facelift for built-in pi tool output — open-right rounded frames, syntax-highlighted reads, status-aware bash summaries, find/grep.",
   "type": "module",
   "files": [

package/write-edit-concurrency.test.ts

@@ -0,0 +1,300 @@
+/**
+ * Regression tests for the per-path mutation lock that pairs the
+ * pre-execute file snapshot with the underlying tool's execute() call.
+ *
+ * Bug history: when two `write` (or `edit`) calls for the *same* path
+ * landed in the agent's tool queue concurrently, both wrappers ran
+ *
+ *     oldContent = readFileSync(path)
+ *     result     = await origExecute(...)
+ *
+ * in parallel. The built-in tool serialized the actual writes, but both
+ * wrappers had already captured the same pre-state snapshot — so the
+ * second wrapper rendered its diff against the *original* file rather
+ * than against whatever the first wrapper had just persisted.
+ *
+ * The fix routes snapshot + execute through `serializePerPath()`, an
+ * in-process per-path promise chain. These tests pin that ordering by
+ * making the mock origExecute() the thing that actually updates the
+ * file on disk: if the snapshot ran before our chained predecessor's
+ * execute(), the second snapshot would be the original contents and
+ * the diff details would prove it.
+ */
+
+import { describe, expect, it } from "vitest";
+
+import { mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+
+import piFaceliftExtension from "./index.ts";
+
+class MockText {
+	private text = "";
+	constructor(_text = "", _x = 0, _y = 0) {}
+	setText(value: string) {
+		this.text = value;
+	}
+	getText() {
+		return this.text;
+	}
+}
+
+function mockToolFactory(exec: any) {
+	return (_cwd: string) => ({
+		name: "mock",
+		description: "mock",
+		parameters: { type: "object", properties: {} },
+		execute: exec,
+	});
+}
+
+function loadTools(opts: { writeExec?: any; editExec?: any } = {}) {
+	const noopExec = async () => ({ content: [{ type: "text", text: "" }] });
+	const tools = new Map<string, any>();
+	const pi = {
+		registerTool: (tool: any) => tools.set(tool.name, tool),
+		registerCommand: () => {},
+		on: () => {},
+	};
+
+	piFaceliftExtension(pi, {
+		sdk: {
+			createReadToolDefinition: mockToolFactory(noopExec),
+			createBashToolDefinition: mockToolFactory(noopExec),
+			createLsToolDefinition: mockToolFactory(noopExec),
+			createFindToolDefinition: mockToolFactory(noopExec),
+			createGrepToolDefinition: mockToolFactory(noopExec),
+			createWriteToolDefinition: mockToolFactory(opts.writeExec ?? noopExec),
+			createEditToolDefinition: mockToolFactory(opts.editExec ?? noopExec),
+			getAgentDir: () => "/tmp/pi-facelift-test",
+		},
+		TextComponent: MockText,
+	});
+
+	return tools;
+}
+
+interface DiffLine {
+	type: "ctx" | "add" | "del" | "sep";
+	oldNum: number | null;
+	newNum: number | null;
+	content: string;
+}
+
+interface WriteDiffDetails {
+	_type: "writeDiff";
+	filePath: string;
+	diff: { lines: DiffLine[]; added: number; removed: number };
+}
+
+interface EditDiffDetails {
+	_type: "editDiff";
+	edits: Array<{ oldText: string; newText: string; diff: { lines: DiffLine[] } }>;
+}
+
+describe("write execute: per-path snapshot/execute serialization", () => {
+	it("the second concurrent write diffs against the first write's result, not the original file", async () => {
+		const dir = mkdtempSync(join(tmpdir(), "facelift-write-concur-"));
+		const file = join(dir, "state.txt");
+		writeFileSync(file, "A\n", "utf-8");
+
+		// `writeExec` is what the agent SDK's built-in write tool would
+		// do: actually persist `params.content` to disk. By doing this
+		// synchronously inside the mock we make the "second snapshot
+		// must run after the first execute" property *observable* —
+		// the resulting `details.diff` shows which baseline was used.
+		const writeExec = async (
+			_tid: string,
+			params: { path: string; content: string },
+		) => {
+			// Yield once so a parallel caller has a chance to race with us.
+			await Promise.resolve();
+			writeFileSync(params.path, params.content, "utf-8");
+			return { content: [{ type: "text", text: `wrote ${params.path}` }] };
+		};
+		const tools = loadTools({ writeExec });
+		const writeTool = tools.get("write");
+
+		// Fire both calls without awaiting — the lock should serialize
+		// them so the second one sees "B\n", not "A\n".
+		const [r1, r2] = await Promise.all([
+			writeTool.execute("t1", { path: file, content: "B\n" }, null, null, {}),
+			writeTool.execute("t2", { path: file, content: "C\n" }, null, null, {}),
+		]);
+
+		const d1 = r1.details as WriteDiffDetails;
+		const d2 = r2.details as WriteDiffDetails;
+
+		expect(d1._type).toBe("writeDiff");
+		expect(d2._type).toBe("writeDiff");
+
+		// First call: A → B. One deletion of "A", one addition of "B".
+		const d1Del = d1.diff.lines.find((l) => l.type === "del");
+		const d1Add = d1.diff.lines.find((l) => l.type === "add");
+		expect(d1Del?.content).toBe("A");
+		expect(d1Add?.content).toBe("B");
+
+		// Second call: with the bug, oldContent would still be "A" and
+		// the diff would say "A → C". With the lock, oldContent is the
+		// state the first call left behind ("B"), so the diff is B → C.
+		const d2Del = d2.diff.lines.find((l) => l.type === "del");
+		const d2Add = d2.diff.lines.find((l) => l.type === "add");
+		expect(d2Del?.content).toBe("B");
+		expect(d2Add?.content).toBe("C");
+
+		// And the file on disk reflects the second write.
+		expect(readFileSync(file, "utf-8")).toBe("C\n");
+
+		rmSync(dir, { recursive: true, force: true });
+	});
+
+	it("a failing first write does not poison subsequent writes to the same path", async () => {
+		const dir = mkdtempSync(join(tmpdir(), "facelift-write-poison-"));
+		const file = join(dir, "state.txt");
+		writeFileSync(file, "A\n", "utf-8");
+
+		let call = 0;
+		const writeExec = async (
+			_tid: string,
+			params: { path: string; content: string },
+		) => {
+			call += 1;
+			await Promise.resolve();
+			if (call === 1) throw new Error("boom");
+			writeFileSync(params.path, params.content, "utf-8");
+			return { content: [{ type: "text", text: "ok" }] };
+		};
+		const tools = loadTools({ writeExec });
+		const writeTool = tools.get("write");
+
+		// First call rejects; the chain must still let the second call
+		// proceed and snapshot the (untouched) "A\n" baseline.
+		const p1 = writeTool
+			.execute("t1", { path: file, content: "B\n" }, null, null, {})
+			.catch((e: Error) => e);
+		const p2 = writeTool.execute("t2", { path: file, content: "C\n" }, null, null, {});
+		const [r1, r2] = await Promise.all([p1, p2]);
+
+		expect(r1).toBeInstanceOf(Error);
+		const d2 = (r2 as any).details as WriteDiffDetails;
+		expect(d2._type).toBe("writeDiff");
+		// First call never wrote, so the snapshot for the second call
+		// is the original "A\n".
+		const del = d2.diff.lines.find((l) => l.type === "del");
+		const add = d2.diff.lines.find((l) => l.type === "add");
+		expect(del?.content).toBe("A");
+		expect(add?.content).toBe("C");
+
+		rmSync(dir, { recursive: true, force: true });
+	});
+});
+
+describe("edit execute: per-path snapshot/execute serialization", () => {
+	it("the second concurrent edit sees the file state produced by the first edit", async () => {
+		const dir = mkdtempSync(join(tmpdir(), "facelift-edit-concur-"));
+		const file = join(dir, "src.ts");
+		const original = ["alpha", "beta", "gamma"].join("\n");
+		writeFileSync(file, original, "utf-8");
+
+		// Mock built-in edit: replace oldText → newText in the file so
+		// the second snapshot can prove it ran after the first execute.
+		const editExec = async (
+			_tid: string,
+			params: { path: string; edits: Array<{ oldText: string; newText: string }> },
+		) => {
+			await Promise.resolve();
+			let body = readFileSync(params.path, "utf-8");
+			for (const e of params.edits) body = body.replace(e.oldText, e.newText);
+			writeFileSync(params.path, body, "utf-8");
+			return { content: [{ type: "text", text: "ok" }] };
+		};
+		const tools = loadTools({ editExec });
+		const editTool = tools.get("edit");
+
+		const [r1, r2] = await Promise.all([
+			editTool.execute(
+				"t1",
+				{ path: file, edits: [{ oldText: "alpha", newText: "ALPHA" }] },
+				null,
+				null,
+				{},
+			),
+			// Targets the *result* of edit 1. Only succeeds if the
+			// snapshot for edit 2 happens after edit 1's execute().
+			editTool.execute(
+				"t2",
+				{ path: file, edits: [{ oldText: "ALPHA", newText: "Alpha!" }] },
+				null,
+				null,
+				{},
+			),
+		]);
+
+		const d2 = r2.details as EditDiffDetails;
+		expect(d2._type).toBe("editDiff");
+		// The edit operated on the line that previously was "alpha"
+		// (file line 1). Without the lock the snapshot would still show
+		// "alpha", `findLineOffset` for "ALPHA" would miss, and the
+		// diff would silently number from line 1 of a fake baseline.
+		const del = d2.edits[0].diff.lines.find((l) => l.type === "del");
+		const add = d2.edits[0].diff.lines.find((l) => l.type === "add");
+		expect(del?.content).toBe("ALPHA");
+		expect(add?.content).toBe("Alpha!");
+		expect(del?.oldNum).toBe(1);
+		expect(add?.newNum).toBe(1);
+
+		// Final on-disk content reflects both edits applied in order.
+		expect(readFileSync(file, "utf-8")).toBe(["Alpha!", "beta", "gamma"].join("\n"));
+
+		rmSync(dir, { recursive: true, force: true });
+	});
+
+	it("does not serialize unrelated paths against each other", async () => {
+		const dir = mkdtempSync(join(tmpdir(), "facelift-edit-indep-"));
+		const fileA = join(dir, "a.txt");
+		const fileB = join(dir, "b.txt");
+		writeFileSync(fileA, "a\n", "utf-8");
+		writeFileSync(fileB, "b\n", "utf-8");
+
+		// First call to `fileA` is intentionally slow. Second call to
+		// `fileB` must NOT wait for it.
+		let order: string[] = [];
+		const editExec = async (
+			_tid: string,
+			params: { path: string; edits: Array<{ oldText: string; newText: string }> },
+		) => {
+			order.push(`start:${params.path}`);
+			if (params.path === fileA) await new Promise((r) => setTimeout(r, 30));
+			order.push(`end:${params.path}`);
+			return { content: [{ type: "text", text: "ok" }] };
+		};
+		const tools = loadTools({ editExec });
+		const editTool = tools.get("edit");
+
+		await Promise.all([
+			editTool.execute(
+				"a",
+				{ path: fileA, edits: [{ oldText: "a", newText: "aa" }] },
+				null,
+				null,
+				{},
+			),
+			editTool.execute(
+				"b",
+				{ path: fileB, edits: [{ oldText: "b", newText: "bb" }] },
+				null,
+				null,
+				{},
+			),
+		]);
+
+		// fileB should finish before fileA — proves the lock is
+		// per-path, not a global mutex.
+		const endA = order.indexOf(`end:${fileA}`);
+		const endB = order.indexOf(`end:${fileB}`);
+		expect(endB).toBeLessThan(endA);
+
+		rmSync(dir, { recursive: true, force: true });
+	});
+});