diff --git a/packages/gitbook/src/components/Search/SearchAskAnswer.tsx b/packages/gitbook/src/components/Search/SearchAskAnswer.tsx index 44a49ad5d0..d21ef6a898 100644 --- a/packages/gitbook/src/components/Search/SearchAskAnswer.tsx +++ b/packages/gitbook/src/components/Search/SearchAskAnswer.tsx @@ -167,15 +167,18 @@ function AnswerBody(props: { query: string; answer: AskAnswerResult }) { <>
{answer.body ?? t(language, 'search_ask_no_answer')} - {answer.sources.length > 0 ? ( + {answer.error ? ( +
{t(language, 'search_ask_error')}
+ ) : null} + {!answer.error && answer.sources.length > 0 ? ( // @TODO: Add responseId once search uses new AI endpoint ) : null} - {answer.followupQuestions.length > 0 ? ( + {!answer.error && answer.followupQuestions.length > 0 ? ( ) : null}
- {answer.sources.length > 0 ? ( + {!answer.error && answer.sources.length > 0 ? ( { - const responseStream = createStreamableValue(); - - (async () => { - const context = await fetchServerActionSiteContext( - await getServerActionBaseContext({ isEmbeddable: asEmbeddable }) - ); - - const apiClient = await context.dataFetcher.api(); + return runStreamableServerAction({ + onError: (_, lastValue) => ({ + ...(lastValue ?? { followupQuestions: [], sources: [] }), + error: true, + }), + run: async (push) => { + const context = await fetchServerActionSiteContext( + await getServerActionBaseContext({ isEmbeddable: asEmbeddable }) + ); - const stream = apiClient.orgs.streamAskInSite( - context.organizationId, - context.site.id, - { - question, - context: { - siteSpaceId: context.siteSpace.id, - }, - scope: { - mode: 'default', - currentSiteSpace: context.siteSpace.id, + 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', + currentSiteSpace: context.siteSpace.id, + }, }, - }, - { format: 'document' } - ); - - const spacePromises = new Map>(); - 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, - throwIfDataError( - context.dataFetcher.getRevision({ - spaceId: source.space, - revisionId: source.revision, - }) - ) - ); - } - - 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 revision = await spacePromises.get(space); - return { space, pages: revision?.pages }; - }) - ).then((results) => { - return results.reduce((map, result) => { - if (result.pages) { - map.set(result.space, result.pages); - } - return map; - }, new Map()); - }); - responseStream.update( - await transformAnswer(context, { - answer: chunk.answer, - asEmbeddable: Boolean(asEmbeddable), - spacePages: pages, - }) + { format: 'document' } ); - } - })() - .then(() => { - responseStream.done(); - }) - .catch((error) => { - responseStream.error(error); - }); - return { - stream: responseStream.value, - }; + const spacePromises = new Map>(); + 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, + throwIfDataError( + context.dataFetcher.getRevision({ + spaceId: source.space, + revisionId: source.revision, + }) + ) + ); + } + + 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 revision = await spacePromises.get(space); + return { space, pages: revision?.pages }; + }) + ).then((results) => { + return results.reduce((map, result) => { + if (result.pages) { + map.set(result.space, result.pages); + } + return map; + }, new Map()); + }); + + push( + await transformAnswer(context, { + answer: chunk.answer, + asEmbeddable: Boolean(asEmbeddable), + spacePages: pages, + }) + ); + } + }, + }); }); } @@ -142,32 +142,25 @@ export async function streamRecommendedQuestions(args: { siteSpaceId?: string }) const siteURLData = await getSiteURLDataFromMiddleware(); const context = await getServerActionBaseContext(); - const responseStream = createStreamableValue< - SearchAIRecommendedQuestionStream | undefined - >(); + return runStreamableServerAction({ + // On mid-stream error, pass the last value through to stop cleanly without a throw. + // On pre-stream error, fail() is called so the existing silent catch in the client handles it. + onError: (_, lastValue) => lastValue, + run: async (push) => { + const apiClient = await context.dataFetcher.api(); + const apiStream = apiClient.orgs.streamRecommendedQuestionsInSite( + siteURLData.organization, + siteURLData.site, + { + siteSpaceId: args.siteSpaceId, + } + ); - (async () => { - const apiClient = await context.dataFetcher.api(); - const apiStream = apiClient.orgs.streamRecommendedQuestionsInSite( - siteURLData.organization, - siteURLData.site, - { - siteSpaceId: args.siteSpaceId, + for await (const chunk of apiStream) { + push(chunk); } - ); - - for await (const chunk of apiStream) { - responseStream.update(chunk); - } - })() - .then(() => { - responseStream.done(); - }) - .catch((error) => { - responseStream.error(error); - }); - - return { stream: responseStream.value }; + }, + }); }); } diff --git a/packages/gitbook/src/lib/graceful-stream.ts b/packages/gitbook/src/lib/graceful-stream.ts new file mode 100644 index 0000000000..51942273a0 --- /dev/null +++ b/packages/gitbook/src/lib/graceful-stream.ts @@ -0,0 +1,43 @@ +/** + * Creates a mid-stream error handler that gracefully downgrades errors to values + * when the stream has already started delivering content. + * + * - If the stream has started: calls `update(onError(error, lastValue))` (when non-null) + * then `done()`, preserving partial content on the client. + * - If the stream has not started: calls `fail(error)`, propagating the error normally + * so the client receives a full error state. + */ +export function createMidStreamErrorHandler( + onError: (error: unknown, lastValue: T | undefined) => T | undefined +): { + track: (value: T) => void; + handleError: ( + error: unknown, + callbacks: { + update: (value: T) => void; + done: () => void; + fail: (error: unknown) => void; + } + ) => void; +} { + let hasStarted = false; + let lastValue: T | undefined; + + return { + track(value) { + hasStarted = true; + lastValue = value; + }, + handleError(error, { update, done, fail }) { + if (hasStarted) { + const errorValue = onError(error, lastValue); + if (errorValue !== undefined) { + update(errorValue); + } + done(); + } else { + fail(error); + } + }, + }; +} diff --git a/packages/gitbook/src/lib/server-actions.ts b/packages/gitbook/src/lib/server-actions.ts index 0e8c8ed937..52a73a9662 100644 --- a/packages/gitbook/src/lib/server-actions.ts +++ b/packages/gitbook/src/lib/server-actions.ts @@ -1,5 +1,8 @@ +import { createStreamableValue } from 'ai/rsc'; +import type { StreamableValue } from 'ai/rsc'; import { type GitBookBaseContext, fetchSiteContextByURLLookup, getBaseContext } from './context'; import { getEmbeddableLinker } from './embeddable-linker'; +import { createMidStreamErrorHandler } from './graceful-stream'; import { getSiteURLDataFromMiddleware, getSiteURLFromMiddleware, @@ -39,3 +42,41 @@ export async function fetchServerActionSiteContext(baseContext: GitBookBaseConte const siteURLData = await getSiteURLDataFromMiddleware(); return fetchSiteContextByURLLookup(baseContext, siteURLData); } + +/** + * Run a server action that streams values to the client using `createStreamableValue`. + * + * When an error occurs after the stream has started delivering content, it is + * converted into a final value via `onError` and the stream closes cleanly — + * preserving partial content on the client. + * + * When an error occurs before any value has been pushed, it is propagated + * normally so the client receives a full error state. + */ +export function runStreamableServerAction({ + onError, + run, +}: { + onError: (error: unknown, lastValue: T | undefined) => T | undefined; + run: (push: (value: T) => void) => Promise; +}): { stream: StreamableValue } { + const responseStream = createStreamableValue(); + const errorHandler = createMidStreamErrorHandler(onError); + + run((value) => { + errorHandler.track(value); + responseStream.update(value); + }) + .then(() => { + responseStream.done(); + }) + .catch((error) => { + errorHandler.handleError(error, { + update: (value) => responseStream.update(value), + done: () => responseStream.done(), + fail: (err) => responseStream.error(err), + }); + }); + + return { stream: responseStream.value }; +}