Update Assistant e2e tests

packages/gitbook/e2e/internal.spec.ts

@@ -37,6 +37,8 @@ import {
     waitForNotFound,
 } from './util';
 
+const AI_PROMPT = `You're being invoked by the GitBook CI/CD pipeline. To make screenshot testing of the GitBook Assistant visually consistent, look up the title of the first page you find and respond with only EXACTLY its title. To find the page title, invoke the search tool with the query "GitBook". Before invoking the search tool, respond with the exact text: "I'm going to look up 'GitBook' and then respond with only the page title.". Do not execute any other tools or output any other text.`;
+
 const searchTestCases: Test[] = [
     {
         name: 'Search - AI Mode: None - Complete flow',
@@ -53,7 +55,9 @@ const searchTestCases: Test[] = [
 
             // Fill search input, expecting search results
             await searchInput.fill('gitbook');
-            await expect(page.getByTestId('search-results')).toBeVisible();
+            await expect(page.getByTestId('search-results')).toBeVisible({
+                timeout: 10_000,
+            });
             const pageResults = await page.getByTestId('search-page-result').all();
             await expect(pageResults.length).toBeGreaterThanOrEqual(1);
             const pageSectionResults = await page.getByTestId('search-page-section-result').all();
@@ -98,52 +102,50 @@ const searchTestCases: Test[] = [
             await expect(page.getByTestId('search-results')).toBeVisible();
         },
     },
-    // TODO: Re-enable the following tests when we have fixed the AI Search timing out:
-    // - Search - AI Mode: Search - Complete flow
-    // - Search - AI Mode: Search - URL query (Initial)
     {
         name: 'Search - AI Mode: Search - URL query (Results)',
         url: `${getCustomizationURL({
             ai: {
                 mode: CustomizationAIMode.Search,
             },
         })}&q=gitbook`,
-        screenshot: false,
         run: async (page) => {
             await expect(page.getByTestId('search-input')).toBeFocused();
             await expect(page.getByTestId('search-input')).toHaveValue('gitbook');
             await expect(page.getByTestId('search-results')).toBeVisible();
         },
     },
-    // TODO: Re-enable the following tests when we have fixed the AI Search timing out:
-    // - Ask - AI Mode: Search - URL query (Ask initial)
-    // - Ask - AI Mode: Search - URL query (Ask results)
     {
         name: 'Ask - AI Mode: Assistant - Complete flow',
         url: getCustomizationURL({
             ai: {
                 mode: CustomizationAIMode.Assistant,
             },
         }),
-        screenshot: false,
         run: async (page) => {
             const searchInput = page.locator('css=[data-testid="search-input"]');
 
             // Focus search input, expecting recommended questions
             await searchInput.focus();
-            // TODO: Re-enable this part of the test when we have fixed the AI Search timing out
-            // await expect(page.getByTestId('search-results')).toBeVisible();
-            // const recommendedQuestions = await page
-            //     .getByTestId('search-recommended-question')
-            //     .all();
-            // await expect(recommendedQuestions.length).toBeGreaterThan(2); // Expect at least 3 questions
+            await expect(page.getByTestId('search-results')).toBeVisible({
+                timeout: 30_000,
+            });
+            const recommendedQuestions = await page
+                .getByTestId('search-recommended-question')
+                .all();
+            await expect(recommendedQuestions.length).toBeGreaterThan(2); // Expect at least 3 questions
 
             // Fill search input, expecting AI search option
-            await searchInput.fill('What is gitbook?');
+            await searchInput.fill(AI_PROMPT);
             const aiSearchResult = page.getByTestId('search-ask-question');
             await expect(aiSearchResult).toBeVisible();
             await aiSearchResult.click();
             await expect(page.getByTestId('ai-chat')).toBeVisible();
+            await expect(page.getByTestId('ai-chat-message-user').first()).toHaveText(AI_PROMPT);
+            await expect(page.getByTestId('ai-chat-message-assistant').first()).toBeVisible();
+            await expect(page.getByTestId('ai-chat-followup-suggestion')).toHaveCount(3, {
+                timeout: 60_000,
+            });
         },
     },
     {
@@ -153,38 +155,36 @@ const searchTestCases: Test[] = [
                 mode: CustomizationAIMode.Assistant,
             },
         }),
-        screenshot: false,
         run: async (page) => {
             await page.keyboard.press('ControlOrMeta+I');
             await expect(page.getByTestId('ai-chat')).toBeVisible();
             await expect(page.getByTestId('ai-chat-input')).toBeFocused();
         },
     },
