Add event listener to AIChatController

To switch to AI tab in embed upon posting a new message

Author: zenoachtig

packages/gitbook/src/components/AI/useAIChat.tsx

@@ -96,6 +96,8 @@ export type AIChatController = {
     postMessage: (input: { message: string }) => void;
     /** Clear the conversation */
     clear: () => void;
+    /** Register an event listener */
+    on: (event: 'postMessage', listener: (input: { message: string }) => void) => () => void;
 };
 
 const AIChatControllerContext = React.createContext<AIChatController | null>(null);
@@ -137,6 +139,11 @@ export function AIChatProvider(props: {
     const [, setSearchState] = useSearch();
     const language = useLanguage();
 
+    // Event listeners storage
+    const eventsRef = React.useRef<Map<'postMessage', Array<(input: { message: string }) => void>>>(
+        new Map()
+    );
+
     // Open AI chat and sync with search state
     const onOpen = React.useCallback(() => {
         const { initialQuery } = globalState.getState();
@@ -379,8 +386,18 @@ export function AIChatProvider(props: {
                 }));
             }
 
+            // Defer event listeners to next tick so React can process state updates first
+            setTimeout(() => {
+                const listeners = eventsRef.current.get('postMessage') || [];
+                listeners.forEach((listener) => listener(input));
+            }, 0);
+
             if (query === input.message) {
                 // Return early if the message is the same as the previous message
+                globalState.setState((state) => ({
+                    ...state,
+                    opened: true,
+                }));
                 return;
             }
 
@@ -440,14 +457,31 @@ export function AIChatProvider(props: {
         }));
     }, [setSearchState]);
 
+    const onEvent = React.useCallback(
+        (event: 'postMessage', listener: (input: { message: string }) => void) => {
+            const listeners = eventsRef.current.get(event) || [];
+            listeners.push(listener);
+            eventsRef.current.set(event, listeners);
+            return () => {
+                const currentListeners = eventsRef.current.get(event) || [];
+                eventsRef.current.set(
+                    event,
+                    currentListeners.filter((l) => l !== listener)
+                );
+            };
+        },
+        []
+    );
+
     const controller = React.useMemo(() => {
         return {
             open: onOpen,
             close: onClose,
             clear: onClear,
             postMessage: onPostMessage,
+            on: onEvent,
         };
-    }, [onOpen, onClose, onClear, onPostMessage]);
+    }, [onOpen, onClose, onClear, onPostMessage, onEvent]);
 
     return (
         <AIChatControllerContext.Provider value={controller}>

packages/gitbook/src/components/Embeddable/EmbeddableAIChat.tsx

@@ -37,6 +37,10 @@ export function EmbeddableAIChat() {
     const configuration = useEmbeddableConfiguration();
     const language = useLanguage();
 
+    React.useEffect(() => {
+        chatController.open();
+    }, [chatController]);
+
     // Track the view of the AI chat
     const trackEvent = useTrackEvent();
     React.useEffect(() => {

packages/gitbook/src/components/Embeddable/EmbeddableIframeAPI.tsx

@@ -40,6 +40,12 @@ export function EmbeddableIframeAPI(props: {
         embeddableConfiguration.setState({ siteTitle });
     }, [baseURL, siteTitle]);
 
+    React.useEffect(() => {
+        return chatController.on('postMessage', () => {
+            router.push(`${baseURL}/assistant`);
+        });
+    }, [router, baseURL, chatController]);
+
     React.useEffect(() => {
         if (window.parent === window) {
             return;
@@ -62,6 +68,7 @@ export function EmbeddableIframeAPI(props: {
                     chatController.postMessage({
                         message: message.message,
                     });
+                    router.push(`${baseURL}/assistant`);
                     break;
                 }
                 case 'configure': {
@@ -192,6 +199,7 @@ export function EmbeddableIframeTabs(props: { active?: string }) {
                 ? tabs.map((tab) => (
                       <Button
                           key={tab.key}
+                          data-testid={`embed-tab-${tab.key}`}
                           label={tab.label}
                           size="default"
                           variant="blank"