From 96602eb552e3056b453cf20487d94a65dcae8fb7 Mon Sep 17 00:00:00 2001
From: Nicolas Dorseuil <nicolas@gitbook.io>
Date: Tue, 16 Jun 2026 13:08:59 +0200
Subject: [PATCH] Refactor AIChatButton and StructurePreview components to
 enhance AI assistant integration, improving button rendering and
 functionality

---
 .../src/components/AIChat/AIChatButton.tsx    | 73 +++++++++++++------
 .../StructurePreview/StructurePreview.tsx     | 45 +++++++++++-
 2 files changed, 95 insertions(+), 23 deletions(-)

diff --git a/packages/gitbook/src/components/AIChat/AIChatButton.tsx b/packages/gitbook/src/components/AIChat/AIChatButton.tsx
index abf533d34..642d6f6ce 100644
--- a/packages/gitbook/src/components/AIChat/AIChatButton.tsx
+++ b/packages/gitbook/src/components/AIChat/AIChatButton.tsx
@@ -1,6 +1,9 @@
 'use client';
+import type { ReactNode } from 'react';
+
 import { useLanguage } from '@/intl/client';
 import { t, tString } from '@/intl/translate';
+import { tcls } from '@/lib/tailwind';
 import type { Assistant } from '../AI';
 import { useIsMobile } from '../hooks/useIsMobile';
 import { Button } from '../primitives';
@@ -8,6 +11,49 @@ import { KeyboardShortcut } from '../primitives/KeyboardShortcut';
 
 const MOBILE_BREAKPOINT = 688; // 43rem, equal to Tailwind's @max-2xl container breakpoint
 
+/**
+ * Button visual for an AI assistant in the header.
+ */
+export function AIChatButtonView(props: {
+    icon: ReactNode;
+    label: string;
+    onClick?: () => void;
+    showLabel?: boolean;
+    withShortcut?: boolean;
+    inert?: boolean;
+}) {
+    const { icon, label, onClick, showLabel = true, withShortcut = true, inert = false } = props;
+    const language = useLanguage();
+    const isMobile = useIsMobile(MOBILE_BREAKPOINT, '[data-gb-header-content]');
+
+    return (
+        <Button
+            icon={icon}
+            data-testid="ai-chat-button"
+            iconOnly={!showLabel || isMobile}
+            size="medium"
+            variant="header"
+            label={
+                <div className="flex items-center gap-2">
+                    {t(language, 'ai_chat_ask', label)}
+                    {withShortcut ? (
+                        <KeyboardShortcut
+                            keys={['mod', 'i']}
+                            className="border-tint-11 text-tint-1"
+                        />
+                    ) : null}
+                </div>
+            }
+            aria-label={tString(language, 'ai_chat_ask', label)}
+            onClick={inert ? undefined : onClick}
+            tabIndex={inert ? -1 : undefined}
+            className={tcls(inert ? 'pointer-events-none select-none' : null)}
+        >
+            {showLabel ? t(language, 'ask') : null}
+        </Button>
+    );
+}
+
 /**
  * Button to open/close the AI chat.
  */
@@ -17,31 +63,14 @@ export function AIChatButton(props: {
     withShortcut?: boolean;
 }) {
     const { assistant, showLabel = true, withShortcut = true } = props;
-    const language = useLanguage();
-    const isMobile = useIsMobile(MOBILE_BREAKPOINT, '[data-gb-header-content]');
 
     return (
-        <Button
+        <AIChatButtonView
             icon={assistant.icon}
-            data-testid="ai-chat-button"
-            iconOnly={!showLabel || isMobile}
-            size="medium"
-            variant="header"
-            label={
-                <div className="flex items-center gap-2">
-                    {t(language, 'ai_chat_ask', assistant.label)}
-                    {withShortcut ? (
-                        <KeyboardShortcut
-                            keys={['mod', 'i']}
-                            className="border-tint-11 text-tint-1"
-                        />
-                    ) : null}
-                </div>
-            }
-            aria-label={tString(language, 'ai_chat_ask', assistant.label)}
+            label={assistant.label}
             onClick={() => assistant.open()}
-        >
-            {showLabel ? t(language, 'ask') : null}
-        </Button>
+            showLabel={showLabel}
+            withShortcut={withShortcut}
+        />
     );
 }
diff --git a/packages/gitbook/src/components/StructurePreview/StructurePreview.tsx b/packages/gitbook/src/components/StructurePreview/StructurePreview.tsx
index d4a43fffd..f369b1d53 100644
--- a/packages/gitbook/src/components/StructurePreview/StructurePreview.tsx
+++ b/packages/gitbook/src/components/StructurePreview/StructurePreview.tsx
@@ -1,13 +1,18 @@
 'use client';
 
 import type { CustomizationContentLink, CustomizationHeaderItem, SiteSpace } from '@gitbook/api';
-import { CustomizationHeaderPreset } from '@gitbook/api';
+import {
+    CustomizationAIMode,
+    CustomizationHeaderPreset,
+    CustomizationSearchStyle,
+} from '@gitbook/api';
 import * as React from 'react';
 
 import { SiteSectionTabs } from '@/components/SiteSections';
 import { tcls } from '@/lib/tailwind';
 
 import { tString, useLanguage } from '@/intl/client';
+import { AIChatButtonView, AIChatIcon, getAIChatName } from '../AIChat';
 import {
     HeaderLinkItem,
     HeaderLinkMenuItem,
@@ -98,6 +103,7 @@ function StructurePreviewHeader(props: { snapshot: StructurePreviewSnapshot }) {
     const variants = getPreviewVariants(snapshot);
     const sections = encodePreviewSiteSections(snapshot);
     const headerSocialAccounts = getHeaderSocialAccounts(customization);
+    const previewAssistants = getPreviewAssistants(snapshot, language);
     const withTopHeader = customization.header.preset !== CustomizationHeaderPreset.None;
     const withSections = Boolean(
         sections &&
@@ -192,6 +198,20 @@ function StructurePreviewHeader(props: { snapshot: StructurePreviewSnapshot }) {
                             )}
                         >
                             <StructurePreviewSearch />
+                            {previewAssistants.map((assistant, index) => (
+                                <AIChatButtonView
+                                    key={assistant.id}
+                                    icon={assistant.icon}
+                                    label={assistant.label}
+                                    withShortcut={index === 0}
+                                    showLabel={
+                                        previewAssistants.length === 1 &&
+                                        customization.styling.search ===
+                                            CustomizationSearchStyle.Prominent
+                                    }
+                                    inert
+                                />
+                            ))}
                         </div>
 
                         {customization.header.links.length > 0 ||
@@ -315,6 +335,29 @@ function StructurePreviewSearch() {
     return <SearchHeaderInput interactive={false} />;
 }
 
+function getPreviewAssistants(
+    snapshot: StructurePreviewSnapshot,
+    language: ReturnType<typeof useLanguage>
+) {
+    if (snapshot.customization.ai?.mode !== CustomizationAIMode.Assistant) {
+        return [];
+    }
+
+    return [
+        {
+            id: 'gitbook-assistant',
+            label: getAIChatName(language, snapshot.customization.trademark.enabled),
+            icon: (
+                <AIChatIcon
+                    state="default"
+                    trademark={snapshot.customization.trademark.enabled}
+                    className="size-text-lg"
+                />
+            ),
+        },
+    ];
+}
+
 function StructurePreviewHeaderLink(props: {
     snapshot: StructurePreviewSnapshot;
     link: CustomizationHeaderItem;