refactor chat input context in favor of zustand store

Automattic · Jan 13, 2025 · a7efc5d · a7efc5d
1 parent 33b9b0b
commit a7efc5d
Show file tree

Hide file tree

Showing 4 changed files with 41 additions and 29 deletions.
diff --git a/src/components/content-tab-assistant.tsx b/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(

diff --git a/src/hooks/tests/use-chat-context.test.tsx b/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( '' );
 		} );
 	} );

diff --git a/src/hooks/use-chat-context.tsx b/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,

diff --git a/src/stores/tests/chat-input-store.test.ts b/src/stores/tests/chat-input-store.test.ts
@@ -0,0 +1,27 @@
+import { act, renderHook } from '@testing-library/react';
+import { useChatInputStore } from '../chat-input-store';
+
+describe( 'useChatInputStore', () => {
+	beforeEach( () => {
+		act( () => {
+			useChatInputStore.setState( { inputBySite: {} } );
+		} );
+	} );
+
+	it( 'saves and retrieves chat input by site', () => {
+		const { result } = renderHook( () => useChatInputStore() );
+
+		act( () => {
+			result.current.saveChatInput( 'test input 1', 'site-1' );
+			result.current.saveChatInput( 'test input 2', 'site-2' );
+		} );
+
+		expect( result.current.getChatInput( 'site-1' ) ).toBe( 'test input 1' );
+		expect( result.current.getChatInput( 'site-2' ) ).toBe( 'test input 2' );
+	} );
+
+	it( 'returns empty string for non-existent site input', () => {
+		const { result } = renderHook( () => useChatInputStore() );
+		expect( result.current.getChatInput( 'non-existent-site' ) ).toBe( '' );
+	} );
+} );