Gate AI features server-actions (#4353)

packages/gitbook/src/app/sites/dynamic/[mode]/[siteURL]/[siteData]/~gitbook/embed/assistant/page.tsx

@@ -1,7 +1,7 @@
 import type { RouteLayoutParams } from '@/app/utils';
 import { EmbeddableAssistantPage } from '@/components/Embeddable';
+import { isAIChatEnabled } from '@/components/utils/isAIChatEnabled';
 import { getEmbeddableDynamicContext } from '@/lib/embeddable';
-import { CustomizationAIMode } from '@gitbook/api';
 import { redirect } from 'next/navigation';
 
 type PageProps = {
@@ -15,7 +15,7 @@ export default async function Page(props: PageProps) {
     const { context } = await getEmbeddableDynamicContext(params);
 
     // If the assistant is not enabled, redirect to the docs
-    if (context.customization.ai.mode !== CustomizationAIMode.Assistant) {
+    if (!isAIChatEnabled(context.customization.ai.mode)) {
         redirect(`${context.linker.toPathInSite('~gitbook/embed/page/')}`);
     }
 

packages/gitbook/src/app/sites/dynamic/[mode]/[siteURL]/[siteData]/~gitbook/embed/page.tsx

@@ -1,6 +1,6 @@
 import type { RouteLayoutParams } from '@/app/utils';
+import { isAIChatEnabled } from '@/components/utils/isAIChatEnabled';
 import { getEmbeddableDynamicContext } from '@/lib/embeddable';
-import { CustomizationAIMode } from '@gitbook/api';
 import { redirect } from 'next/navigation';
 
 type PageProps = {
@@ -13,7 +13,7 @@ export default async function Page(props: PageProps) {
     const baseURL = context.linker.toPathInSite('~gitbook/embed/');
 
     // If assistant is enabled, redirect to assistant, otherwise to docs
-    if (context.customization.ai.mode === CustomizationAIMode.Assistant) {
+    if (isAIChatEnabled(context.customization.ai.mode)) {
         redirect(`${baseURL}/assistant`);
     } else {
         redirect(`${baseURL}/page/`);

packages/gitbook/src/app/sites/static/[mode]/[siteURL]/[siteData]/~gitbook/embed/assistant/page.tsx

@@ -1,7 +1,7 @@
 import type { RouteParams } from '@/app/utils';
 import { EmbeddableAssistantPage } from '@/components/Embeddable';
+import { isAIChatEnabled } from '@/components/utils/isAIChatEnabled';
 import { getEmbeddableStaticContext } from '@/lib/embeddable';
-import { CustomizationAIMode } from '@gitbook/api';
 import { redirect } from 'next/navigation';
 
 export const dynamic = 'force-static';
@@ -15,7 +15,7 @@ export default async function Page(props: PageProps) {
     const { context } = await getEmbeddableStaticContext(params);
 
     // If the assistant is not enabled, redirect to the docs
-    if (context.customization.ai.mode !== CustomizationAIMode.Assistant) {
+    if (!isAIChatEnabled(context.customization.ai.mode)) {
         redirect(`${context.linker.toPathInSite('~gitbook/embed/page/')}`);
     }
 

packages/gitbook/src/app/sites/static/[mode]/[siteURL]/[siteData]/~gitbook/embed/page.tsx

@@ -1,6 +1,6 @@
 import type { RouteLayoutParams } from '@/app/utils';
+import { isAIChatEnabled } from '@/components/utils/isAIChatEnabled';
 import { getEmbeddableStaticContext } from '@/lib/embeddable';
-import { CustomizationAIMode } from '@gitbook/api';
 import { redirect } from 'next/navigation';
 
 export const dynamic = 'force-static';
@@ -15,7 +15,7 @@ export default async function Page(props: PageProps) {
     const baseURL = context.linker.toPathInSite('~gitbook/embed/');
 
     // If assistant is enabled, redirect to assistant, otherwise to docs
-    if (context.customization.ai.mode === CustomizationAIMode.Assistant) {
+    if (isAIChatEnabled(context.customization.ai.mode)) {
         redirect(`${baseURL}/assistant`);
     } else {
         redirect(`${baseURL}/page/`);

packages/gitbook/src/components/AI/server-actions/chat.ts

@@ -1,6 +1,7 @@
 'use server';
+import { isAIChatEnabled } from '@/components/utils/isAIChatEnabled';
 import { getSiteURLDataFromMiddleware } from '@/lib/middleware';
-import { getServerActionBaseContext } from '@/lib/server-actions';
+import { fetchServerActionSiteContext, getServerActionBaseContext } from '@/lib/server-actions';
 import { traceErrorOnly } from '@/lib/tracing';
 import {
     type AIMessageContext,
@@ -39,6 +40,11 @@ export async function* streamAIChatResponse({
             isEmbeddable: options?.asEmbeddable,
         });
 
+        const siteContext = await fetchServerActionSiteContext(context);
+        if (!isAIChatEnabled(siteContext.customization.ai.mode)) {
+            throw new Error('The AI Assistant is not enabled for this site.');
+        }
+
         const siteURLData = await getSiteURLDataFromMiddleware();
 
         const api = await context.dataFetcher.api();

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

@@ -1,10 +1,11 @@
 'use client';
 
-import { CustomizationAIMode } from '@gitbook/api';
+import type { CustomizationAIMode } from '@gitbook/api';
 import { Icon, type IconName } from '@gitbook/icons';
 import * as React from 'react';
 import type { ReactNode } from 'react';
 
+import { isAIChatEnabled, isAISearchEnabled } from '@/components/utils/isAIChatEnabled';
 import { tString, useLanguage } from '@/intl/client';
 import type { GitBookAssistant } from '@gitbook/browser-types';
 import { useAIChatController, useAIChatState } from '.';
@@ -89,7 +90,7 @@ export function useAI(): AIContext {
 
     const assistants: Assistant[] = [];
 
-    if (config.aiMode === CustomizationAIMode.Assistant) {
+    if (isAIChatEnabled(config.aiMode)) {
         assistants.push({
             id: 'gitbook-assistant',
             label: config.assistantName ?? getAIChatName(language, config.trademark),
@@ -110,7 +111,7 @@ export function useAI(): AIContext {
             ui: true,
             mode: 'sidebar',
         });
-    } else if (config.aiMode === CustomizationAIMode.Search) {
+    } else if (isAISearchEnabled(config.aiMode)) {
         assistants.push({
             id: 'gitbook-ai-search',
             label: tString(language, 'ai_chat_context_badge'),

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

@@ -1,6 +1,5 @@
 'use client';
 
-import { CustomizationAIMode } from '@gitbook/api';
 import fnv1a from '@sindresorhus/fnv1a';
 import { AnimatePresence, motion } from 'motion/react';
 import * as React from 'react';
@@ -10,6 +9,7 @@ import { useAIChatController, useAIConfig } from '@/components/AI';
 import { useIsMobile } from '@/components/hooks/useIsMobile';
 import { useIsMounted } from '@/components/hooks/useIsMounted';
 import { Button } from '@/components/primitives';
+import { isAIChatEnabled } from '@/components/utils/isAIChatEnabled';
 import { t, useLanguage } from '@/intl/client';
 
 import { AIChatIcon } from '../AIChatIcon';
@@ -31,7 +31,7 @@ export function AskAITextSelection() {
     const isMobile = useIsMobile();
     const isMounted = useIsMounted();
 
-    const enabled = config.aiMode === CustomizationAIMode.Assistant && !isMobile;
+    const enabled = isAIChatEnabled(config.aiMode) && !isMobile;
 
     const toolbarRef = React.useRef<HTMLDivElement>(null);
     const { selection, clear } = useStableTextSelection({

packages/gitbook/src/components/DocumentView/CodeBlock/AskAICodeButton.tsx

@@ -1,10 +1,9 @@
 'use client';
 
-import { CustomizationAIMode } from '@gitbook/api';
-
 import { useAIChatController, useAIConfig } from '@/components/AI';
 import { AIChatIcon } from '@/components/AIChat';
 import { Button } from '@/components/primitives';
+import { isAIChatEnabled } from '@/components/utils/isAIChatEnabled';
 import { t, useLanguage } from '@/intl/client';
 import { type ClassValue, tcls } from '@/lib/tailwind';
 
@@ -22,7 +21,7 @@ export function AskAICodeButton(props: {
     const config = useAIConfig();
     const chatController = useAIChatController();
 
-    if (config.aiMode !== CustomizationAIMode.Assistant) {
+    if (!isAIChatEnabled(config.aiMode)) {
         return null;
     }
 

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

@@ -4,8 +4,8 @@ import type { GitBookEmbeddableConfiguration, ParentToFrameMessage } from '@gitb
 import React, { useEffect, useRef } from 'react';
 
 import { useAI, useAIChatController } from '@/components/AI';
+import { isAIChatEnabled } from '@/components/utils/isAIChatEnabled';
 import { tString, useLanguage } from '@/intl/client';
-import { CustomizationAIMode } from '@gitbook/api';
 import { useRouter } from 'next/navigation';
 import { createStore, useStore } from 'zustand';
 import { integrationsAssistantTools } from '../Integrations';
@@ -186,9 +186,7 @@ export function EmbeddableIframeTabs(props: {
     const router = useRouter();
 
     const enabledTabs = [
-        config.aiMode === CustomizationAIMode.Assistant &&
-        assistants[0] &&
-        tabs.includes('assistant')
+        isAIChatEnabled(config.aiMode) && assistants[0] && tabs.includes('assistant')
             ? {
                   key: 'assistant',
                   label: assistants[0].label,

packages/gitbook/src/components/Search/server-actions.tsx

@@ -1,5 +1,6 @@
 'use server';
 
+import { isAIEnabled, isAISearchEnabled } from '@/components/utils/isAIChatEnabled';
 import type { GitBookSiteContext } from '@/lib/context';
 import { resolvePageId } from '@/lib/pages';
 import { fetchServerActionSiteContext, getServerActionBaseContext } from '@/lib/server-actions';
@@ -55,6 +56,10 @@ export async function streamAskQuestion({
                 await getServerActionBaseContext({ isEmbeddable: asEmbeddable })
             );
 
+            if (!isAISearchEnabled(context.customization.ai.mode)) {
+                throw new Error('AI Search is not enabled for this site.');
+            }
+
             const apiClient = await context.dataFetcher.api();
 
             const stream = apiClient.orgs.streamAskInSite(
@@ -151,6 +156,11 @@ export async function streamRecommendedQuestions(args: { siteSpaceId?: string })
         >();
 
         (async () => {
+            const siteContext = await fetchServerActionSiteContext(context);
+            if (!isAIEnabled(siteContext.customization.ai.mode)) {
+                throw new Error('AI is not enabled for this site.');
+            }
+
             const apiClient = await context.dataFetcher.api();
             const apiStream = apiClient.orgs.streamRecommendedQuestionsInSite(
                 siteURLData.organization,

packages/gitbook/src/components/SpaceLayout/SpaceLayout.tsx

@@ -1,14 +1,11 @@
 import type { GitBookSiteContext } from '@/lib/context';
-import {
-    CustomizationAIMode,
-    CustomizationHeaderPreset,
-    CustomizationSearchStyle,
-} from '@gitbook/api';
+import { CustomizationHeaderPreset, CustomizationSearchStyle } from '@gitbook/api';
 import type React from 'react';
 
 import { Footer } from '@/components/Footer';
 import { Header, HeaderLogo } from '@/components/Header';
 import { TableOfContents } from '@/components/TableOfContents';
+import { isAIChatEnabled } from '@/components/utils/isAIChatEnabled';
 import type { VisitorAuthClaims } from '@/lib/adaptive';
 import { GITBOOK_APP_URL } from '@/lib/env';
 import { tcls } from '@/lib/tailwind';
@@ -128,7 +125,7 @@ export function SpaceLayout(props: SpaceLayoutProps) {
             <Announcement context={context} />
             <Header withTopHeader={withTopHeader} variants={variants} context={context} />
             <NavigationLoader />
-            {customization.ai?.mode === CustomizationAIMode.Assistant ? (
+            {isAIChatEnabled(customization.ai?.mode) ? (
                 <>
                     <AIChat />
                     <AskAITextSelection />

packages/gitbook/src/components/utils/isAIChatEnabled.ts

@@ -1,5 +1,11 @@
-import type { GitBookSiteContext } from '@/lib/context';
 import { CustomizationAIMode } from '@gitbook/api';
 
-export const isAIChatEnabled = (context: GitBookSiteContext) =>
-    context.customization.ai.mode === CustomizationAIMode.Assistant;
+// `ai.mode` is exclusive: Assistant (chat) and Search (ask) are distinct features; None means no AI.
+export const isAIChatEnabled = (mode: CustomizationAIMode | undefined) =>
+    mode === CustomizationAIMode.Assistant;
+
+export const isAISearchEnabled = (mode: CustomizationAIMode | undefined) =>
+    mode === CustomizationAIMode.Search;
+
+export const isAIEnabled = (mode: CustomizationAIMode | undefined) =>
+    isAIChatEnabled(mode) || isAISearchEnabled(mode);

packages/gitbook/src/routes/llms.ts

@@ -1,3 +1,4 @@
+import { isAIEnabled } from '@/components/utils/isAIChatEnabled';
 import { type GitBookSiteContext, checkIsRootSiteContext } from '@/lib/context';
 import { throwIfDataError } from '@/lib/data';
 import { type GitBookLinker, linkerWithMarkdownPages } from '@/lib/links';
@@ -203,6 +204,10 @@ export async function getMarkdownForPagesTree(
 }
 
 function renderAskFooter(context: GitBookSiteContext) {
+    if (!isAIEnabled(context.customization.ai.mode)) {
+        return '';
+    }
+
     return `\n\n---\n\n# Agent Instructions
 This documentation is published with GitBook. GitBook is the documentation platform designed so that both humans and AI agents can read, navigate, and reason over technical content effectively. Learn more at gitbook.com.
 

packages/gitbook/src/routes/markdownAsk.ts

@@ -1,3 +1,4 @@
+import { isAIEnabled } from '@/components/utils/isAIChatEnabled';
 import type { GitBookSiteContext } from '@/lib/context';
 import { throwIfDataError } from '@/lib/data';
 import { linkerWithMarkdownPages } from '@/lib/links';
@@ -27,6 +28,10 @@ export async function serveAskMarkdown(
     rawQuestion: string,
     options: ServeAskMarkdownOptions = {}
 ) {
+    if (!isAIEnabled(context.customization.ai.mode)) {
+        return new Response('Not Found', { status: 404 });
+    }
+
     return serveMarkdown(async () => {
         const question = rawQuestion.trim();
         const goal = options.goal?.trim() || undefined;

packages/gitbook/src/routes/markdownPage.ts

@@ -1,3 +1,4 @@
+import { isAIEnabled } from '@/components/utils/isAIChatEnabled';
 import type { GitBookSiteContext } from '@/lib/context';
 import { getExposableError } from '@/lib/data';
 import { linkerWithMarkdownPages } from '@/lib/links';
@@ -100,6 +101,10 @@ function renderAskFooter(
     context: GitBookSiteContext,
     pageLookup: ResolvedPagePath<RevisionPageDocument | RevisionPageGroup>
 ) {
+    if (!isAIEnabled(context.customization.ai.mode)) {
+        return '';
+    }
+
     const pageUrl = context.linker.toAbsoluteURL(
         context.linker.toPathForPage({
             page: pageLookup.page,