feat: add tracing to server actions (#3543)

Co-authored-by: Nicolas Dorseuil <nicolas@gitbook.io>

packages/gitbook/src/components/AI/server-actions/api.tsx

@@ -1,6 +1,7 @@
 'use server';
 import type { GitBookBaseContext } from '@/lib/context';
 import { fetchServerActionSiteContext } from '@/lib/server-actions';
+import { traceErrorOnly } from '@/lib/tracing';
 import {
     type AIMessage,
     AIMessageRole,
@@ -19,97 +20,100 @@ export async function streamRenderAIMessage(
     rawStream: AsyncIterable<AIStreamResponse>,
     options?: RenderAIMessageOptions
 ) {
-    const message: AIMessage = {
-        id: '',
-        role: AIMessageRole.Assistant,
-        steps: [],
-    };
-
-    const updateProcessingMessageStep = (
-        stepIndex: number,
-        callback: (step: AIMessageStep) => void
-    ) => {
-        if (stepIndex > message.steps.length) {
-            throw new Error(
-                `Step index out of bounds ${stepIndex} (${message.steps.length} steps)`
-            );
-        }
-
-        if (message.steps[stepIndex]) {
-            message.steps = [...message.steps];
-            message.steps[stepIndex] = { ...message.steps[stepIndex] };
-            callback(message.steps[stepIndex]);
-        } else {
-            message.steps = [
-                ...message.steps,
-                {
-                    content: {
-                        object: 'document',
-                        data: {},
-                        nodes: [],
-                    },
-                },
-            ];
-            callback(message.steps[stepIndex]);
-        }
-    };
-
-    // Fetch the full-context in the background to avoid blocking the stream.
-    const promiseContext = fetchServerActionSiteContext(baseContext);
-
-    return parseResponse<{
-        content: React.ReactNode;
-        event: AIStreamResponse;
-    }>(rawStream, async (event) => {
-        switch (event.type) {
-            /**
-             * The agent is processing a tool call in a new message.
-             */
-            case 'response_tool_call': {
-                updateProcessingMessageStep(event.stepIndex, (step) => {
-                    step.toolCalls ??= [];
-                    step.toolCalls.push(event.toolCall);
-                });
-                break;
-            }
-
-            /**
-             * The agent is writing the content of a new message.
-             */
-            case 'response_reasoning':
-            case 'response_document': {
-                updateProcessingMessageStep(event.stepIndex, (step) => {
-                    const container = event.type === 'response_reasoning' ? 'reasoning' : 'content';
-
-                    step[container] ??= {
-                        object: 'document',
-                        data: {},
-                        nodes: [],
-                    };
-                    step[container] = {
-                        ...step[container],
-                        nodes: [...step[container].nodes],
-                    };
-                    if (event.operation === 'insert') {
-                        step[container].nodes.push(...event.blocks);
-                    } else {
-                        step[container].nodes.splice(
-                            -event.blocks.length,
-                            event.blocks.length,
-                            ...event.blocks
-                        );
-                    }
-                });
-                break;
-            }
-        }
-
-        return {
-            event,
-            content: (
-                <AIMessageView message={message} context={await promiseContext} {...options} />
-            ),
+    return traceErrorOnly('AI.streamRenderAIMessage', async () => {
+        const message: AIMessage = {
+            id: '',
+            role: AIMessageRole.Assistant,
+            steps: [],
         };
+
+        const updateProcessingMessageStep = (
+            stepIndex: number,
+            callback: (step: AIMessageStep) => void
+        ) => {
+            if (stepIndex > message.steps.length) {
+                throw new Error(
+                    `Step index out of bounds ${stepIndex} (${message.steps.length} steps)`
+                );
+            }
+
+            if (message.steps[stepIndex]) {
+                message.steps = [...message.steps];
+                message.steps[stepIndex] = { ...message.steps[stepIndex] };
+                callback(message.steps[stepIndex]);
+            } else {
+                message.steps = [
+                    ...message.steps,
+                    {
+                        content: {
+                            object: 'document',
+                            data: {},
+                            nodes: [],
+                        },
+                    },
+                ];
+                callback(message.steps[stepIndex]);
+            }
+        };
+
+        // Fetch the full-context in the background to avoid blocking the stream.
+        const promiseContext = fetchServerActionSiteContext(baseContext);
+
+        return parseResponse<{
+            content: React.ReactNode;
+            event: AIStreamResponse;
+        }>(rawStream, async (event) => {
+            switch (event.type) {
+                /**
+                 * The agent is processing a tool call in a new message.
+                 */
+                case 'response_tool_call': {
+                    updateProcessingMessageStep(event.stepIndex, (step) => {
+                        step.toolCalls ??= [];
+                        step.toolCalls.push(event.toolCall);
+                    });
+                    break;
+                }
+
+                /**
+                 * The agent is writing the content of a new message.
+                 */
+                case 'response_reasoning':
+                case 'response_document': {
+                    updateProcessingMessageStep(event.stepIndex, (step) => {
+                        const container =
+                            event.type === 'response_reasoning' ? 'reasoning' : 'content';
+
+                        step[container] ??= {
+                            object: 'document',
+                            data: {},
+                            nodes: [],
+                        };
+                        step[container] = {
+                            ...step[container],
+                            nodes: [...step[container].nodes],
+                        };
+                        if (event.operation === 'insert') {
+                            step[container].nodes.push(...event.blocks);
+                        } else {
+                            step[container].nodes.splice(
+                                -event.blocks.length,
+                                event.blocks.length,
+                                ...event.blocks
+                            );
+                        }
+                    });
+                    break;
+                }
+            }
+
+            return {
+                event,
+                content: (
+                    <AIMessageView message={message} context={await promiseContext} {...options} />
+                ),
+            };
+        });
     });
 }
 

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

@@ -1,6 +1,7 @@
 'use server';
 import { getSiteURLDataFromMiddleware } from '@/lib/middleware';
 import { getServerActionBaseContext } from '@/lib/server-actions';
+import { traceErrorOnly } from '@/lib/tracing';
 import { type AIMessageContext, AIMessageRole, AIModel } from '@gitbook/api';
 import { streamRenderAIMessage } from './api';
 import type { RenderAIMessageOptions } from './types';
@@ -19,25 +20,31 @@ export async function* streamAIChatResponse({
     previousResponseId?: string;
     options?: RenderAIMessageOptions;
 }) {
-    const context = await getServerActionBaseContext();
-    const siteURLData = await getSiteURLDataFromMiddleware();
+    const { stream } = await traceErrorOnly('AI.streamAIChatResponse', async () => {
+        const context = await getServerActionBaseContext();
+        const siteURLData = await getSiteURLDataFromMiddleware();
 
-    const api = await context.dataFetcher.api();
-    const rawStream = api.orgs.streamAiResponseInSite(siteURLData.organization, siteURLData.site, {
-        mode: 'assistant',
-        input: [
+        const api = await context.dataFetcher.api();
+        const rawStream = api.orgs.streamAiResponseInSite(
+            siteURLData.organization,
+            siteURLData.site,
             {
-                role: AIMessageRole.User,
-                content: message,
-                context: messageContext,
-            },
-        ],
-        output: { type: 'document' },
-        model: AIModel.ReasoningLow,
-        previousResponseId,
-    });
+                mode: 'assistant',
+                input: [
+                    {
+                        role: AIMessageRole.User,
+                        content: message,
+                        context: messageContext,
+                    },
+                ],
+                output: { type: 'document' },
+                model: AIModel.ReasoningLow,
+                previousResponseId,
+            }
+        );
 
-    const { stream } = await streamRenderAIMessage(context, rawStream, options);
+        return await streamRenderAIMessage(context, rawStream, options);
+    });
 
     for await (const output of stream) {
         yield output;

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

@@ -1,6 +1,7 @@
 'use server';
 import { getSiteURLDataFromMiddleware } from '@/lib/middleware';
 import { getServerActionBaseContext } from '@/lib/server-actions';
+import { traceErrorOnly } from '@/lib/tracing';
 import { streamRenderAIMessage } from './api';
 import type { RenderAIMessageOptions } from './types';
 
@@ -14,16 +15,18 @@ export async function* streamAIResponseById({
     responseId: string;
     options?: RenderAIMessageOptions;
 }) {
-    const context = await getServerActionBaseContext();
-    const siteURLData = await getSiteURLDataFromMiddleware();
+    const { stream } = await traceErrorOnly('AI.streamAIResponseById', async () => {
+        const context = await getServerActionBaseContext();
+        const siteURLData = await getSiteURLDataFromMiddleware();
 
-    const api = await context.dataFetcher.api();
-    const rawStream = api.orgs.streamExistingAiResponseInSite(
-        siteURLData.organization,
-        siteURLData.site,
-        responseId
-    );
-    const { stream } = await streamRenderAIMessage(context, rawStream, options);
+        const api = await context.dataFetcher.api();
+        const rawStream = api.orgs.streamExistingAiResponseInSite(
+            siteURLData.organization,
+            siteURLData.site,
+            responseId
+        );
+        return await streamRenderAIMessage(context, rawStream, options);
+    });
 
     for await (const output of stream) {
         yield output;

packages/gitbook/src/components/Ads/renderAd.tsx

@@ -4,6 +4,7 @@ import type { SiteInsightsAd, SiteInsightsAdPlacement } from '@gitbook/api';
 import { headers } from 'next/headers';
 
 import { getServerActionBaseContext } from '@/lib/server-actions';
+import { traceErrorOnly } from '@/lib/tracing';
 import { AdClassicRendering } from './AdClassicRendering';
 import { AdCoverRendering } from './AdCoverRendering';
 import { AdPixels } from './AdPixels';
@@ -40,40 +41,42 @@ interface FetchPlaceholderAdOptions {
  * and properly access user-agent and IP.
  */
 export async function renderAd(options: FetchAdOptions) {
-    const [context, result] = await Promise.all([
-        getServerActionBaseContext(),
-        options.source === 'live' ? fetchAd(options) : getPlaceholderAd(),
-    ]);
+    return traceErrorOnly('Ads.renderAd', async () => {
+        const [context, result] = await Promise.all([
+            getServerActionBaseContext(),
+            options.source === 'live' ? fetchAd(options) : getPlaceholderAd(),
+        ]);
 
-    const mode = options.source === 'live' ? options.mode : 'classic';
-    if (!result || !result.ad.description || !result.ad.statlink) {
-        return null;
-    }
+        const mode = options.source === 'live' ? options.mode : 'classic';
+        if (!result || !result.ad.description || !result.ad.statlink) {
+            return null;
+        }
 
-    const { ad } = result;
+        const { ad } = result;
 
-    const insightsAd: SiteInsightsAd | null =
-        options.source === 'live'
-            ? {
-                  placement: options.placement,
-                  zoneId: options.zoneId,
-                  domain: 'company' in ad ? ad.company : '',
-              }
-            : null;
+        const insightsAd: SiteInsightsAd | null =
+            options.source === 'live'
+                ? {
+                      placement: options.placement,
+                      zoneId: options.zoneId,
+                      domain: 'company' in ad ? ad.company : '',
+                  }
+                : null;
 
-    return {
-        children: (
-            <>
-                {mode === 'classic' || !('callToAction' in ad) ? (
-                    <AdClassicRendering ad={ad} insightsAd={insightsAd} context={context} />
-                ) : (
-                    <AdCoverRendering ad={ad} insightsAd={insightsAd} context={context} />
-                )}
-                {ad.pixel ? <AdPixels rawPixel={ad.pixel} /> : null}
-            </>
-        ),
-        insightsAd,
-    };
+        return {
+            children: (
+                <>
+                    {mode === 'classic' || !('callToAction' in ad) ? (
+                        <AdClassicRendering ad={ad} insightsAd={insightsAd} context={context} />
+                    ) : (
+                        <AdCoverRendering ad={ad} insightsAd={insightsAd} context={context} />
+                    )}
+                    {ad.pixel ? <AdPixels rawPixel={ad.pixel} /> : null}
+                </>
+            ),
+            insightsAd,
+        };
+    });
 }
 
 async function fetchAd({

packages/gitbook/src/components/AutoRefreshContent/server-actions.ts

@@ -2,6 +2,7 @@
 
 import { getDataOrNull } from '@/lib/data';
 import { getServerActionBaseContext } from '@/lib/server-actions';
+import { traceErrorOnly } from '@/lib/tracing';
 
 /**
  * Return true if a change-request has been updated.
@@ -11,15 +12,17 @@ export async function hasContentBeenUpdated(props: {
     changeRequestId: string;
     revisionId: string;
 }) {
-    const context = await getServerActionBaseContext();
-    const changeRequest = await getDataOrNull(
-        context.dataFetcher.getChangeRequest({
-            spaceId: props.spaceId,
-            changeRequestId: props.changeRequestId,
-        })
-    );
-    if (!changeRequest) {
-        return false;
-    }
-    return changeRequest.revision !== props.revisionId;
+    return traceErrorOnly('AutoRefreshContent.hasContentBeenUpdated', async () => {
+        const context = await getServerActionBaseContext();
+        const changeRequest = await getDataOrNull(
+            context.dataFetcher.getChangeRequest({
+                spaceId: props.spaceId,
+                changeRequestId: props.changeRequestId,
+            })
+        );
+        if (!changeRequest) {
+            return false;
+        }
+        return changeRequest.revision !== props.revisionId;
+    });
 }

packages/gitbook/src/components/DocumentView/Integration/server-actions.tsx

@@ -1,6 +1,7 @@
 'use server';
 
 import { getServerActionBaseContext } from '@/lib/server-actions';
+import { traceErrorOnly } from '@/lib/tracing';
 import type { RenderIntegrationUI } from '@gitbook/api';
 import { ContentKitOutput } from '@gitbook/react-contentkit';
 import { contentKitServerContext } from './contentkit';
@@ -19,20 +20,22 @@ export async function renderIntegrationUi({
     };
     request: RenderIntegrationUI;
 }) {
-    const serverAction = await getServerActionBaseContext();
-    const output = await fetchSafeIntegrationUI(serverAction, {
-        integrationName: renderContext.integrationName,
-        request,
-    });
+    return traceErrorOnly('DocumentView.renderIntegrationUi', async () => {
+        const serverAction = await getServerActionBaseContext();
+        const output = await fetchSafeIntegrationUI(serverAction, {
+            integrationName: renderContext.integrationName,
+            request,
+        });
+
+        if (output.error) {
+            return {
+                error: output.error.message,
+            };
+        }
 
-    if (output.error) {
         return {
-            error: output.error.message,
+            children: <ContentKitOutput output={output.data} context={contentKitServerContext} />,
+            output: output.data,
         };
-    }
-
-    return {
-        children: <ContentKitOutput output={output.data} context={contentKitServerContext} />,
-        output: output.data,
-    };
+    });
 }

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

@@ -22,6 +22,7 @@ import type * as React from 'react';
 import { throwIfDataError } from '@/lib/data';
 import { getSiteURLDataFromMiddleware } from '@/lib/middleware';
 import { joinPathWithBaseURL } from '@/lib/paths';
+import { traceErrorOnly } from '@/lib/tracing';
 import type { IconName } from '@gitbook/icons';
 import { DocumentView } from '../DocumentView';
 
@@ -68,11 +69,13 @@ export interface AskAnswerResult {
  * Server action to search content in the entire site.
  */
 export async function searchAllSiteContent(query: string): Promise<OrderedComputedResult[]> {
-    const context = await getServerActionBaseContext();
+    return traceErrorOnly('Search.searchAllSiteContent', async () => {
+        const context = await getServerActionBaseContext();
 
-    return await searchSiteContent(context, {
-        query,
-        scope: { mode: 'all' },
+        return await searchSiteContent(context, {
+            query,
+            scope: { mode: 'all' },
+        });
     });
 }
 
@@ -80,16 +83,18 @@ export async function searchAllSiteContent(query: string): Promise<OrderedComput
  * Server action to search content in a space.
  */
 export async function searchSiteSpaceContent(query: string): Promise<OrderedComputedResult[]> {
-    const context = await getServerActionBaseContext();
-    const siteURLData = await getSiteURLDataFromMiddleware();
+    return traceErrorOnly('Search.searchSiteSpaceContent', async () => {
+        const context = await getServerActionBaseContext();
+        const siteURLData = await getSiteURLDataFromMiddleware();
 
-    return await searchSiteContent(context, {
-        query,
-        // If we have a siteSectionId that means its a sections site use `current` mode
-        // which searches in the current space + all default spaces of sections
-        scope: siteURLData.siteSection
-            ? { mode: 'current', siteSpaceId: siteURLData.siteSpace }
-            : { mode: 'specific', siteSpaceIds: [siteURLData.siteSpace] },
+        return await searchSiteContent(context, {
+            query,
+            // If we have a siteSectionId that means its a sections site use `current` mode
+            // which searches in the current space + all default spaces of sections
+            scope: siteURLData.siteSection
+                ? { mode: 'current', siteSpaceId: siteURLData.siteSpace }
+                : { mode: 'specific', siteSpaceIds: [siteURLData.siteSpace] },
+        });
     });
 }
 
@@ -101,86 +106,88 @@ export async function streamAskQuestion({
 }: {
     question: string;
 }) {
-    const responseStream = createStreamableValue<AskAnswerResult | undefined>();
+    return traceErrorOnly('Search.streamAskQuestion', async () => {
+        const responseStream = createStreamableValue<AskAnswerResult | undefined>();
 
-    (async () => {
-        const context = await fetchServerActionSiteContext(await getServerActionBaseContext());
+        (async () => {
+            const context = await fetchServerActionSiteContext(await getServerActionBaseContext());
 
-        const apiClient = await context.dataFetcher.api();
+            const apiClient = await context.dataFetcher.api();
 
-        const stream = apiClient.orgs.streamAskInSite(
-            context.organizationId,
-            context.site.id,
-            {
-                question,
-                context: {
-                    siteSpaceId: context.siteSpace.id,
+            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],
+                    },
                 },
-                scope: {
-                    mode: 'default',
-                    // Include the current site space regardless.
-                    includedSiteSpaces: [context.siteSpace.id],
-                },
-            },
-            { format: 'document' }
-        );
-
-        const spacePromises = new Map<string, Promise<Revision>>();
-        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<string, RevisionPage[]>());
-            });
-            responseStream.update(
-                await transformAnswer(context, { answer: chunk.answer, spacePages: pages })
+                { format: 'document' }
             );
-        }
-    })()
-        .then(() => {
-            responseStream.done();
-        })
-        .catch((error) => {
-            responseStream.error(error);
-        });
 
-    return {
-        stream: responseStream.value,
-    };
+            const spacePromises = new Map<string, Promise<Revision>>();
+            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<string, RevisionPage[]>());
+                });
+                responseStream.update(
+                    await transformAnswer(context, { answer: chunk.answer, spacePages: pages })
+                );
+            }
+        })()
+            .then(() => {
+                responseStream.done();
+            })
+            .catch((error) => {
+                responseStream.error(error);
+            });
+
+        return {
+            stream: responseStream.value,
+        };
+    });
 }
 
 /**
@@ -188,33 +195,37 @@ export async function streamAskQuestion({
  * Optionally scoped to a specific space.
  */
 export async function streamRecommendedQuestions(spaceId?: string) {
-    const siteURLData = await getSiteURLDataFromMiddleware();
-    const context = await getServerActionBaseContext();
+    return traceErrorOnly('Search.streamRecommendedQuestions', async () => {
+        const siteURLData = await getSiteURLDataFromMiddleware();
+        const context = await getServerActionBaseContext();
 
-    const responseStream = createStreamableValue<SearchAIRecommendedQuestionStream | undefined>();
+        const responseStream = createStreamableValue<
+            SearchAIRecommendedQuestionStream | undefined
+        >();
 
-    (async () => {
-        const apiClient = await context.dataFetcher.api();
-        const apiStream = apiClient.orgs.streamRecommendedQuestionsInSite(
-            siteURLData.organization,
-            siteURLData.site,
-            {
-                spaceId,
+        (async () => {
+            const apiClient = await context.dataFetcher.api();
+            const apiStream = apiClient.orgs.streamRecommendedQuestionsInSite(
+                siteURLData.organization,
+                siteURLData.site,
+                {
+                    spaceId,
+                }
+            );
+
+            for await (const chunk of apiStream) {
+                responseStream.update(chunk);
             }
-        );
+        })()
+            .then(() => {
+                responseStream.done();
+            })
+            .catch((error) => {
+                responseStream.error(error);
+            });
 
-        for await (const chunk of apiStream) {
-            responseStream.update(chunk);
-        }
-    })()
-        .then(() => {
-            responseStream.done();
-        })
-        .catch((error) => {
-            responseStream.error(error);
-        });
-
-    return { stream: responseStream.value };
+        return { stream: responseStream.value };
+    });
 }
 
 /**

packages/gitbook/src/lib/tracing.ts

@@ -49,6 +49,41 @@ export async function trace<T>(
     }
 }
 
+/**
+ * Record a performance trace for the given function, but only log errors.
+ * This is useful to not output too much noise in the logs, while still capturing important errors.
+ */
+export async function traceErrorOnly<T>(
+    name: string | TraceName,
+    fn: (span: TraceSpan) => Promise<T>
+): Promise<T> {
+    const { operation, name: executionName } =
+        typeof name === 'string' ? { operation: name, name: undefined } : name;
+    const completeName = executionName ? `${operation}(${executionName})` : operation;
+
+    const attributes: Record<string, boolean | string | number> = {};
+    const span: TraceSpan = {
+        setAttribute(label, value) {
+            attributes[label] = value;
+        },
+    };
+
+    const start = now();
+    let traceError: null | Error = null;
+    try {
+        return await fn(span);
+    } catch (error) {
+        span.setAttribute('error', true);
+        traceError = error as Error;
+        const logger = getLogger().subLogger(operation);
+        logger.error(
+            `trace ${completeName} failed with ${traceError.message} in ${now() - start}ms`,
+            attributes
+        );
+        throw error;
+    }
+}
+
 /**
  * Return the current time in milliseconds.
  */