Preserve temporary message in Studio Code UI

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

@@ -0,0 +1,75 @@
+import { DEFAULT_MODEL } from '@studio/common/ai/models';
+import { QueryClient, QueryClientProvider } from '@tanstack/react-query';
+import { fireEvent, render, screen, waitFor } from '@testing-library/react';
+import { beforeEach, describe, expect, it, vi } from 'vitest';
+import { Composer } from '.';
+import type { ComposerSendAttachments } from './use-composer-attachments';
+
+const defaultProps = {
+	busy: false,
+	error: null,
+	model: DEFAULT_MODEL,
+	onSend: vi.fn< ( prompt: string, attachments: ComposerSendAttachments ) => Promise< void > >(),
+	onInterrupt: vi.fn< () => Promise< void > >(),
+	entries: [],
+};
+
+function renderComposer( props: Partial< Parameters< typeof Composer >[ 0 ] > = {} ) {
+	const queryClient = new QueryClient();
+	return render(
+		<QueryClientProvider client={ queryClient }>
+			<Composer { ...defaultProps } sessionId="session-1" { ...props } />
+		</QueryClientProvider>
+	);
+}
+
+describe( 'Composer', () => {
+	beforeEach( () => {
+		localStorage.clear();
+		defaultProps.onSend.mockReset();
+		defaultProps.onInterrupt.mockReset();
+		defaultProps.onSend.mockResolvedValue( undefined );
+		defaultProps.onInterrupt.mockResolvedValue( undefined );
+	} );
+
+	it( 'restores unsent drafts for the same session', () => {
+		const { unmount } = renderComposer();
+
+		fireEvent.change( screen.getByRole( 'combobox' ), {
+			target: { value: 'Update the homepage copy' },
+		} );
+		unmount();
+
+		renderComposer();
+
+		expect( screen.getByRole( 'combobox' ) ).toHaveValue( 'Update the homepage copy' );
+	} );
+
+	it( 'keeps drafts scoped to their session', () => {
+		const { unmount } = renderComposer();
+
+		fireEvent.change( screen.getByRole( 'combobox' ), {
+			target: { value: 'Update the homepage copy' },
+		} );
+		unmount();
+
+		renderComposer( { sessionId: 'session-2' } );
+
+		expect( screen.getByRole( 'combobox' ) ).toHaveValue( '' );
+	} );
+
+	it( 'clears the stored draft after sending', async () => {
+		const { unmount } = renderComposer();
+
+		fireEvent.change( screen.getByRole( 'combobox' ), {
+			target: { value: 'Update the homepage copy' },
+		} );
+		fireEvent.click( screen.getByRole( 'button', { name: 'Send' } ) );
+
+		await waitFor( () => expect( defaultProps.onSend ).toHaveBeenCalled() );
+		unmount();
+		renderComposer();
+
+		expect( screen.getByRole( 'combobox' ) ).toHaveValue( '' );
+	} );
+} );

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

@@ -6,7 +6,7 @@ import { createInterpolateElement } from '@wordpress/element';
 import { __ } from '@wordpress/i18n';
 import { arrowUp, chevronDownSmall, closeSmall, page } from '@wordpress/icons';
 import { Icon } from '@wordpress/ui';
-import { useCallback, useEffect, useRef, useState } from 'react';
+import { useCallback, useEffect, useRef, useState, type SetStateAction } from 'react';
 import { cx } from 'src/lib/cx';
 import { getIpcApi } from 'src/lib/get-ipc-api';
 import * as Menu from '../menu';
@@ -98,6 +98,36 @@ function formatAttachmentSize( bytes: number ): string {
 	return `${ ( bytes / ( 1024 * 1024 ) ).toFixed( 1 ) } MB`;
 }
 
