fix(askai): Fixes issue with askai state being outside DocSearchModal (#2778)

Fixes an issue with Docusaurus after the new Ask AI screens were added.
This was caused from the state and callbacks being set outside of
DocSearchModal.

- Move askAiState inside of DocSearchModal
- Handle changing the state when needed
- Updates color styling of dropdown menu

packages/docsearch-css/src/_variables.css

@@ -67,6 +67,9 @@
     var(--docsearch-muted-color) 60%,
     rgb(224, 227, 232) 95%
   );
+
+  --docsearch-dropdown-menu-background: var(--docsearch-hit-background);
+  --docsearch-dropdown-menu-item-hover-background: var(--docsearch-modal-background);
 }
 
 /* Darkmode */
@@ -105,4 +108,5 @@ html[data-theme='dark'] {
     var(--docsearch-muted-color) 60%,
     rgb(224, 227, 232) 95%
   );
+  --docsearch-dropdown-menu-item-hover-background: #10121e;
 }

packages/docsearch-css/src/modal.css

@@ -1506,7 +1506,7 @@ assistive tech users */
   padding: 8px 0;
   flex-direction: column;
   border-radius: var(--docsearch-border-radius);
-  background-color: var(--docsearch-modal-background);
+  background-color: var(--docsearch-dropdown-menu-background);
   box-shadow: 0px 0px 0px 1px #21243D0D, 0px 8px 16px -4px #21243D40;
   min-width: 195px;
   top: calc(100% + 12px);
@@ -1531,5 +1531,5 @@ assistive tech users */
 }
 
 .DocSearch-Menu-item:hover {
-  background-color: var(--docsearch-hit-background);
+  background-color: var(--docsearch-dropdown-menu-item-hover-background);
 }

packages/docsearch-react/src/DocSearch.tsx

@@ -12,7 +12,6 @@ import type {
   KeyboardShortcuts,
   StoredDocSearchHit,
 } from './types';
-import type { AskAiState } from './types/AskiAi';
 import { useDocSearchKeyboardEvents } from './useDocSearchKeyboardEvents';
 import { useTheme } from './useTheme';
 
