Zustand Proof of Concept

package.json

@@ -138,7 +138,8 @@
 		"unzipper": "0.10.11",
 		"url-loader": "^4.1.1",
 		"wpcom": "^5.4.2",
-		"yargs": "17.7.2"
+		"yargs": "17.7.2",
+		"zustand": "^5.0.3"
 	},
 	"optionalDependencies": {
 		"appdmg": "^0.6.6"

src/components/content-tab-assistant.tsx

@@ -17,6 +17,7 @@ import { usePromptUsage } from '../hooks/use-prompt-usage';
 import { useWelcomeMessages } from '../hooks/use-welcome-messages';
 import { cx } from '../lib/cx';
 import { getIpcApi } from '../lib/get-ipc-api';
+import { useChatInputStore } from '../stores/chat-input-store';
 import ClearHistoryReminder from './ai-clear-history-reminder';
 import { AIInput } from './ai-input';
 import { MessageThinking } from './assistant-thinking';
@@ -368,19 +369,20 @@ export function ContentTabAssistant( { selectedSite }: ContentTabAssistantProps
 	const { __ } = useI18n();
 	const lastMessage = messages.length === 0 ? undefined : messages[ messages.length - 1 ];
 	const hasFailedMessage = messages.some( ( msg ) => msg.failedMessage );
+	const { getChatInput, saveChatInput } = useChatInputStore();
 
 	// Restore prompt input when site changes
 	useEffect( () => {
-		setCurrentInput( chatContext.getChatInput( selectedSite.id ) );
-	}, [ selectedSite.id, chatContext ] );
+		setCurrentInput( getChatInput( selectedSite.id ) );
+	}, [ selectedSite.id, getChatInput ] );
 
 	// Save prompt input when it changes
 	const setInput = useCallback(
 		( input: string ) => {
-			chatContext.saveChatInput( input, selectedSite.id );
+			saveChatInput( input, selectedSite.id );
 			setCurrentInput( input );
 		},
-		[ selectedSite.id, chatContext ]
+		[ selectedSite.id, saveChatInput ]
 	);
 
 	const submitPrompt = useCallback(

src/hooks/tests/use-chat-context.test.tsx

@@ -1,5 +1,6 @@
 import { renderHook, act } from '@testing-library/react';
 import { getIpcApi } from '../../lib/get-ipc-api';
+import { useChatInputStore } from '../../stores/chat-input-store';
 import { ChatProvider, useChatContext } from '../use-chat-context';
 import { useCheckInstalledApps } from '../use-check-installed-apps';
 import { useGetWpVersion } from '../use-get-wp-version';
@@ -76,6 +77,11 @@ beforeEach( () => {
 		terminalWpCliEnabled: false,
 	} );
 	setupWpCliResult( { themes: [], plugins: [] } );
+
+	// Reset Zustand store
+	act( () => {
+		useChatInputStore.setState( { inputBySite: {} } );
+	} );
 } );
 
 afterEach( () => {
@@ -102,8 +108,6 @@ describe( 'useChatContext hook', () => {
 				availableEditors: [],
 				siteName: 'Test Site',
 				os: 'darwin',
-				getChatInput: expect.any( Function ),
-				saveChatInput: expect.any( Function ),
 			} );
 		} );
 
@@ -212,8 +216,6 @@ describe( 'useChatContext hook', () => {
 				availableEditors: [],
 				siteName: 'Another Test Site',
 				os: 'darwin',
-				getChatInput: expect.any( Function ),
-				saveChatInput: expect.any( Function ),
 			} );
 		} );
 
@@ -370,8 +372,6 @@ describe( 'useChatContext hook', () => {
 				availableEditors: [],
 				siteName: 'New Test Site',
 				os: 'darwin',
-				getChatInput: expect.any( Function ),
-				saveChatInput: expect.any( Function ),
 			} );
 		} );
 