+function getDraftStorageKey( sessionId: string | undefined ): string | null {
+	return sessionId ? `studio_code_session_draft:${ sessionId }` : null;
+}
+
+function loadDraft( storageKey: string | null ): string {
+	if ( ! storageKey ) {
+		return '';
+	}
+	try {
+		return localStorage.getItem( storageKey ) ?? '';
+	} catch {
+		return '';
+	}
+}
+
+function saveDraft( storageKey: string | null, value: string ): void {
+	if ( ! storageKey ) {
+		return;
+	}
+	try {
+		if ( value ) {
+			localStorage.setItem( storageKey, value );
+		} else {
+			localStorage.removeItem( storageKey );
+		}
+	} catch {
+		// Ignore storage errors.
+	}
+}
+
 export function Composer( {
 	busy,
 	isInterrupting = false,
@@ -112,11 +142,23 @@ export function Composer( {
 	draftPrompt,
 	previewPrompt,
 }: ComposerProps ) {
-	const [ value, setValue ] = useState( '' );
+	const draftStorageKey = getDraftStorageKey( sessionId );
+	const [ value, setValue ] = useState( () => loadDraft( draftStorageKey ) );
 	const textareaRef = useRef< HTMLTextAreaElement | null >( null );
 	const fileInputRef = useRef< HTMLInputElement | null >( null );
 	const appliedDraftPromptIdRef = useRef< number | null >( null );
 	const queryClient = useQueryClient();
+	const setDraftValue = useCallback(
+		( nextValue: SetStateAction< string > ) => {
+			setValue( ( previousValue ) => {
+				const resolvedValue =
+					typeof nextValue === 'function' ? nextValue( previousValue ) : nextValue;
+				saveDraft( draftStorageKey, resolvedValue );
+				return resolvedValue;
+			} );
+		},
+		[ draftStorageKey ]
+	);
 
 	// File/image attachments (attach button + drag-and-drop). Images ride as
 	// base64 content blocks; other files are referenced by disk path.
@@ -136,7 +178,7 @@ export function Composer( {
 			return;
 		}
 		appliedDraftPromptIdRef.current = draftPrompt.id;
-		setValue( draftPrompt.prompt );
+		setDraftValue( draftPrompt.prompt );
 		queueMicrotask( () => {
 			const node = textareaRef.current;
 			if ( ! node ) {
@@ -146,11 +188,15 @@ export function Composer( {
 			const length = node.value.length;
 			node.setSelectionRange( length, length );
 		} );
-	}, [ draftPrompt ] );
+	}, [ draftPrompt, setDraftValue ] );
+
+	useEffect( () => {
+		setValue( loadDraft( draftStorageKey ) );
+	}, [ draftStorageKey ] );
 
 	// Inline slash-command autocomplete (popup, keyboard nav, ARIA wiring, and
 	// the toolbar "/" toggle). Kept in its own hook so the Composer stays lean.
-	const slash = useSlashCommands( { value, setValue, textareaRef, previewPrompt } );
+	const slash = useSlashCommands( { value, setValue: setDraftValue, textareaRef, previewPrompt } );
 
 	// Cross-family swap state. We hold the picked model here while the
 	// confirmation dialog is open; nothing is persisted until the user
@@ -167,7 +213,7 @@ export function Composer( {
 		}
 		const prompt = trimmed || __( 'Please review the attached files.' );
 		const sentAttachments = attachments;
-		setValue( '' );
+		setDraftValue( '' );
 		clearAttachments();
 		try {
 			await onSend( prompt, toComposerSendAttachments( sentAttachments ) );
@@ -176,10 +222,10 @@ export function Composer( {
 			// surfaces the error message via `error`. Queued sends never throw from
 			// onSend (the parent swallows the failure and clears the queue instead),
 			// so this path only trips for direct sends from the idle state.
-			setValue( trimmed );
+			setDraftValue( trimmed );
 			restoreAttachments( sentAttachments );
 		}
-	}, [ value, attachments, clearAttachments, restoreAttachments, onSend ] );
+	}, [ value, attachments, clearAttachments, restoreAttachments, onSend, setDraftValue ] );
 
 	const openFilePicker = useCallback( () => {
 		fileInputRef.current?.click();
@@ -355,7 +401,9 @@ export function Composer( {
 							value={ previewPrompt ?? value }
 							data-preview={ previewPrompt ? 'true' : 'false' }
 							{ ...slash.comboboxProps }
-							onChange={ ( event ) => setValue( event.target.value ) }
+							onChange={ ( event ) => {
+								setDraftValue( event.target.value );
+							} }
 							onKeyDown={ ( event ) => {
 								if ( slash.handleKeyDown( event ) ) {
 									return;