@@ -179,8 +178,7 @@ export function DocSearch(props: DocSearchProps): JSX.Element {
   const searchButtonRef = React.useRef<HTMLButtonElement>(null);
   const [isOpen, setIsOpen] = React.useState(false);
   const [initialQuery, setInitialQuery] = React.useState<string | undefined>(props?.initialQuery || undefined);
-  const [askAiState, setAskAiState] = React.useState<AskAiState>('initial');
-  const isAskAiActive = askAiState !== 'initial';
+  const [isAskAiActive, setIsAskAiActive] = React.useState(false);
 
   let currentPlaceholder =
     props?.translations?.modal?.searchBox?.placeholderText || props?.placeholder || 'Search docs';
@@ -198,9 +196,9 @@ export function DocSearch(props: DocSearchProps): JSX.Element {
 
   const onAskAiToggle = React.useCallback(
     (askAitoggle: boolean) => {
-      setAskAiState(askAitoggle ? 'conversation' : 'initial');
+      setIsAskAiActive(askAitoggle);
     },
-    [setAskAiState],
+    [setIsAskAiActive],
   );
 
   const onOpen = React.useCallback(() => {
@@ -211,9 +209,9 @@ export function DocSearch(props: DocSearchProps): JSX.Element {
     setIsOpen(false);
     setInitialQuery(props?.initialQuery);
     if (isAskAiActive) {
-      setAskAiState('initial');
+      setIsAskAiActive(false);
     }
-  }, [setIsOpen, props.initialQuery, isAskAiActive, setAskAiState]);
+  }, [setIsOpen, props.initialQuery, isAskAiActive, setIsAskAiActive]);
 
   const onInput = React.useCallback(
     (event: KeyboardEvent) => {
@@ -254,8 +252,6 @@ export function DocSearch(props: DocSearchProps): JSX.Element {
             translations={props?.translations?.modal}
             isAskAiActive={isAskAiActive}
             canHandleAskAi={canHandleAskAi}
-            askAiState={askAiState}
-            setAskAiState={setAskAiState}
             onAskAiToggle={onAskAiToggle}
             onClose={onClose}
           />,

packages/docsearch-react/src/DocSearchModal.tsx

@@ -45,8 +45,6 @@ export type DocSearchModalProps = DocSearchProps & {
   isAskAiActive?: boolean;
   canHandleAskAi?: boolean;
   translations?: ModalTranslations;
-  askAiState: AskAiState;
-  setAskAiState: (state: AskAiState) => void;
 };
 
 /**
@@ -302,8 +300,6 @@ export function DocSearchModal({
   indices = [],
   indexName,
   searchParameters,
-  askAiState,
-  setAskAiState,
 }: DocSearchModalProps): JSX.Element {
   const { footer: footerTranslations, searchBox: searchBoxTranslations, ...screenStateTranslations } = translations;
   const [state, setState] = React.useState<DocSearchState<InternalDocSearchHit>>({
@@ -332,6 +328,7 @@ export function DocSearchModal({
   const askAiConfig = typeof askAi === 'object' ? askAi : null;
   const askAiConfigurationId = typeof askAi === 'string' ? askAi : askAiConfig?.assistantId || null;
   const askAiSearchParameters = askAiConfig?.searchParameters;
+  const [askAiState, setAskAiState] = React.useState<AskAiState>('initial');
 
   // Format the `indexes` to be used until `indexName` and `searchParameters` props are fully removed.
   const indexes: DocSearchIndex[] = [];
@@ -507,7 +504,7 @@ export function DocSearchModal({
       >
     >(undefined);
 
-  const handleAskAiToggle = React.useCallback(
+  const handleSelectAskAiQuestion = React.useCallback(
     (toggle: boolean, query: string) => {
       onAskAiToggle(toggle);
       setStoppedStream(false);
@@ -660,7 +657,7 @@ export function DocSearchModal({
                 },
                 onSelect({ item }): void {
                   if (item.type === 'askAI' && item.query) {
-                    handleAskAiToggle(true, item.query);
+                    handleSelectAskAiQuestion(true, item.query);
                   }
                 },
               },
@@ -775,6 +772,11 @@ export function DocSearchModal({
     }
   }, [isAskAiActive, autocomplete, setMessages]);
 
+  // Track external state in order to manage internal askAiState
+  React.useEffect(() => {
+    setAskAiState('initial');
+  }, [isAskAiActive, setAskAiState]);
+
   const onStopAskAiStreaming = async (): Promise<void> => {
     setStoppedStream(true);
 
@@ -830,10 +832,11 @@ export function DocSearchModal({
             isAskAiActive={isAskAiActive}
             askAiStatus={status}
             askAiState={askAiState}
+            setAskAiState={setAskAiState}
             onClose={onClose}
             onAskAiToggle={onAskAiToggle}
             onAskAgain={(query) => {
-              handleAskAiToggle(true, query);
+              handleSelectAskAiQuestion(true, query);
             }}
             onStopAskAiStreaming={onStopAskAiStreaming}
             onNewConversation={handleNewConversation}
@@ -864,7 +867,7 @@ export function DocSearchModal({
               status={status}
               hasCollections={hasCollections}
               askAiState={askAiState}
-              selectAskAiQuestion={handleAskAiToggle}
+              selectAskAiQuestion={handleSelectAskAiQuestion}
               onAskAiToggle={onAskAiToggle}
               onItemClick={(item, event) => {
                 // if the item is askAI toggle the screen
@@ -874,8 +877,9 @@ export function DocSearchModal({
                     setMessages(item.messages as any);
                     onAskAiToggle(true);
                   } else {
-                    handleAskAiToggle(true, item.query);
+                    handleSelectAskAiQuestion(true, item.query);
                   }
+                  setAskAiState('initial');
                   event.preventDefault();
                   return;
                 }

packages/docsearch-react/src/ScreenState.tsx

@@ -58,11 +58,11 @@ export interface ScreenStateProps<TItem extends BaseItem>
 
 export const ScreenState = React.memo(
   ({ translations = {}, ...props }: ScreenStateProps<InternalDocSearchHit>) => {
-    if (props.canHandleAskAi && props.askAiState === 'conversation-history') {
+    if (props.canHandleAskAi && props.isAskAiActive && props.askAiState === 'conversation-history') {
       return <ConversationHistoryScreen {...props} />;
     }
 
-    if (props.canHandleAskAi && props.askAiState === 'new-conversation') {
+    if (props.canHandleAskAi && props.isAskAiActive && props.askAiState === 'new-conversation') {
       return (
         <NewConversationScreen
           translations={translations?.newConversation}

packages/docsearch-react/src/SearchBox.tsx

@@ -52,11 +52,18 @@ interface SearchBoxProps
   isFromSelection: boolean;
   translations?: SearchBoxTranslations;
   askAiState: AskAiState;
+  setAskAiState: (state: AskAiState) => void;
   onNewConversation: () => void;
   onViewConversationHistory: () => void;
 }
 
-export function SearchBox({ translations = {}, askAiState, onAskAiToggle, ...props }: SearchBoxProps): JSX.Element {
+export function SearchBox({
+  translations = {},
+  askAiState,
+  onAskAiToggle,
+  setAskAiState,
+  ...props
+}: SearchBoxProps): JSX.Element {
   const {
     clearButtonTitle = 'Clear',
     clearButtonAriaLabel = 'Clear the query',
@@ -173,11 +180,12 @@ export function SearchBox({ translations = {}, askAiState, onAskAiToggle, ...pro
   const handleAskAiBackClick = React.useCallback((): void => {
     if (askAiState === 'conversation-history') {
       onAskAiToggle(true);
+      setAskAiState('initial');
       return;
     }
 
     onAskAiToggle(false);
-  }, [askAiState, onAskAiToggle]);
+  }, [askAiState, onAskAiToggle, setAskAiState]);
 
   return (
     <>