@@ -511,7 +511,7 @@ describe( 'useChatContext hook', () => {
 
 	describe( 'chat input functionality', () => {
 		it( 'saves and retrieves chat input by site', () => {
-			const { result } = renderHook( () => useChatContext(), { wrapper } );
+			const { result } = renderHook( () => useChatInputStore(), { wrapper } );
 
 			act( () => {
 				result.current.saveChatInput( 'test input 1', 'site-1' );
@@ -523,7 +523,7 @@ describe( 'useChatContext hook', () => {
 		} );
 
 		it( 'returns empty string for non-existent site input', () => {
-			const { result } = renderHook( () => useChatContext(), { wrapper } );
+			const { result } = renderHook( () => useChatInputStore(), { wrapper } );
 			expect( result.current.getChatInput( 'non-existent-site' ) ).toBe( '' );
 		} );
 	} );

src/hooks/use-chat-context.tsx

@@ -28,8 +28,6 @@ export interface ChatContextType {
 	os: string;
 	availableEditors: string[];
 	siteName?: string;
-	getChatInput: ( siteId: string ) => string;
-	saveChatInput: ( input: string, siteId: string ) => void;
 }
 const ChatContext = createContext< ChatContextType >( {
 	currentURL: '',
@@ -43,10 +41,6 @@ const ChatContext = createContext< ChatContextType >( {
 	availableEditors: [] as string[],
 	os: '',
 	siteName: '',
-	getChatInput: () => '',
-	saveChatInput: () => {
-		// noop
-	},
 } );
 
 interface ChatProviderProps {
@@ -67,7 +61,6 @@ const parseWpCliOutput = ( stdout: string, defaultValue: string[] ): string[] =>
 
 export const ChatProvider: React.FC< ChatProviderProps > = ( { children } ) => {
 	const initialLoad = useRef< Record< string, boolean > >( {} );
-	const inputBySite = useRef< Record< string, string > >( {} );
 	const installedApps = useCheckInstalledApps();
 	const { data: sites, loadingSites, selectedSite } = useSiteDetails();
 	const wpVersion = useGetWpVersion( selectedSite || ( {} as SiteDetails ) );
@@ -82,14 +75,6 @@ export const ChatProvider: React.FC< ChatProviderProps > = ( { children } ) => {
 		return installedApps[ app as keyof InstalledApps ];
 	} );
 
-	const getChatInput = useCallback( ( siteId: string ) => {
-		return inputBySite.current[ siteId ] ?? '';
-	}, [] );
-
-	const saveChatInput = useCallback( ( input: string, siteId: string ) => {
-		inputBySite.current[ siteId ] = input;
-	}, [] );
-
 	const fetchPluginList = useCallback( async ( siteId: string ) => {
 		const { stdout, stderr } = await getIpcApi().executeWPCLiInline( {
 			siteId,
@@ -165,13 +150,11 @@ export const ChatProvider: React.FC< ChatProviderProps > = ( { children } ) => {
 			value={ {
 				availableEditors,
 				currentURL: `http://localhost:${ sitePort }`,
-				getChatInput,
 				isBlockTheme: themeDetails?.isBlockTheme,
 				numberOfSites,
 				os: window.appGlobals?.platform,
 				phpVersion: selectedSite?.phpVersion ?? DEFAULT_PHP_VERSION,
 				pluginList: selectedSite?.id ? pluginsList[ selectedSite.id ] || [] : [],
-				saveChatInput,
 				siteName: selectedSite?.name,
 				themeList: selectedSite?.id ? themesList[ selectedSite.id ] || [] : [],
 				themeName: themeDetails?.name,

src/stores/chat-input-store.ts

@@ -0,0 +1,19 @@
+import { create } from 'zustand';
+
+type ChatInputState = {
+	inputBySite: Record< string, string >;
+	getChatInput: ( siteId: string ) => string;
+	saveChatInput: ( input: string, siteId: string ) => void;
+};
+
+export const useChatInputStore = create< ChatInputState >( ( set, get ) => ( {
+	inputBySite: {},
+	getChatInput: ( siteId ) => get().inputBySite[ siteId ] ?? '',
+	saveChatInput: ( input, siteId ) =>
+		set( ( state ) => ( {
+			inputBySite: {
+				...state.inputBySite,
+				[ siteId ]: input,
+			},
+		} ) ),
+} ) );