Merge PR #6738: fix(tui): preserve prompt input while ask opens (@roboomp)

packages/coding-agent/CHANGELOG.md

@@ -1,6 +1,10 @@
 # Changelog
 
 ## [Unreleased]
+### Fixed
+
+- Fixed an Ask form appearing while the main prompt contains a draft hiding that text and consuming the next in-flight keystroke. The draft now remains visible and keeps receiving input until it is submitted or cleared; only then do form controls activate ([#6737](https://github.com/can1357/oh-my-pi/issues/6737)).
+
 
 ### Fixed
 

packages/coding-agent/src/modes/components/ask-dialog.ts

@@ -79,10 +79,21 @@ interface AskDialogCallbacks {
 	onPrompt(title: string, prefill?: string): Promise<string | undefined>;
 }
 
+interface AskDialogInputGuard {
+	isBlocked(): boolean;
+	handleInput(keyData: string): void;
+	hint: string;
+	/** Mirror the guard's blocked state onto the proxied draft surface each
+	 *  render, so a draft that owns input shows a visible insertion cursor even
+	 *  though this dialog holds TUI focus. */
+	syncPresentation?(): void;
+}
+
 interface AskDialogOptions {
 	timeout?: number;
 	onTimeout?: () => void;
 	tui?: TUI;
+	inputGuard?: AskDialogInputGuard;
 }
 
 interface QuestionState {
@@ -394,6 +405,12 @@ export class AskDialogComponent implements Component {
 			this.#finishCancel();
 			return;
 		}
+		const inputGuard = this.options.inputGuard;
+		if (inputGuard?.isBlocked()) {
+			inputGuard.handleInput(keyData);
+			this.#requestRender();
+			return;
+		}
 		if (this.#hasSubmitTab() && handleTabSwitchKey(keyData, direction => this.#switchTab(direction))) {
 			this.#requestRender();
 			return;
@@ -406,6 +423,10 @@ export class AskDialogComponent implements Component {
 	}
 
 	render(width: number): readonly string[] {
+		// Keep the proxied draft's cursor visible while it owns input (the editor
+		// renders as the next sibling in the same container, so this lands in the
+		// same frame).
+		this.options.inputGuard?.syncPresentation?.();
 		const innerWidth = Math.max(1, width - 4);
 		// Fixed panel height: measured from the tallest tab at spawn and
 		// re-measured only when the viewport changes. Tab switches, cursor
@@ -532,6 +553,8 @@ export class AskDialogComponent implements Component {
 
 	#footerHintText(indicator: string): string {
 		const cancel = `${cancelKeyLabel()} cancel`;
+		const inputGuard = this.options.inputGuard;
+		if (inputGuard?.isBlocked()) return `${inputGuard.hint} · ${cancel}`;
 		if (this.#isSubmitTab()) {
 			const scroll = indicator ? ` ${indicator} scroll ·` : "";
 			return `Enter submit · ↑/↓ scroll ·${scroll} ${cancel}`;

packages/coding-agent/src/modes/components/custom-editor.ts

@@ -974,4 +974,31 @@ export class CustomEditor extends Editor {
 			this.insertText("\n");
 		}
 	}
+
+	/**
+	 * Route a keystroke through the base text-editor pipeline only, skipping the
+	 * app-level shortcut interception in {@link handleInput} (Agent Hub, model
+	 * selector, history search, external editor, …). Used when the editor is
+	 * mounted for draft editing beneath another focused surface — e.g. an Ask
+	 * dialog opened over a non-empty prompt — so finishing or submitting the
+	 * draft can never fire an editor-slot shortcut that clears `editorContainer`
+	 * and orphans the overlay. Only text editing, cursor movement, submission,
+	 * and the clear action reach the buffer.
+	 */
+	handleDraftEdit(data: string): void {
+		// The base editor reserves Ctrl+C for parent handling and returns without
+		// touching the buffer, so the configured clear action must be dispatched
+		// explicitly here — otherwise the guard's "finish or clear the prompt"
+		// instruction has no working clear key. onClear (Ctrl+C → handleCtrlC)
+		// clears the draft on first press without swapping the editor slot; a
+		// standalone editor with no callback clears its own text.
+		const parsed = parseKey(data);
+		const canonical = parsed !== undefined ? canonicalKeyId(parsed) : undefined;
+		if (canonical !== undefined && this.#matchesAction(canonical, "app.clear")) {
+			if (this.onClear) this.onClear();
+			else this.setText("");
+			return;
+		}
+		super.handleInput(data);
+	}
 }

packages/coding-agent/src/modes/controllers/extension-ui-controller.test.ts

@@ -1,20 +1,39 @@
-import { describe, expect, it, vi } from "bun:test";
-import type { ExtensionUIContext } from "../../extensibility/extensions";
+import { afterEach, beforeAll, describe, expect, it, vi } from "bun:test";
+import { Container, setKeybindings } from "@oh-my-pi/pi-tui";
+import { KeybindingsManager } from "../../config/keybindings";
+import type { ExtensionAskDialogQuestion, ExtensionUIContext } from "../../extensibility/extensions";
+import { AskDialogComponent } from "../components/ask-dialog";
 import { CustomEditor } from "../components/custom-editor";
-import { getEditorTheme } from "../theme/theme";
+import { getEditorTheme, getThemeByName, setThemeInstance } from "../theme/theme";
 import type { InteractiveModeContext } from "../types";
 import { ExtensionUiController } from "./extension-ui-controller";
 
+afterEach(() => {
+	setKeybindings(KeybindingsManager.inMemory());
+});
+
+beforeAll(async () => {
+	const dark = await getThemeByName("dark");
+	if (!dark) throw new Error("Failed to load dark theme");
+	setThemeInstance(dark);
+});
+
 function makeHarness() {
 	const editor = new CustomEditor(getEditorTheme());
+	const editorContainer = new Container();
+	editorContainer.addChild(editor);
 	const requestRender = vi.fn();
+	const setFocus = vi.fn();
 	const addAutocompleteProvider = vi.fn();
 	let uiContext: ExtensionUIContext | undefined;
 	const ctx = {
 		editor,
 		ui: {
 			requestRender,
+			setFocus,
+			terminal: { rows: 40 },
 		},
+		editorContainer,
 		session: {
 			extensionRunner: undefined,
 			setUsageFallbackConfirmer: vi.fn(),
@@ -26,12 +45,17 @@ function makeHarness() {
 		addAutocompleteProvider,
 	} as unknown as InteractiveModeContext;
 
+	const controller = new ExtensionUiController(ctx);
+
 	return {
 		editor,
 		requestRender,
 		addAutocompleteProvider,
+		editorContainer,
+		setFocus,
+		controller,
 		async init(): Promise<ExtensionUIContext> {
-			await new ExtensionUiController(ctx).initHooksAndCustomTools();
+			await controller.initHooksAndCustomTools();
 			expect(uiContext).toBeDefined();
 			return uiContext!;
 		},
@@ -60,6 +84,123 @@ describe("ExtensionUiController editor UI", () => {
 		expect(harness.requestRender).toHaveBeenCalledTimes(1);
 	});
 
+	it("keeps a populated prompt visible and routes input to it until the draft is cleared", async () => {
+		const harness = makeHarness();
+		harness.editor.setText("finish this wor");
+		const questions: ExtensionAskDialogQuestion[] = [
+			{ id: "confirm", question: "Continue?", options: [{ label: "Yes" }, { label: "No" }] },
+		];
+
+		const pending = harness.controller.showAskDialog(questions);
+		const ask = harness.editorContainer.children[0];
+		expect(ask).toBeInstanceOf(AskDialogComponent);
+		expect(harness.editorContainer.children).toEqual([ask, harness.editor]);
+
+		ask?.handleInput?.("d");
+		expect(harness.editor.getText()).toBe("finish this word");
+
+		harness.editor.setText("");
+		ask?.handleInput?.("\n");
+		expect(await pending).toEqual({
+			kind: "submit",
+			results: [
+				{
+					id: "confirm",
+					question: "Continue?",
+					options: ["Yes", "No"],
+					multi: false,
+					selectedOptions: ["Yes"],
+					customInput: undefined,
+					note: undefined,
+					timedOut: undefined,
+				},
+			],
+		});
+		expect(harness.editorContainer.children).toEqual([harness.editor]);
+	});
+
+	it("does not fire editor-slot shortcuts that would orphan the ask dialog (#6738)", () => {
+		const harness = makeHarness();
+		harness.editor.setText("draft in progress");
+		// Simulate an editor-slot shortcut like the Agent Hub binding, whose
+		// handler clears editorContainer and would strand the pending ask.
+		let hubOpened = false;
+		harness.editor.setCustomKeyHandler("ctrl+s", () => {
+			hubOpened = true;
+			harness.editorContainer.clear();
+		});
+		const questions: ExtensionAskDialogQuestion[] = [
+			{ id: "confirm", question: "Continue?", options: [{ label: "Yes" }, { label: "No" }] },
+		];
+
+		harness.controller.showAskDialog(questions);
+		const ask = harness.editorContainer.children[0];
+		expect(ask).toBeInstanceOf(AskDialogComponent);
+
+		// Ctrl+S reaches the draft editor while ask is open; the shortcut must be
+		// swallowed, the draft untouched, and the ask surface preserved.
+		ask?.handleInput?.("\x13");
+		expect(hubOpened).toBe(false);
+		expect(harness.editor.getText()).toBe("draft in progress");
+		expect(harness.editorContainer.children).toEqual([ask, harness.editor]);
+	});
+
+	it("exposes the draft editor cursor while it proxies input, and drops it once cleared (#6738)", () => {
+		const harness = makeHarness();
+		harness.editor.setText("finish this wor");
+		const questions: ExtensionAskDialogQuestion[] = [
+			{ id: "confirm", question: "Continue?", options: [{ label: "Yes" }, { label: "No" }] },
+		];
+
+		harness.controller.showAskDialog(questions);
+		const ask = harness.editorContainer.children[0];
+		expect(ask).toBeInstanceOf(AskDialogComponent);
+
+		// The ask dialog holds TUI focus, but rendering it must mirror focus onto
+		// the draft editor so its insertion cursor is visible.
+		ask?.render?.(80);
+		expect(harness.editor.focused).toBe(true);
+
+		// Once the draft clears, the ask controls take over and the editor cursor
+		// must not linger.
+		harness.editor.setText("");
+		ask?.render?.(80);
+		expect(harness.editor.focused).toBe(false);
+	});
+
+	it("lets the clear action empty the draft and lift the ask guard (#6738)", () => {
+		const harness = makeHarness();
+		// Route Ctrl+C to the guard: keep app.clear on Ctrl+C but move the ask
+		// cancel key off it, so Ctrl+C reaches draft editing instead of cancelling.
+		setKeybindings(KeybindingsManager.inMemory({ "tui.select.cancel": "ctrl+g" }));
+		harness.editor.setActionKeys("app.clear", ["ctrl+c"]);
+		let cleared = 0;
+		// Mirror interactive wiring: app.clear (Ctrl+C) clears the draft.
+		harness.editor.onClear = () => {
+			cleared++;
+			harness.editor.setText("");
+		};
+		harness.editor.setText("half typed prompt");
+		const questions: ExtensionAskDialogQuestion[] = [
+			{ id: "confirm", question: "Continue?", options: [{ label: "Yes" }, { label: "No" }] },
+		];
+
+		harness.controller.showAskDialog(questions);
+		const ask = harness.editorContainer.children[0];
+		expect(ask).toBeInstanceOf(AskDialogComponent);
+
+		// Ctrl+C is reserved by the base editor and never clears; the guard must
+		// dispatch the configured clear action so the "finish or clear" hint works.
+		ask?.handleInput?.("\x03");
+		expect(cleared).toBe(1);
+		expect(harness.editor.getText()).toBe("");
+
+		// With the draft gone the guard releases: the next key reaches the ask
+		// controls and submits the highlighted option.
+		ask?.handleInput?.("\n");
+		expect(harness.editorContainer.children).toEqual([harness.editor]);
+	});
+
 	it("bridges addAutocompleteProvider factories to the interactive mode context (#4919)", async () => {
 		const harness = makeHarness();
 		const ui = await harness.init();

packages/coding-agent/src/modes/controllers/extension-ui-controller.ts

@@ -623,6 +623,20 @@ export class ExtensionUiController {
 			let promptEditor: HookEditorComponent | undefined;
 			let promptResolve: ((value: string | undefined) => void) | undefined;
 			let closed = false;
+			const draftEditor = this.ctx.editor;
+			const inputGuard =
+				draftEditor.getText().length > 0
+					? {
+							isBlocked: () => draftEditor.getText().length > 0,
+							handleInput: (keyData: string) => draftEditor.handleDraftEdit(keyData),
+							hint: "Finish or clear the current prompt to answer",
+							// Show the draft's insertion cursor while it owns input; drop it
+							// once the draft clears and the ask controls take over.
+							syncPresentation: () => {
+								draftEditor.focused = draftEditor.getText().length > 0;
+							},
+						}
+					: undefined;
 
 			const restoreAskDialog = (): void => {
 				if (closed || !askDialog) return;
@@ -670,10 +684,12 @@ export class ExtensionUiController {
 					timeout: dialogOptions?.timeout,
 					onTimeout: dialogOptions?.onTimeout,
 					tui: this.ctx.ui,
+					inputGuard,
 				},
 			);
 			this.ctx.editorContainer.clear();
 			this.ctx.editorContainer.addChild(askDialog);
+			if (inputGuard) this.ctx.editorContainer.addChild(this.ctx.editor);
 			this.ctx.ui.setFocus(askDialog);
 			this.ctx.ui.requestRender();