fix(coding-agent): stabilize edit diff previews closes #3134

packages/coding-agent/src/core/extensions/types.ts

@@ -379,6 +379,8 @@ export interface ToolDefinition<TParams extends TSchema = TSchema, TDetails = un
 	promptGuidelines?: string[];
 	/** Parameter schema (TypeBox) */
 	parameters: TParams;
+	/** Controls whether ToolExecutionComponent renders the standard colored shell or the tool renders its own framing. */
+	renderShell?: "default" | "self";
 
 	/** Optional compatibility shim to prepare raw tool call arguments before schema validation. Must return an object conforming to TParams. */
 	prepareArguments?: (args: unknown) => Static<TParams>;

packages/coding-agent/src/core/tools/edit.ts

@@ -1,5 +1,5 @@
 import type { AgentTool } from "@mariozechner/pi-agent-core";
-import { Container, Text } from "@mariozechner/pi-tui";
+import { Box, Container, Spacer, Text } from "@mariozechner/pi-tui";
 import { type Static, Type } from "@sinclair/typebox";
 import { constants } from "fs";
 import { access as fsAccess, readFile as fsReadFile, writeFile as fsWriteFile } from "fs/promises";
@@ -7,8 +7,11 @@ import { renderDiff } from "../../modes/interactive/components/diff.js";
 import type { ToolDefinition } from "../extensions/types.js";
 import {
 	applyEditsToNormalizedContent,
+	computeEditsDiff,
 	detectLineEnding,
 	type Edit,
+	type EditDiffError,
+	type EditDiffResult,
 	generateDiffString,
 	normalizeToLF,
 	restoreLineEndings,
@@ -19,7 +22,11 @@ import { resolveToCwd } from "./path-utils.js";
 import { invalidArgText, shortenPath, str } from "./render-utils.js";
 import { wrapToolDefinition } from "./tool-definition-wrapper.js";
 
-type EditRenderState = Record<string, never>;
+type EditPreview = EditDiffResult | EditDiffError;
+
+type EditRenderState = {
+	callComponent?: EditCallRenderComponent;
+};
 
 const replaceEditSchema = Type.Object(
 	{
@@ -111,6 +118,66 @@ type RenderableEditArgs = {
 	newText?: string;
 };
 
+type EditToolResultLike = {
+	content: Array<{ type: string; text?: string; data?: string; mimeType?: string }>;
+	details?: EditToolDetails;
+};
+
+type EditCallRenderComponent = Box & {
+	preview?: EditPreview;
+	previewArgsKey?: string;
+	previewPending?: boolean;
+	settledError?: boolean;
+};
+
+function createEditCallRenderComponent(): EditCallRenderComponent {
+	return Object.assign(new Box(1, 1, (text: string) => text), {
+		preview: undefined as EditPreview | undefined,
+		previewArgsKey: undefined as string | undefined,
+		previewPending: false,
+		settledError: false,
+	});
+}
+
+function getEditCallRenderComponent(state: EditRenderState, lastComponent: unknown): EditCallRenderComponent {
+	if (lastComponent instanceof Box) {
+		const component = lastComponent as EditCallRenderComponent;
+		state.callComponent = component;
+		return component;
+	}
+	if (state.callComponent) {
+		return state.callComponent;
+	}
+	const component = createEditCallRenderComponent();
+	state.callComponent = component;
+	return component;
+}
+
+function getRenderablePreviewInput(args: RenderableEditArgs | undefined): { path: string; edits: Edit[] } | null {
+	if (!args) {
+		return null;
+	}
+
+	const path = typeof args.path === "string" ? args.path : typeof args.file_path === "string" ? args.file_path : null;
+	if (!path) {
+		return null;
+	}
+
+	if (
+		Array.isArray(args.edits) &&
+		args.edits.length > 0 &&
+		args.edits.every((edit) => typeof edit?.oldText === "string" && typeof edit?.newText === "string")
+	) {
+		return { path, edits: args.edits };
+	}
+
+	if (typeof args.oldText === "string" && typeof args.newText === "string") {
+		return { path, edits: [{ oldText: args.oldText, newText: args.newText }] };
+	}
+
+	return null;
+}
+
 function formatEditCall(
 	args: RenderableEditArgs | undefined,
 	theme: typeof import("../../modes/interactive/theme/theme.js").theme,
@@ -124,30 +191,88 @@ function formatEditCall(
 
 function formatEditResult(
 	args: RenderableEditArgs | undefined,
-	result: {
-		content: Array<{ type: string; text?: string; data?: string; mimeType?: string }>;
-		details?: EditToolDetails;
-	},
+	preview: EditPreview | undefined,
+	result: EditToolResultLike,
 	theme: typeof import("../../modes/interactive/theme/theme.js").theme,
 	isError: boolean,
 ): string | undefined {
 	const rawPath = str(args?.file_path ?? args?.path);
+	const previewDiff = preview && !("error" in preview) ? preview.diff : undefined;
+	const previewError = preview && "error" in preview ? preview.error : undefined;
 	if (isError) {
 		const errorText = result.content
 			.filter((c) => c.type === "text")
 			.map((c) => c.text || "")
 			.join("\n");
-		if (!errorText) {
+		if (!errorText || errorText === previewError) {
 			return undefined;
 		}
-		return `\n${theme.fg("error", errorText)}`;
+		return theme.fg("error", errorText);
 	}
 
 	const resultDiff = result.details?.diff;
-	if (!resultDiff) {
-		return undefined;
+	if (resultDiff && resultDiff !== previewDiff) {
+		return renderDiff(resultDiff, { filePath: rawPath ?? undefined });
 	}
-	return `\n${renderDiff(resultDiff, { filePath: rawPath ?? undefined })}`;
+
+	return undefined;
+}
+
+function getEditHeaderBg(
+	preview: EditPreview | undefined,
+	settledError: boolean | undefined,
+	theme: typeof import("../../modes/interactive/theme/theme.js").theme,
+): (text: string) => string {
+	if (preview) {
+		if ("error" in preview) {
+			return (text: string) => theme.bg("toolErrorBg", text);
+		}
+		return (text: string) => theme.bg("toolSuccessBg", text);
+	}
+	if (settledError) {
+		return (text: string) => theme.bg("toolErrorBg", text);
+	}
+	return (text: string) => theme.bg("toolPendingBg", text);
+}
+
+function buildEditCallComponent(
+	component: EditCallRenderComponent,
+	args: RenderableEditArgs | undefined,
+	theme: typeof import("../../modes/interactive/theme/theme.js").theme,
+): EditCallRenderComponent {
+	component.setBgFn(getEditHeaderBg(component.preview, component.settledError, theme));
+	component.clear();
+	component.addChild(new Text(formatEditCall(args, theme), 0, 0));
+
+	if (!component.preview) {
+		return component;
+	}
+
+	const body =
+		"error" in component.preview ? theme.fg("error", component.preview.error) : renderDiff(component.preview.diff);
+	component.addChild(new Spacer(1));
+	component.addChild(new Text(body, 0, 0));
+	return component;
+}
+
+function setEditPreview(
+	component: EditCallRenderComponent,
+	preview: EditPreview,
+	argsKey: string | undefined,
+): boolean {
+	const current = component.preview;
+	const changed =
+		current === undefined ||
+		("error" in current && "error" in preview
+			? current.error !== preview.error
+			: "error" in current !== "error" in preview) ||
+		(!("error" in current) &&
+			!("error" in preview) &&
+			(current.diff !== preview.diff || current.firstChangedLine !== preview.firstChangedLine));
+	component.preview = preview;
+	component.previewArgsKey = argsKey;
+	component.previewPending = false;
+	return changed;
 }
 
 export function createEditToolDefinition(
@@ -169,6 +294,7 @@ export function createEditToolDefinition(
 			"Keep edits[].oldText as small as possible while still being unique in the file. Do not pad with large unchanged regions.",
 		],
 		parameters: editSchema,
+		renderShell: "self",
 		prepareArguments: prepareEditArguments,
 		async execute(_toolCallId, input: EditToolInput, signal?: AbortSignal, _onUpdate?, _ctx?) {
 			const { path, edits } = validateEditInput(input);
@@ -280,20 +406,68 @@ export function createEditToolDefinition(
 			);
 		},
 		renderCall(args, theme, context) {
-			const text = (context.lastComponent as Text | undefined) ?? new Text("", 0, 0);
-			text.setText(formatEditCall(args, theme));
-			return text;
+			const component = getEditCallRenderComponent(context.state, context.lastComponent);
+			const previewInput = getRenderablePreviewInput(args as RenderableEditArgs | undefined);
+			const argsKey = previewInput
+				? JSON.stringify({ path: previewInput.path, edits: previewInput.edits })
+				: undefined;
+
+			if (component.previewArgsKey !== argsKey) {
+				component.preview = undefined;
+				component.previewArgsKey = argsKey;
+				component.previewPending = false;
+				component.settledError = false;
+			}
+
+			if (context.argsComplete && previewInput && !component.preview && !component.previewPending) {
+				component.previewPending = true;
+				const requestKey = argsKey;
+				void computeEditsDiff(previewInput.path, previewInput.edits, context.cwd).then((preview) => {
+					if (component.previewArgsKey === requestKey) {
+						setEditPreview(component, preview, requestKey);
+						context.invalidate();
+					}
+				});
+			}
+
+			return buildEditCallComponent(component, args, theme);
 		},
 		renderResult(result, _options, theme, context) {
-			const output = formatEditResult(context.args, result as any, theme, context.isError);
+			const callComponent = context.state.callComponent;
+			const previewInput = getRenderablePreviewInput(context.args as RenderableEditArgs | undefined);
+			const argsKey = previewInput
+				? JSON.stringify({ path: previewInput.path, edits: previewInput.edits })
+				: undefined;
+			const typedResult = result as EditToolResultLike;
+			const resultDiff = !context.isError ? typedResult.details?.diff : undefined;
+			let changed = false;
+			if (callComponent) {
+				if (typeof resultDiff === "string") {
+					changed =
+						setEditPreview(
+							callComponent,
+							{ diff: resultDiff, firstChangedLine: typedResult.details?.firstChangedLine },
+							argsKey,
+						) || changed;
+				}
+				if (callComponent.settledError !== context.isError) {
+					callComponent.settledError = context.isError;
+					changed = true;
+				}
+				if (changed) {
+					context.invalidate();
+				}
+			}
+
+			const output = formatEditResult(context.args, callComponent?.preview, typedResult, theme, context.isError);
+			const component = (context.lastComponent as Container | undefined) ?? new Container();
+			component.clear();
 			if (!output) {
-				const component = (context.lastComponent as Container | undefined) ?? new Container();
-				component.clear();
 				return component;
 			}
-			const text = (context.lastComponent as Text | undefined) ?? new Text("", 0, 0);
-			text.setText(output);
-			return text;
+			component.addChild(new Spacer(1));
+			component.addChild(new Text(output, 1, 0));
+			return component;
 		},
 	};
 }