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

packages/coding-agent/docs/extensions.md

@@ -1763,7 +1763,25 @@ export default function (pi: ExtensionAPI) {
 
 Tools can provide `renderCall` and `renderResult` for custom TUI display. See [tui.md](tui.md) for the full component API and [tool-execution.ts](https://github.com/badlogic/pi-mono/blob/main/packages/coding-agent/src/modes/interactive/components/tool-execution.ts) for how tool rows are composed.
 
-Tool output is wrapped in a `Box` that handles padding and background. A defined `renderCall` or `renderResult` must return a `Component`. If a slot renderer is not defined, `tool-execution.ts` uses fallback rendering for that slot.
+By default, tool output is wrapped in a `Box` that handles padding and background. A defined `renderCall` or `renderResult` must return a `Component`. If a slot renderer is not defined, `tool-execution.ts` uses fallback rendering for that slot.
+
+Set `renderShell: "self"` when the tool should render its own shell instead of using the default `Box`. This is useful for tools that need complete control over framing or background behavior, for example large previews that must stay visually stable after the tool settles.
+
+```typescript
+pi.registerTool({
+  name: "my_tool",
+  label: "My Tool",
+  description: "Custom shell example",
+  parameters: Type.Object({}),
+  renderShell: "self",
+  async execute() {
+    return { content: [{ type: "text", text: "ok" }], details: undefined };
+  },
+  renderCall(args, theme, context) {
+    return new Text(theme.fg("accent", "my custom shell"), 0, 0);
+  },
+});
+```
 
 `renderCall` and `renderResult` each receive a `context` object with:
 - `args` - the current tool call arguments
@@ -1850,7 +1868,7 @@ Custom editors and `ctx.ui.custom()` components receive `keybindings: Keybinding
 
 #### Best Practices
 
-- Use `Text` with padding `(0, 0)`. The Box handles padding.
+- Use `Text` with padding `(0, 0)`. The default Box handles padding.
 - Use `\n` for multi-line content.
 - Handle `isPartial` for streaming progress.
 - Support `expanded` for detail on demand.
@@ -1858,6 +1876,7 @@ Custom editors and `ctx.ui.custom()` components receive `keybindings: Keybinding
 - Read `context.args` in `renderResult` instead of copying args into `context.state`.
 - Use `context.state` only for data that must be shared across call and result slots.
 - Reuse `context.lastComponent` when the same component instance can be updated in place.
+- Use `renderShell: "self"` only when the default boxed shell gets in the way. In self-shell mode the tool is responsible for its own framing, padding, and background.
 
 #### Fallback
 

packages/coding-agent/examples/extensions/built-in-tool-renderer.ts

@@ -16,6 +16,8 @@
  *   and delegate execute() to them
  * - renderCall() controls what's shown when the tool is invoked
  * - renderResult() controls what's shown after execution completes
+ * - renderShell: "self" lets a tool render its own outer shell instead of
+ *   using the default boxed shell from ToolExecutionComponent
  * - The `expanded` flag in renderResult indicates whether the user has
  *   toggled the tool output open (via ctrl+e or clicking)
  *
@@ -155,6 +157,7 @@ export default function (pi: ExtensionAPI) {
 		label: "edit",
 		description: originalEdit.description,
 		parameters: originalEdit.parameters,
+		renderShell: "self",
 
 		async execute(toolCallId, params, signal, onUpdate) {
 			return originalEdit.execute(toolCallId, params, signal, onUpdate);

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;
 		},
 	};
 }

packages/coding-agent/src/modes/interactive/components/tool-execution.ts

@@ -12,6 +12,7 @@ export interface ToolExecutionOptions {
 export class ToolExecutionComponent extends Container {
 	private contentBox: Box;
 	private contentText: Text;
+	private selfRenderContainer: Container;
 	private callRendererComponent?: Component;
 	private resultRendererComponent?: Component;
 	private rendererState: any = {};
@@ -58,13 +59,15 @@ export class ToolExecutionComponent extends Container {
 
 		this.addChild(new Spacer(1));
 
-		// Always create both. contentBox is used for tools with renderer-based call/result composition.
+		// Always create all shell variants. contentBox is used for default renderer-based composition.
+		// selfRenderContainer is used when the tool renders its own framing.
 		// contentText is reserved for generic fallback rendering when no tool definition exists.
 		this.contentBox = new Box(1, 1, (text: string) => theme.bg("toolPendingBg", text));
 		this.contentText = new Text("", 1, 1, (text: string) => theme.bg("toolPendingBg", text));
+		this.selfRenderContainer = new Container();
 
 		if (this.hasRendererDefinition()) {
-			this.addChild(this.contentBox);
+			this.addChild(this.getRenderShell() === "self" ? this.selfRenderContainer : this.contentBox);
 		} else {
 			this.addChild(this.contentText);
 		}
@@ -96,6 +99,16 @@ export class ToolExecutionComponent extends Container {
 		return this.builtInToolDefinition !== undefined || this.toolDefinition !== undefined;
 	}
 
+	private getRenderShell(): "default" | "self" {
+		if (!this.builtInToolDefinition) {
+			return this.toolDefinition?.renderShell ?? "default";
+		}
+		if (!this.toolDefinition) {
+			return this.builtInToolDefinition.renderShell ?? "default";
+		}
+		return this.toolDefinition.renderShell ?? this.builtInToolDefinition.renderShell ?? "default";
+	}
+
 	private getRenderContext(lastComponent: Component | undefined): ToolRenderContext {
 		return {
 			args: this.args,
@@ -214,22 +227,25 @@ export class ToolExecutionComponent extends Container {
 		let hasContent = false;
 		this.hideComponent = false;
 		if (this.hasRendererDefinition()) {
-			this.contentBox.setBgFn(bgFn);
-			this.contentBox.clear();
+			const renderContainer = this.getRenderShell() === "self" ? this.selfRenderContainer : this.contentBox;
+			if (renderContainer instanceof Box) {
+				renderContainer.setBgFn(bgFn);
+			}
+			renderContainer.clear();
 
 			const callRenderer = this.getCallRenderer();
 			if (!callRenderer) {
-				this.contentBox.addChild(this.createCallFallback());
+				renderContainer.addChild(this.createCallFallback());
 				hasContent = true;
 			} else {
 				try {
 					const component = callRenderer(this.args, theme, this.getRenderContext(this.callRendererComponent));
 					this.callRendererComponent = component;
-					this.contentBox.addChild(component);
+					renderContainer.addChild(component);
 					hasContent = true;
 				} catch {
 					this.callRendererComponent = undefined;
-					this.contentBox.addChild(this.createCallFallback());
+					renderContainer.addChild(this.createCallFallback());
 					hasContent = true;
 				}
 			}
@@ -239,7 +255,7 @@ export class ToolExecutionComponent extends Container {
 				if (!resultRenderer) {
 					const component = this.createResultFallback();
 					if (component) {
-						this.contentBox.addChild(component);
+						renderContainer.addChild(component);
 						hasContent = true;
 					}
 				} else {
@@ -251,13 +267,13 @@ export class ToolExecutionComponent extends Container {
 							this.getRenderContext(this.resultRendererComponent),
 						);
 						this.resultRendererComponent = component;
-						this.contentBox.addChild(component);
+						renderContainer.addChild(component);
 						hasContent = true;
 					} catch {
 						this.resultRendererComponent = undefined;
 						const component = this.createResultFallback();
 						if (component) {
-							this.contentBox.addChild(component);
+							renderContainer.addChild(component);
 							hasContent = true;
 						}
 					}

packages/coding-agent/test/edit-tool-no-full-redraw.test.ts

@@ -56,7 +56,7 @@ describe("edit tool TUI rendering", () => {
 		await Promise.all(tempDirs.splice(0).map((dir) => rm(dir, { recursive: true, force: true })));
 	});
 
-	it("renders the large diff only in the settled result without triggering a full TUI redraw", async () => {
+	it("renders the large diff in the call preview and does not full-redraw when the result settles", async () => {
 		const dir = await mkdtemp(join(tmpdir(), "pi-edit-redraw-"));
 		tempDirs.push(dir);
 		const filePath = join(dir, "large-edit.txt");
@@ -97,11 +97,12 @@ describe("edit tool TUI rendering", () => {
 		component.setArgsComplete();
 		tui.requestRender();
 		await waitForRender();
+		await waitForRender();
 
 		const callOnlyRender = component.render(80).join("\n");
 		expect(callOnlyRender).toContain("edit");
-		expect(callOnlyRender).not.toContain("line 50 changed");
-		expect(callOnlyRender).not.toContain("+  51");
+		expect(callOnlyRender).toContain("line 50 changed");
+		expect(callOnlyRender).toContain("line 950 changed");
 
 		const redrawsBeforeResult = tui.fullRedraws;
 		const clearsBeforeResult = terminal.fullClearCount;
@@ -122,5 +123,87 @@ describe("edit tool TUI rendering", () => {
 		const settledRender = component.render(80).join("\n");
 		expect(settledRender).toContain("line 50 changed");
 		expect(settledRender).toContain("line 950 changed");
+		expect(settledRender).not.toContain("Successfully replaced");
+	});
+
+	it("reconstructs the boxed preview from a settled result without argsComplete", async () => {
+		const dir = await mkdtemp(join(tmpdir(), "pi-edit-replay-"));
+		tempDirs.push(dir);
+		const filePath = join(dir, "replay-edit.txt");
+		await writeFile(
+			filePath,
+			`${Array.from({ length: 200 }, (_, i) => `line ${i}`).join("\n")}
+`,
+			"utf8",
+		);
+		const lines = (await readFile(filePath, "utf8")).trimEnd().split("\n");
+		const edits = createLargeEdits(lines).slice(0, 2);
+		const diff = await computeEditsDiff(filePath, edits, process.cwd());
+		if ("error" in diff) {
+			throw new Error(diff.error);
+		}
+		await rm(filePath, { force: true });
+
+		const terminal = new FakeTerminal();
+		const tui = new TUI(terminal);
+		const component = new ToolExecutionComponent(
+			"edit",
+			"tool-call-replay",
+			{ path: filePath, edits },
+			{},
+			createEditToolDefinition(process.cwd()),
+			tui,
+			process.cwd(),
+		);
+		tui.addChild(component);
+		tui.start();
+		await waitForRender();
+
+		component.updateResult(
+			{
+				content: [{ type: "text", text: `Successfully replaced ${edits.length} block(s) in ${filePath}.` }],
+				details: diff,
+				isError: false,
+			},
+			false,
+		);
+		await waitForRender();
+		await waitForRender();
+
+		const rendered = component.render(80).join("\n");
+		expect(rendered).toContain("line 50 changed");
+		expect(rendered).toContain("line 150 changed");
+	});
+
+	it("shows a preflight error without rendering a diff when the edits do not apply", async () => {
+		const dir = await mkdtemp(join(tmpdir(), "pi-edit-preflight-"));
+		tempDirs.push(dir);
+		const filePath = join(dir, "missing-edit.txt");
+		await writeFile(filePath, "line 0\nline 1\n", "utf8");
+
+		const terminal = new FakeTerminal();
+		const tui = new TUI(terminal);
+		const component = new ToolExecutionComponent(
+			"edit",
+			"tool-call-2",
+			{ path: filePath, edits: [{ oldText: "does not exist", newText: "replacement" }] },
+			{},
+			createEditToolDefinition(process.cwd()),
+			tui,
+			process.cwd(),
+		);
+		tui.addChild(component);
+		tui.start();
+		await waitForRender();
+
+		component.setArgsComplete();
+		tui.requestRender();
+		await waitForRender();
+		await waitForRender();
+
+		const rendered = component.render(80).join("\n");
+		expect(rendered).toContain("Could not find");
+		expect(rendered).not.toContain("+1 ");
+		expect(rendered).not.toContain("-1 ");
 	});
 });