Fix AI CLI interrupt handling and immediate turn recovery

apps/cli/ai/tests/ui.test.ts

@@ -206,6 +206,55 @@ describe( 'AiChatUI.clearTranscript', () => {
 	} );
 } );
 
+describe( 'AiChatUI interrupt handling', () => {
+	it( 'centralizes ESC interruption cleanup and only calls the interrupt callback once', () => {
+		const ui = Object.create( AiChatUI.prototype ) as {
+			requestInterrupt: () => boolean;
+			[ key: string ]: unknown;
+		};
+		const interruptCallback = vi.fn();
+		const submitResolve = vi.fn();
+
+		ui.interruptCallback = interruptCallback;
+		ui.wasInterrupted = false;
+		ui.closeSitePicker = vi.fn();
+		ui.cancelOptionPicker = vi.fn();
+		ui.showInterruptedNotice = vi.fn();
+		ui.submitResolve = submitResolve;
+		ui.updateHints = vi.fn();
+
+		expect( ui.requestInterrupt() ).toBe( true );
+		expect( ui.wasInterrupted ).toBe( true );
+		expect( ui.closeSitePicker ).toHaveBeenCalledTimes( 1 );
+		expect( ui.cancelOptionPicker ).toHaveBeenCalledTimes( 1 );
+		expect( submitResolve ).toHaveBeenCalledWith( '' );
+		expect( ui.showInterruptedNotice ).toHaveBeenCalledTimes( 1 );
+		expect( interruptCallback ).toHaveBeenCalledTimes( 1 );
+
+		expect( ui.requestInterrupt() ).toBe( true );
+		expect( interruptCallback ).toHaveBeenCalledTimes( 1 );
+		expect( ui.showInterruptedNotice ).toHaveBeenCalledTimes( 1 );
+	} );
+
+	it( 'does not advertise ESC as interrupt when no interrupt callback is active', () => {
+		const ui = Object.create( AiChatUI.prototype ) as {
+			updateHints: () => void;
+			[ key: string ]: unknown;
+		};
+		const editor = { hints: [] as string[] };
+
+		ui.editor = editor;
+		ui.interruptCallback = null;
+		ui._inAgentTurn = false;
+		ui.activeExpandablePreview = null;
+		ui.queuedPrompts = [];
+
+		ui.updateHints();
+
+		expect( editor.hints ).not.toContain( 'esc to interrupt' );
+	} );
+} );
+
 describe( 'AiChatUI.handleMessage', () => {
 	beforeEach( () => {
 		vi.clearAllMocks();

apps/cli/ai/ui.ts

@@ -386,6 +386,7 @@ export class AiChatUI implements AiOutputAdapter {
 	private editorVisible = false;
 	private interruptCallback: ( () => void ) | null = null;
 	private wasInterrupted = false;
+	private interruptionNoticeShown = false;
 	private usageCapReached = false;
 	private hasShownResponseMarker = false;
 	private turnStartTime = 0;
@@ -594,6 +595,9 @@ export class AiChatUI implements AiOutputAdapter {
 				this.stop();
 				process.exit( 0 );
 			}
+			if ( matchesKey( data, 'escape' ) && this.requestInterrupt() ) {
+				return { consume: true };
+			}
 			// Option picker navigation (must be checked before site picker)
 			if ( this.optionPickerSelectList ) {
 				// When "Other" is active, let the inline input handle most keys
@@ -704,10 +708,6 @@ export class AiChatUI implements AiOutputAdapter {
 				this.renderQueuedContainer();
 				return { consume: true };
 			}
-			if ( matchesKey( data, 'escape' ) && this.interruptCallback ) {
-				this.wasInterrupted = true;
-				this.interruptCallback();
-			}
 			if ( matchesKey( data, 'ctrl+o' ) && this.activeExpandablePreview ) {
 				this.toggleExpandablePreview();
 				return { consume: true };
@@ -1174,6 +1174,13 @@ export class AiChatUI implements AiOutputAdapter {
 		this.tui.requestRender();
 	}
 
+	private cancelOptionPicker(): void {
+		const resolve = this.optionPickerResolve;
+		this.optionPickerResolve = null;
+		this.closeOptionPicker();
+		resolve?.( '' );
+	}
+
 	start(): void {
 		this.tui.start();
 	}
@@ -1246,6 +1253,55 @@ export class AiChatUI implements AiOutputAdapter {
 
 	set onInterrupt( fn: ( () => void ) | null ) {
 		this.interruptCallback = fn;
+		this.updateHints();
+	}
+
+	private requestInterrupt(): boolean {
+		if ( ! this.interruptCallback ) {
+			return false;
+		}
+
+		if ( this.wasInterrupted ) {
+			return true;
+		}
+
+		this.wasInterrupted = true;
+		this.closeSitePicker();
+		this.cancelOptionPicker();
+		if ( this.submitResolve ) {
+			const resolve = this.submitResolve;
+			this.submitResolve = null;
+			resolve( '' );
+		}
+		this.showInterruptedNotice();
+		this.interruptCallback();
+		this.updateHints();
+		return true;
+	}
+
+	private showInterruptedNotice(): void {
+		if ( this.interruptionNoticeShown ) {
+			return;
+		}
+
+		this.interruptionNoticeShown = true;
+		this.hideLoader();
+		this.stopToolDotBlink();
+		this.toolDotText = null;
+		this.currentMarkdown = null;
+		this.currentResponseText = '';
+
+		const thinkingSec = Math.round( ( this.nowMs() - this.turnStartTime ) / 1000 );
+		this.messages.addChild(
+			new Text( '\n ' + chalk.yellow( '⏺' ) + ' ' + chalk.yellow( __( 'Interrupted' ) ), 0, 0 )
+		);
+		this.showInfo(
+			sprintf(
+				/* translators: %d: number of seconds */
+				__( 'Ran for %ds before interruption' ),
+				thinkingSec
+			)
+		);
 	}
 
 	stop(): void {
@@ -1353,7 +1409,9 @@ export class AiChatUI implements AiOutputAdapter {
 		if ( this.queuedPrompts.length > 0 ) {
 			hints.push( __( 'backspace to unqueue' ) );
 		}
-		hints.push( __( 'esc to interrupt' ) );
+		if ( this.interruptCallback ) {
+			hints.push( __( 'esc to interrupt' ) );
+		}
 		this.editor.hints = hints;
 	}
 
@@ -1387,6 +1445,7 @@ export class AiChatUI implements AiOutputAdapter {
 		this.currentResponseText = '';
 		this.hasShownResponseMarker = false;
 		this.wasInterrupted = false;
+		this.interruptionNoticeShown = false;
 		this.usageCapReached = false;
 		this.turnStartTime = this.nowMs();
 		this.todoSnapshot = [];
@@ -2005,12 +2064,21 @@ export class AiChatUI implements AiOutputAdapter {
 						this.closeOptionPicker();
 						resolve( item.value );
 					};
+					selectList.onCancel = () => {
+						this.cancelOptionPicker();
+					};
 				} );
 
+				if ( ! selected ) {
+					return answers;
+				}
 				answers[ q.question ] = selected;
 			} else {
 				// Free-form text input
 				const answer = await this.waitForInput();
+				if ( ! answer ) {
+					return answers;
+				}
 				answers[ q.question ] = answer;
 			}
 		}
@@ -2025,6 +2093,10 @@ export class AiChatUI implements AiOutputAdapter {
 	 * Returns session result when the agent turn is complete.
 	 */
 	handleMessage( message: SDKMessage ): HandleMessageResult | undefined {
+		if ( this.wasInterrupted && message.type !== 'result' ) {
+			return undefined;
+		}
+
 		switch ( message.type ) {
 			case 'assistant': {
 				// Detect the AI usage cap response from the WordPress.com proxy.
@@ -2168,21 +2240,7 @@ export class AiChatUI implements AiOutputAdapter {
 
 				// User-initiated interruption: friendly message, suppress retry prompt.
 				if ( this.wasInterrupted ) {
-					const thinkingSec = Math.round( ( this.nowMs() - this.turnStartTime ) / 1000 );
-					this.messages.addChild(
-						new Text(
-							'\n ' + chalk.yellow( '⏺' ) + ' ' + chalk.yellow( __( 'Interrupted' ) ),
-							0,
-							0
-						)
-					);
-					this.showInfo(
-						sprintf(
-							/* translators: %d: number of seconds */
-							__( 'Ran for %ds before interruption' ),
-							thinkingSec
-						)
-					);
+					this.showInterruptedNotice();
 					return {
 						type: 'result',
 						sessionId: message.session_id,

apps/cli/commands/ai/index.ts

@@ -461,16 +461,29 @@ export async function runCommand( options: {
 			onAskUser: ( questions ) => askUserAndPersistAnswers( questions ),
 		} );
 
+		let interruptRequested = false;
+		let resolveInterrupt: () => void = () => undefined;
+		const interruptPromise = new Promise< 'interrupted' >( ( resolve ) => {
+			resolveInterrupt = () => resolve( 'interrupted' );
+		} );
 		ui.onInterrupt = () => {
+			if ( interruptRequested ) {
+				return;
+			}
+			interruptRequested = true;
 			void agentQuery.interrupt();
+			resolveInterrupt();
 		};
 
 		let maxTurnsResult: { numTurns: number } | undefined;
-		let turnStatus: TurnStatus = 'interrupted';
+		const turnState: { status: TurnStatus } = { status: 'interrupted' };
 
-		try {
+		const consumeAgentTurn = async (): Promise< void > => {
 			for await ( const message of agentQuery ) {
 				const timestamp = new Date().toISOString();
+				if ( interruptRequested ) {
+					continue;
+				}
 				const result = ui.handleMessage( message );
 				await persist( ( recorder ) => recorder.recordSdkMessage( message, timestamp ) );
 				if ( result ) {
@@ -481,30 +494,45 @@ export async function runCommand( options: {
 						maxTurnsResult = {
 							numTurns: result.numTurns,
 						};
-						turnStatus = 'max_turns';
+						turnState.status = 'max_turns';
 					} else if ( result.interrupted ) {
-						turnStatus = 'interrupted';
+						turnState.status = 'interrupted';
 					} else {
-						turnStatus = result.success ? 'success' : 'error';
+						turnState.status = result.success ? 'success' : 'error';
 					}
 				}
 			}
-		} catch ( error ) {
-			turnStatus = 'error';
-			// In JSON mode there's no interactive retry, so re-throw and let
-			// the caller record the error. In interactive mode, fall through
-			// so the post-loop retry prompt offers the user a chance to retry.
-			if ( isJsonMode ) {
-				throw error;
+		};
+
+		const consumeAgentTurnResult = consumeAgentTurn().catch( ( error ) => {
+			if ( interruptRequested ) {
+				turnState.status = 'interrupted';
+				return;
 			}
+			turnState.status = 'error';
 			// If the UI already surfaced a descriptive terminal error (e.g.
 			// the AI usage cap was reached), suppress the generic SDK exit
 			// error (e.g. "Claude Code process exited with code 1").
 			if ( ! ( ui instanceof AiChatUI && ui.hasErrorBeenSurfaced() ) ) {
 				ui.showError( getErrorMessage( error ) );
 			}
+			// In JSON mode there's no interactive retry, so re-throw and let
+			// the caller record the error.
+			if ( isJsonMode ) {
+				throw error;
+			}
+		} );
+
+		try {
+			const result = await Promise.race( [
+				consumeAgentTurnResult.then( () => 'completed' as const ),
+				interruptPromise,
+			] );
+			if ( result === 'interrupted' ) {
+				turnState.status = 'interrupted';
+			}
 		} finally {
-			await persist( ( recorder ) => recorder.recordTurnClosed( turnStatus ) );
+			await persist( ( recorder ) => recorder.recordTurnClosed( turnState.status ) );
 			ui.endAgentTurn();
 		}
 
@@ -537,7 +565,7 @@ export async function runCommand( options: {
 		// the user has already been told what to do next.
 		const hasTerminalError = ui instanceof AiChatUI && ui.hasErrorBeenSurfaced();
 
-		if ( turnStatus === 'error' && ! isJsonMode && ! hasTerminalError ) {
+		if ( turnState.status === 'error' && ! isJsonMode && ! hasTerminalError ) {
 			if ( retryAttempt >= MAX_RETRY_ATTEMPTS ) {
 				ui.showInfo(
 					__( 'The server has not recovered after multiple attempts. Please try again later.' )
@@ -568,7 +596,7 @@ export async function runCommand( options: {
 		}
 
 		return {
-			status: turnStatus,
+			status: turnState.status,
 			usage: maxTurnsResult,
 		};
 	}