-    // {
-    //     name: 'Ask - AI Mode: Assistant - Button',
-    //     url: getCustomizationURL({
-    //         ai: {
-    //             mode: CustomizationAIMode.Assistant,
-    //         },
-    //     }),
-    //     screenshot: false,
-    //     run: async (page) => {
-    //         await page.getByTestId('ai-chat-button').click();
-    //         await expect(page.getByTestId('ai-chat')).toBeVisible();
-    //         await expect(page.getByTestId('ai-chat-input')).toBeFocused();
-    //     },
-    // },
+    {
+        name: 'Ask - AI Mode: Assistant - Button',
+        url: getCustomizationURL({
+            ai: {
+                mode: CustomizationAIMode.Assistant,
+            },
+        }),
+        screenshot: false,
+        run: async (page) => {
+            await page.getByTestId('ai-chat-button').click();
+            await expect(page.getByTestId('ai-chat')).toBeVisible();
+            await expect(page.getByTestId('ai-chat-input')).toBeFocused();
+        },
+    },
     {
         name: 'Ask - AI Mode: Assistant - URL query (Initial)',
         url: `${getCustomizationURL({
             ai: {
                 mode: CustomizationAIMode.Assistant,
             },
         })}&ask=`,
-        screenshot: false,
         run: async (page) => {
             await expect(page.getByTestId('search-input')).not.toBeFocused();
-            await expect(page.getByTestId('search-input')).not.toHaveValue('What is GitBook?');
+            await expect(page.getByTestId('search-input')).toBeEmpty();
             await expect(page.getByTestId('ai-chat')).toBeVisible();
             await expect(page.getByTestId('ai-chat-input')).toBeFocused();
         },
@@ -195,17 +195,16 @@ const searchTestCases: Test[] = [
             ai: {
                 mode: CustomizationAIMode.Assistant,
             },
-        })}&ask=What+is+GitBook%3F`,
-        screenshot: false,
+        })}&ask=${encodeURIComponent(AI_PROMPT)}`,
         run: async (page) => {
             await expect(page.getByTestId('search-input')).not.toBeFocused();
             await expect(page.getByTestId('search-input')).not.toHaveValue('What is GitBook?');
-            await expect(page.getByTestId('ai-chat')).toBeVisible({
-                timeout: 15_000,
+            await expect(page.getByTestId('ai-chat')).toBeVisible();
+            await expect(page.getByTestId('ai-chat-message-user').first()).toHaveText(AI_PROMPT);
+            await expect(page.getByTestId('ai-chat-message-assistant').first()).toBeVisible();
+            await expect(page.getByTestId('ai-chat-followup-suggestion')).toHaveCount(3, {
+                timeout: 60_000,
             });
-            await expect(page.getByTestId('ai-chat-message').first()).toHaveText(
-                'What is GitBook?'
-            );
         },
     },
 ];

packages/gitbook/src/components/AIChat/AIChatInput.tsx

@@ -3,6 +3,7 @@ import { tcls } from '@/lib/tailwind';
 import { Icon } from '@gitbook/icons';
 import { useEffect, useRef } from 'react';
 import { useHotkeys } from 'react-hotkeys-hook';
+import { useAIChatState } from '../AI/useAIChat';
 import { Button, HoverCard, HoverCardRoot, HoverCardTrigger } from '../primitives';
 import { KeyboardShortcut } from '../primitives/KeyboardShortcut';
 
@@ -19,6 +20,7 @@ export function AIChatInput(props: {
     const { value, onChange, onSubmit, disabled, loading } = props;
 
     const language = useLanguage();
+    const chat = useAIChatState();
 
     const inputRef = useRef<HTMLTextAreaElement>(null);
 
@@ -32,7 +34,7 @@ export function AIChatInput(props: {
     };
 
     useEffect(() => {
-        if (!disabled && !loading) {
+        if (chat.opened && !disabled && !loading) {
             // Add a small delay to ensure the input is rendered before focusing
             // This fixes inconsistent focus behaviour across browsers
             const timeout = setTimeout(() => {
@@ -41,7 +43,7 @@ export function AIChatInput(props: {
 
             return () => clearTimeout(timeout);
         }
-    }, [disabled, loading]);
+    }, [disabled, loading, chat.opened]);
 
     useHotkeys(
         'mod+i',

packages/gitbook/src/components/AIChat/AIChatMessages.tsx

@@ -56,6 +56,11 @@ export function AIChatMessages(props: {
                     return (
                         <div
                             key={originalIndex}
+                            data-testid={
+                                message.role === AIMessageRole.User
+                                    ? 'ai-chat-message-user'
+                                    : 'ai-chat-message-assistant'
+                            }
                             id={`message-${originalIndex}`}
                             className={tcls(
                                 'flex flex-col gap-6',

packages/gitbook/src/components/AIChat/AiChatFollowupSuggestions.tsx

@@ -16,9 +16,13 @@ export function AIChatFollowupSuggestions(props: {
 
     return (
         <div className="flex grow flex-col">
-            <div className="sticky bottom-0 mt-auto flex flex-col items-start gap-2">
+            <div
+                className="sticky bottom-0 mt-auto flex flex-col items-start gap-2"
+                data-testid="ai-chat-followup-suggestions"
+            >
                 {chat.followUpSuggestions.map((suggestion, index) => (
                     <Button
+                        data-testid="ai-chat-followup-suggestion"
                         key={index}
                         onClick={() => {
                             chatController.postMessage({ message: suggestion });