Make the streaming work

packages/gitbook-v2/src/lib/context.ts

@@ -174,8 +174,7 @@ export async function fetchSiteContextByURL(
     const context = await fetchSiteContextByIds(
         {
             ...baseContext,
-            dataFetcher: createDataFetcher({
-                apiEndpoint: dataFetcher.apiEndpoint,
+            dataFetcher: dataFetcher.withToken({
                 apiToken: data.apiToken,
             }),
         },

packages/gitbook-v2/src/lib/data/api.ts

@@ -39,6 +39,13 @@ export function createDataFetcher(input: DataFetcherInput = commonInput): GitBoo
             return getAPI(input);
         },
 
+        withToken({ apiToken }) {
+            return createDataFetcher({
+                ...input,
+                apiToken,
+            });
+        },
+
         //
         // API that are tied to the token
         //

packages/gitbook-v2/src/lib/data/types.ts

@@ -15,6 +15,13 @@ export interface GitBookDataFetcher {
      */
     api(): Promise<api.GitBookAPI>;
 
+    /**
+     * Create a data fetcher authenticated with a specific token.
+     */
+    withToken(input: {
+        apiToken: string;
+    }): GitBookDataFetcher;
+
     /**
      * Get a user by its ID.
      */

packages/gitbook/src/components/Search/SearchAskAnswer.tsx

@@ -1,13 +1,13 @@
 'use client';
 
 import { Icon } from '@gitbook/icons';
+import { readStreamableValue } from 'ai/rsc';
 import React from 'react';
 
 import { Loading } from '@/components/primitives';
 import { useLanguage } from '@/intl/client';
 import { t } from '@/intl/translate';
 import type { TranslationLanguage } from '@/intl/translations';
-import { iterateStreamResponse } from '@/lib/actions';
 import { tcls } from '@/lib/tailwind';
 
 import { useTrackEvent } from '../Insights';
@@ -50,19 +50,19 @@ export function SearchAskAnswer(props: { query: string }) {
                 query,
             });
 
-            const response = streamAskQuestion({ question: query });
-            const stream = iterateStreamResponse(response);
-
             // When we pass in "ask" mode, the query could still be updated by the client
             // we ensure that the query is up-to-date before starting the stream.
             setSearchState((prev) => (prev ? { ...prev, query, ask: true } : null));
 
-            for await (const chunk of stream) {
+            const { stream } = await streamAskQuestion({ question: query });
+            for await (const chunk of readStreamableValue(stream)) {
                 if (cancelled) {
                     return;
                 }
 
-                setAskState({ type: 'answer', answer: chunk });
+                if (chunk) {
+                    setAskState({ type: 'answer', answer: chunk });
+                }
             }
         })().catch(() => {
             if (cancelled) {

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

@@ -13,7 +13,6 @@ import { fetchServerActionSiteContext, getServerActionBaseContext } from '@v2/li
 import { createStreamableValue } from 'ai/rsc';
 import type * as React from 'react';
 
-import { streamResponse } from '@/lib/actions';
 import { getAbsoluteHref } from '@/lib/links';
 import { resolvePageId } from '@/lib/pages';
 import { findSiteSpaceById } from '@/lib/sites';
@@ -93,77 +92,93 @@ export async function searchSiteSpaceContent(query: string): Promise<OrderedComp
 /**
  * Server action to ask a question in a space.
  */
-export const streamAskQuestion = streamResponse(async function* ({
+export async function streamAskQuestion({
     question,
 }: {
     question: string;
 }) {
-    const context = await fetchServerActionSiteContext(
-        isV2() ? await getServerActionBaseContext() : await getV1BaseContext()
-    );
-
-    const apiClient = await context.dataFetcher.api();
-
-    const stream = apiClient.orgs.streamAskInSite(
-        context.organizationId,
-        context.site.id,
-        {
-            question,
-            context: {
-                siteSpaceId: context.siteSpace.id,
-            },
-            scope: {
-                mode: 'default',
-                // Include the current site space regardless.
-                includedSiteSpaces: [context.siteSpace.id],
-            },
-        },
-        { format: 'document' }
-    );
+    const responseStream = createStreamableValue<AskAnswerResult | undefined>();
 
-    const spacePromises = new Map<string, Promise<RevisionPage[]>>();
-    for await (const chunk of stream) {
-        const answer = chunk.answer;
-
-        // Register the space of each page source into the promise queue.
-        const spaces = answer.sources
-            .map((source) => {
-                if (source.type !== 'page') {
-                    return null;
-                }
+    (async () => {
+        const context = await fetchServerActionSiteContext(
+            isV2() ? await getServerActionBaseContext() : await getV1BaseContext()
+        );
 
-                if (!spacePromises.has(source.space)) {
-                    spacePromises.set(
-                        source.space,
-                        context.dataFetcher.getRevisionPages({
-                            spaceId: source.space,
-                            revisionId: source.revision,
-                            metadata: false,
-                        })
-                    );
-                }
+        const apiClient = await context.dataFetcher.api();
 
-                return source.space;
-            })
-            .filter(filterOutNullable);
+        const stream = apiClient.orgs.streamAskInSite(
+            context.organizationId,
+            context.site.id,
+            {
+                question,
+                context: {
+                    siteSpaceId: context.siteSpace.id,
+                },
+                scope: {
+                    mode: 'default',
+                    // Include the current site space regardless.
+                    includedSiteSpaces: [context.siteSpace.id],
+                },
+            },
+            { format: 'document' }
+        );
 
-        // Get the pages for all spaces referenced by this answer.
-        const pages = await Promise.all(
-            spaces.map(async (space) => {
-                const pages = await spacePromises.get(space);
-                return { space, pages };
-            })
-        ).then((results) => {
-            return results.reduce((map, result) => {
-                if (result.pages) {
-                    map.set(result.space, result.pages);
-                }
-                return map;
-            }, new Map<string, RevisionPage[]>());
+        const spacePromises = new Map<string, Promise<RevisionPage[]>>();
+        for await (const chunk of stream) {
+            const answer = chunk.answer;
+
+            // Register the space of each page source into the promise queue.
+            const spaces = answer.sources
+                .map((source) => {
+                    if (source.type !== 'page') {
+                        return null;
+                    }
+
+                    if (!spacePromises.has(source.space)) {
+                        spacePromises.set(
+                            source.space,
+                            context.dataFetcher.getRevisionPages({
+                                spaceId: source.space,
+                                revisionId: source.revision,
+                                metadata: false,
+                            })
+                        );
+                    }
+
+                    return source.space;
+                })
+                .filter(filterOutNullable);
+
+            // Get the pages for all spaces referenced by this answer.
+            const pages = await Promise.all(
+                spaces.map(async (space) => {
+                    const pages = await spacePromises.get(space);
+                    return { space, pages };
+                })
+            ).then((results) => {
+                return results.reduce((map, result) => {
+                    if (result.pages) {
+                        map.set(result.space, result.pages);
+                    }
+                    return map;
+                }, new Map<string, RevisionPage[]>());
+            });
+            responseStream.update(
+                await transformAnswer(context, { answer: chunk.answer, spacePages: pages })
+            );
+        }
+    })()
+        .then(() => {
+            responseStream.done();
+        })
+        .catch((error) => {
+            responseStream.error(error);
         });
-        yield await transformAnswer(context, { answer: chunk.answer, spacePages: pages });
-    }
-});
+
+    return {
+        stream: responseStream.value,
+    };
+}
 
 /**
  * Stream a list of suggested questions for the site.
@@ -173,7 +188,7 @@ export async function streamRecommendedQuestions() {
         isV2() ? await getServerActionBaseContext() : await getV1BaseContext()
     );
 
-    const stream = createStreamableValue<SearchAIRecommendedQuestionStream | undefined>();
+    const responseStream = createStreamableValue<SearchAIRecommendedQuestionStream | undefined>();
 
     (async () => {
         const apiClient = await context.dataFetcher.api();
@@ -183,20 +198,17 @@ export async function streamRecommendedQuestions() {
         );
 
         for await (const chunk of apiStream) {
-            console.log('chunk', chunk);
-            stream.update(chunk);
+            responseStream.update(chunk);
         }
     })()
         .then(() => {
-            console.log('done');
-            stream.done();
+            responseStream.done();
         })
         .catch((error) => {
-            console.log('error', error);
-            stream.error(error);
+            responseStream.error(error);
         });
 
-    return { stream: stream.value };
+    return { stream: responseStream.value };
 }
 
 /**

packages/gitbook/src/lib/v1.ts

@@ -70,20 +70,27 @@ export async function getV1BaseContext(): Promise<GitBookBaseContext> {
 async function getDataFetcherV1(): Promise<GitBookDataFetcher> {
     const apiClient = await api();
 
-    return {
+    const dataFetcher: GitBookDataFetcher = {
         apiEndpoint: apiClient.client.endpoint,
 
         async api() {
             const result = await api();
             return result.client;
         },
 
+        withToken() {
+            // In v1, the token is global and controlled by the middleware.
+            // We don't need to do anything special here.
+            return dataFetcher;
+        },
+
         getUserById(userId) {
             return getUserById(userId);
         },
 
         // @ts-ignore - types are compatible enough, and this will not be called in v1 this way
         getPublishedContentByUrl(params) {
+            console.log('getPublishedContentByUrl', params);
             return getPublishedContentByUrl(
                 params.url,
                 params.visitorAuthToken ?? undefined,
@@ -147,6 +154,8 @@ async function getDataFetcherV1(): Promise<GitBookDataFetcher> {
             return getEmbedByUrlInSpace(params.spaceId, params.url);
         },
     };
+
+    return dataFetcher;
 }
 
 /**