From d57a0ff233857a00c027dd11692b38b33d921bcd Mon Sep 17 00:00:00 2001
From: can1357 <me@can.ac>
Date: Sun, 16 Aug 2026 02:37:46 +0200
Subject: [PATCH] fix(extensions): preserved drafts before custom presentation

---
 .../src/modes/controllers/extension-ui-controller.ts          | 4 +++-
 .../test/modes/controllers/extension-ui-controller.test.ts    | 3 +++
 2 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/packages/coding-agent/src/modes/controllers/extension-ui-controller.ts b/packages/coding-agent/src/modes/controllers/extension-ui-controller.ts
index 727fa9404..fce13e198 100644
--- a/packages/coding-agent/src/modes/controllers/extension-ui-controller.ts
+++ b/packages/coding-agent/src/modes/controllers/extension-ui-controller.ts
@@ -1057,12 +1057,13 @@ export class ExtensionUiController {
 		let component: (Component & { dispose?(): void }) | undefined;
 		let overlayHandle: OverlayHandle | undefined;
 		let closed = false;
+		let editorReplaced = false;
 
 		const cleanup = () => {
 			component?.dispose?.();
 			overlayHandle?.hide();
 			overlayHandle = undefined;
-			if (!options?.overlay) {
+			if (editorReplaced) {
 				this.ctx.editorContainer.clear();
 				this.ctx.editorContainer.addChild(this.ctx.editor);
 				this.ctx.editor.setText(savedText);
@@ -1112,6 +1113,7 @@ export class ExtensionUiController {
 					options.onHandle?.(overlayHandle);
 					return;
 				}
+				editorReplaced = true;
 				this.ctx.editorContainer.clear();
 				this.ctx.editorContainer.addChild(component);
 				this.ctx.ui.setFocus(component);
diff --git a/packages/coding-agent/test/modes/controllers/extension-ui-controller.test.ts b/packages/coding-agent/test/modes/controllers/extension-ui-controller.test.ts
index 116d956cc..7f3aea80a 100644
--- a/packages/coding-agent/test/modes/controllers/extension-ui-controller.test.ts
+++ b/packages/coding-agent/test/modes/controllers/extension-ui-controller.test.ts
@@ -331,12 +331,14 @@ describe("ExtensionUiController custom overlay", () => {
 	it("aborts a pending custom factory and disposes its late component", async () => {
 		const harness = makeHarness();
 		const ui = await harness.init();
+		harness.editor.setText("draft before factory");
 		const controller = new AbortController();
 		const factory = Promise.withResolvers<Container>();
 		const component = new Container() as Container & { dispose: Mock<() => void> };
 		component.dispose = vi.fn();
 
 		const pending = ui.custom(() => factory.promise, { signal: controller.signal });
+		harness.editor.setText("draft typed while factory is pending");
 		controller.abort();
 
 		await expect(pending).rejects.toBe(controller.signal.reason);
@@ -345,5 +347,6 @@ describe("ExtensionUiController custom overlay", () => {
 
 		expect(component.dispose).toHaveBeenCalledTimes(1);
 		expect(harness.editorContainer.children).toEqual([harness.editor]);
+		expect(harness.editor.getText()).toBe("draft typed while factory is pending");
 	});
 });