From a650b58276077e069db0c2153fbccf2935943f89 Mon Sep 17 00:00:00 2001 From: conico974 Date: Fri, 8 Aug 2025 16:10:06 +0200 Subject: [PATCH] feat: add tracing to server actions (#3543) Co-authored-by: Nicolas Dorseuil --- .../src/components/AI/server-actions/api.tsx | 184 +++++++------- .../src/components/AI/server-actions/chat.ts | 39 +-- .../components/AI/server-actions/responses.ts | 21 +- .../gitbook/src/components/Ads/renderAd.tsx | 63 ++--- .../AutoRefreshContent/server-actions.ts | 25 +- .../Integration/server-actions.tsx | 29 ++- .../src/components/Search/server-actions.tsx | 231 +++++++++--------- packages/gitbook/src/lib/tracing.ts | 35 +++ 8 files changed, 348 insertions(+), 279 deletions(-) diff --git a/packages/gitbook/src/components/AI/server-actions/api.tsx b/packages/gitbook/src/components/AI/server-actions/api.tsx index 12e82b62a..99b7424f1 100644 --- a/packages/gitbook/src/components/AI/server-actions/api.tsx +++ b/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, 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: ( - - ), + 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: ( + + ), + }; + }); }); } diff --git a/packages/gitbook/src/components/AI/server-actions/chat.ts b/packages/gitbook/src/components/AI/server-actions/chat.ts index 57900c600..86c67579d 100644 --- a/packages/gitbook/src/components/AI/server-actions/chat.ts +++ b/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; diff --git a/packages/gitbook/src/components/AI/server-actions/responses.ts b/packages/gitbook/src/components/AI/server-actions/responses.ts index 8326fe42f..41b702ac4 100644 --- a/packages/gitbook/src/components/AI/server-actions/responses.ts +++ b/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; diff --git a/packages/gitbook/src/components/Ads/renderAd.tsx b/packages/gitbook/src/components/Ads/renderAd.tsx index ae23047a8..5a964fcc1 100644 --- a/packages/gitbook/src/components/Ads/renderAd.tsx +++ b/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) ? ( - - ) : ( - - )} - {ad.pixel ? : null} - - ), - insightsAd, - }; + return { + children: ( + <> + {mode === 'classic' || !('callToAction' in ad) ? ( + + ) : ( + + )} + {ad.pixel ? : null} + + ), + insightsAd, + }; + }); } async function fetchAd({ diff --git a/packages/gitbook/src/components/AutoRefreshContent/server-actions.ts b/packages/gitbook/src/components/AutoRefreshContent/server-actions.ts index 866bfce35..5c7e2777a 100644 --- a/packages/gitbook/src/components/AutoRefreshContent/server-actions.ts +++ b/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; + }); } diff --git a/packages/gitbook/src/components/DocumentView/Integration/server-actions.tsx b/packages/gitbook/src/components/DocumentView/Integration/server-actions.tsx index 6487ec68c..27689c837 100644 --- a/packages/gitbook/src/components/DocumentView/Integration/server-actions.tsx +++ b/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: , + output: output.data, }; - } - - return { - children: , - output: output.data, - }; + }); } diff --git a/packages/gitbook/src/components/Search/server-actions.tsx b/packages/gitbook/src/components/Search/server-actions.tsx index e1b691f5f..708b799ae 100644 --- a/packages/gitbook/src/components/Search/server-actions.tsx +++ b/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 { - 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 { - 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(); + return traceErrorOnly('Search.streamAskQuestion', async () => { + const responseStream = createStreamableValue(); - (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>(); - 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, 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()); + }); + 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(); + 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 }; + }); } /** diff --git a/packages/gitbook/src/lib/tracing.ts b/packages/gitbook/src/lib/tracing.ts index 925d27e5f..099c35ac8 100644 --- a/packages/gitbook/src/lib/tracing.ts +++ b/packages/gitbook/src/lib/tracing.ts @@ -49,6 +49,41 @@ export async function trace( } } +/** + * 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( + name: string | TraceName, + fn: (span: TraceSpan) => Promise +): Promise { + const { operation, name: executionName } = + typeof name === 'string' ? { operation: name, name: undefined } : name; + const completeName = executionName ? `${operation}(${executionName})` : operation; + + const attributes: Record = {}; + 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. */