Allow editing the previous user message

apps/studio/src/components/studio-code-session/conversation/index.test.tsx

@@ -1,6 +1,6 @@
 import { render, screen } from '@testing-library/react';
 import { beforeEach, describe, expect, it, vi } from 'vitest';
-import { Conversation, entriesToRenderItems } from './index';
+import { Conversation, entriesToRenderItems, isConversationStopped } from './index';
 import type { SessionEntry } from '@earendil-works/pi-coding-agent';
 import type { LoadedAiSession } from '@studio/common/ai/sessions/types';
 
@@ -46,6 +46,14 @@ function answer( text: string ): SessionEntry {
 	return customEntry( 'studio.user_prompt', { text, source: 'ask_user' } );
 }
 
+function prompt( text: string ): SessionEntry {
+	return customEntry( 'studio.user_prompt', { text, source: 'prompt' } );
+}
+
+function turnClosed( status: 'completed' | 'interrupted' ): SessionEntry {
+	return customEntry( 'studio.turn_closed', { status } );
+}
+
 function assistantToolCallEntry( name: string ): SessionEntry {
 	return {
 		type: 'message',
@@ -114,6 +122,31 @@ describe( 'entriesToRenderItems – persisted picked answers', () => {
 	} );
 } );
 
+describe( 'isConversationStopped', () => {
+	it( 'is false for an in-flight or completed turn', () => {
+		expect( isConversationStopped( [ prompt( 'Build me a blog' ) ] ) ).toBe( false );
+		expect(
+			isConversationStopped( [ prompt( 'Build me a blog' ), turnClosed( 'completed' ) ] )
+		).toBe( false );
+	} );
+
+	it( 'is true once the latest turn was interrupted', () => {
+		expect(
+			isConversationStopped( [ prompt( 'Build me a blog' ), turnClosed( 'interrupted' ) ] )
+		).toBe( true );
+	} );
+
+	it( 'is false again once a newer turn has started', () => {
+		expect(
+			isConversationStopped( [
+				prompt( 'Build me a blog' ),
+				turnClosed( 'interrupted' ),
+				prompt( 'Actually, a shop' ),
+			] )
+		).toBe( false );
+	} );
+} );
+
 describe( 'entriesToRenderItems – legacy payload markers', () => {
 	it( 'strips media payload marker lines from any tool output', () => {
 		const items = entriesToRenderItems( [

apps/studio/src/components/studio-code-session/conversation/index.tsx

@@ -22,7 +22,7 @@ import {
 import { __ } from '@wordpress/i18n';
 import { image, page } from '@wordpress/icons';
 import { Icon } from '@wordpress/ui';
-import { useEffect, useMemo, useState } from 'react';
+import { useCallback, useEffect, useMemo, useRef, useState } from 'react';
 import { cx } from 'src/lib/cx';
 import { getIpcApi } from 'src/lib/get-ipc-api';
 import { Markdown } from '../markdown';
@@ -35,6 +35,7 @@ type RenderItem =
 	| {
 			kind: 'user-text';
 			key: string;
+			entryId: string;
 			text: string;
 			attachments?: StudioChatAttachmentSummary[];
 	  }
@@ -125,14 +126,44 @@ export function entriesToRenderItems( entries: SessionEntry[] ): RenderItem[] {
 	// only if that case ever shows wrong highlights.
 	let questionOrdinal = 0;
 
+	// Pre-scan: find entries superseded by edits. For each edit marker, hide
+	// everything from the original entry through the marker itself.
+	const editMarkers = new Map< string, number >();
+	entries.forEach( ( entry, index ) => {
+		if ( isStudioCustomEntryOfType( entry, 'studio.message_edited' ) ) {
+			const data = ( entry as StudioCustomEntry< 'studio.message_edited' > ).data;
+			if ( data?.originalEntryId ) {
+				editMarkers.set( data.originalEntryId, index );
+			}
+		}
+	} );
+	const hiddenIndices = new Set< number >();
+	if ( editMarkers.size > 0 ) {
+		let hideUntilIndex = -1;
+		for ( let i = 0; i < entries.length; i += 1 ) {
+			const entryId = ( entries[ i ] as { id?: string } ).id;
+			if ( entryId && editMarkers.has( entryId ) ) {
+				hideUntilIndex = editMarkers.get( entryId )!;
+			}
+			if ( i <= hideUntilIndex ) {
+				hiddenIndices.add( i );
+			}
+		}
+	}
+
 	const items: RenderItem[] = [];
 	entries.forEach( ( entry, entryIndex ) => {
+		if ( hiddenIndices.has( entryIndex ) ) {
+			return;
+		}
+
 		if ( isStudioCustomEntryOfType( entry, 'studio.user_prompt' ) ) {
 			const data = ( entry as StudioCustomEntry< 'studio.user_prompt' > ).data;
 			if ( ! data || data.source !== 'prompt' ) return;
 			items.push( {
 				kind: 'user-text',
 				key: `${ entryIndex }:user`,
+				entryId: ( entry as { id: string } ).id,
 				text: data.text,
 				attachments: data.attachments,
 			} );
@@ -221,6 +252,21 @@ export function entriesToRenderItems( entries: SessionEntry[] ): RenderItem[] {
 	return items;
 }
 
+// True when the most recent turn was interrupted by the user.
+export function isConversationStopped( entries: SessionEntry[] ): boolean {
+	for ( let index = entries.length - 1; index >= 0; index -= 1 ) {
+		const entry = entries[ index ];
+		if ( isStudioCustomEntryOfType( entry, 'studio.turn_closed' ) ) {
+			const data = ( entry as StudioCustomEntry< 'studio.turn_closed' > ).data;
+			return data?.status === 'interrupted';
+		}
+		if ( isStudioCustomEntryOfType( entry, 'studio.user_prompt' ) ) {
+			return false;
+		}
+	}
+	return false;
+}
+
 // Progress from earlier turns must not leak into the current indicator, so
 // the scan stops at the nearest turn boundary.
 function findLatestProgressMessage( entries: SessionEntry[] ): string | null {
@@ -243,10 +289,83 @@ function findLatestProgressMessage( entries: SessionEntry[] ): string | null {
 function UserTurn( {
 	text,
 	attachments,
+	editable = false,
+	onSubmitEdit,
 }: {
 	text: string;
 	attachments?: StudioChatAttachmentSummary[];
+	editable?: boolean;
+	onSubmitEdit?: ( newText: string ) => void;
 } ) {
+	const [ isEditing, setIsEditing ] = useState( false );
+	const [ editText, setEditText ] = useState( text );
+	const textareaRef = useRef< HTMLTextAreaElement >( null );
+
+	const startEditing = useCallback( () => {
+		setEditText( text );
+		setIsEditing( true );
+	}, [ text ] );
+
+	const cancelEditing = useCallback( () => {
+		setIsEditing( false );
+	}, [] );
+
+	const submitEdit = useCallback( () => {
+		const trimmed = editText.trim();
+		if ( trimmed && onSubmitEdit ) {
+			onSubmitEdit( trimmed );
+		}
+		setIsEditing( false );
+	}, [ editText, onSubmitEdit ] );
+
+	const handleKeyDown = useCallback(
+		( event: React.KeyboardEvent< HTMLTextAreaElement > ) => {
+			if ( event.key === 'Escape' ) {
+				cancelEditing();
+			} else if ( event.key === 'Enter' && ! event.shiftKey ) {
+				event.preventDefault();
+				submitEdit();
+			}
+		},
+		[ cancelEditing, submitEdit ]
+	);
+
+	useEffect( () => {
+		if ( isEditing && textareaRef.current ) {
+			const textarea = textareaRef.current;
+			textarea.focus();
+			textarea.selectionStart = textarea.value.length;
+		}
+	}, [ isEditing ] );
+
+	if ( isEditing ) {
+		return (
+			<div className={ styles.userTurn }>
+				<textarea
+					ref={ textareaRef }
+					className={ styles.userEditTextarea }
+					value={ editText }
+					onChange={ ( e ) => setEditText( e.target.value ) }
+					onKeyDown={ handleKeyDown }
+					rows={ 3 }
+				/>
+				<div className={ styles.userEditActions }>
+					<button type="button" className={ styles.userEditCancelButton } onClick={ cancelEditing }>
+						{ __( 'Cancel' ) }
+					</button>
+					<button
+						type="button"
+						className={ styles.userEditSendButton }
+						onClick={ submitEdit }
+						disabled={ ! editText.trim() }
+					>
+						{ __( 'Send' ) }
+					</button>
+				</div>
+			</div>
+		);
+	}
+
 	return (
 		<div className={ styles.userTurn }>
 			<div className={ styles.userText }>{ text }</div>
@@ -276,6 +395,18 @@ function UserTurn( {
 					) }
 				</ul>
 			) : null }
+			{ editable && onSubmitEdit ? (
+				<div className={ styles.userActions }>
+					<button
+						type="button"
+						className={ styles.userEditButton }
+						onClick={ startEditing }
+						aria-label={ __( 'Edit your last message' ) }
+					>
+						{ __( 'Edit' ) }
+					</button>
+				</div>
+			) : null }
 		</div>
 	);
 }
@@ -492,6 +623,8 @@ export function Conversation( {
 	pendingAnswers,
 	answeredQuestions,
 	onAnswerQuestion,
+	canEditLastUserMessage = false,
+	onEditUserMessage,
 }: {
 	data: LoadedAiSession;
 	isRunning: boolean;
@@ -500,22 +633,74 @@ export function Conversation( {
 	pendingAnswers: Record< string, string >;
 	answeredQuestions: Record< string, string >;
 	onAnswerQuestion: ( question: string, label: string ) => void;
+	canEditLastUserMessage?: boolean;
+	onEditUserMessage?: ( entryId: string, text: string ) => void;
 } ) {
 	const entries = data.entries;
 	const items = useMemo( () => entriesToRenderItems( entries ), [ entries ] );
 	const progressMessage = useMemo(
 		() => ( isRunning ? findLatestProgressMessage( entries ) : null ),
 		[ entries, isRunning ]
 	);
+	// Only the most recent user prompt is offered for editing, and only once
+	// the run behind it has been stopped.
+	const lastUserTextKey = useMemo( () => {
+		for ( let index = items.length - 1; index >= 0; index -= 1 ) {
+			if ( items[ index ].kind === 'user-text' ) {
+				return items[ index ].key;
+			}
+		}
+		return null;
+	}, [ items ] );
+
+	// Optimistic hide: the persisted `studio.message_edited` marker only takes
+	// effect once the cache refreshes from disk. This local set hides the old
+	// turn immediately so the user sees the replacement right away.
+	const [ optimisticHiddenKeys, setOptimisticHiddenKeys ] = useState< Set< string > >(
+		() => new Set()
+	);
+	const handleEditSubmit = useCallback(
+		( entryId: string, newText: string ) => {
+			const keysToHide = new Set< string >();
+			let lastUserIdx = -1;
+			for ( let i = items.length - 1; i >= 0; i -= 1 ) {
+				if ( items[ i ].kind === 'user-text' ) {
+					lastUserIdx = i;
+					break;
+				}
+			}
+			if ( lastUserIdx >= 0 ) {
+				for ( let i = lastUserIdx; i < items.length; i += 1 ) {
+					keysToHide.add( items[ i ].key );
+				}
+				setOptimisticHiddenKeys( keysToHide );
+			}
+			onEditUserMessage?.( entryId, newText );
+		},
+		[ items, onEditUserMessage ]
+	);
 
 	return (
 		<div className={ styles.root }>
 			{ items.map( ( item ) => {
+				if ( optimisticHiddenKeys.has( item.key ) ) {
+					return null;
+				}
 				switch ( item.kind ) {
-					case 'user-text':
+					case 'user-text': {
+						const editable = canEditLastUserMessage && item.key === lastUserTextKey;
 						return (
-							<UserTurn key={ item.key } text={ item.text } attachments={ item.attachments } />
+							<UserTurn
+								key={ item.key }
+								text={ item.text }
+								attachments={ item.attachments }
+								editable={ editable }
+								onSubmitEdit={
+									editable ? ( newText ) => handleEditSubmit( item.entryId, newText ) : undefined
+								}
+							/>
 						);
+					}
 					case 'assistant-text':
 						return <AssistantText key={ item.key } text={ item.text } />;
 					case 'tool-use':