From 56866996ebbc16dc7bf145220a9c461acd8ef683 Mon Sep 17 00:00:00 2001 From: Ara Date: Thu, 5 Dec 2024 07:32:33 +0530 Subject: [PATCH] Adding Distributed Tracing and Smart Apply to cody (#6178) This PR introduces distributed tracing to connect traces originating in the UX with those starting in the extension host. This ensures unified and all-encompassing traces for operations spanning both the web view and the extension host. NOTE: this PR has been rebased to the main and is ready for review ## Changes Made 1. Distributed Tracing for Chat Interaction: Traces starting in the UX are now connected to corresponding traces in the extension host for chat interactions. 2. Distributed Tracing for Smart Apply: Similarly, traces for Smart Apply now span both the UX and the extension host. These changes are in alignment with the milestone goals and aim to deliver foundational support for distributed tracing in these areas. ## Next Steps - Immediate Goal: Ensure distributed tracing is functional and provides value for chat interaction and Smart Apply use cases. - I will rebase the PR and resolve merge conflicts AFTER the proper merge of https://github.com/sourcegraph/cody/pull/6100 it was reverted last time because [of some issues](https://sourcegraph.slack.com/archives/C05AGQYD528/p1732289326228309) - Follow-Up Work for next PR: - Refactor naming conventions for better consistency. - Eliminate redundant metric names. ## Test plan - Run sourcegraph instance locally - Run `sg start otel ` - Run the debugger for vscode cody locally on this branch - Perform some chat operations - Go to `http://localhost:16686` to see if Jaegar is running - Select Cody-Client as the service - See a trace with the title `chat-interaction ` this is a collection of spans that show a single trace for spans from both webview and extension host image ## Changelog --- lib/shared/src/prompt/prompt-string.ts | 4 + lib/shared/src/tracing/index.ts | 25 +- vscode/src/chat/chat-view/ChatController.ts | 82 ++-- vscode/src/chat/protocol.ts | 2 + vscode/src/edit/manager.ts | 384 +++++++++--------- vscode/src/edit/provider.ts | 8 + vscode/src/edit/smart-apply.ts | 1 + .../open-telemetry/CodyTraceExport.ts | 54 ++- .../utils/codeblock-action-tracker.ts | 4 +- vscode/webviews/Chat.tsx | 14 +- vscode/webviews/chat/Transcript.tsx | 10 +- vscode/webviews/utils/spanManager.ts | 4 +- vscode/webviews/utils/telemetry.ts | 7 + 13 files changed, 364 insertions(+), 235 deletions(-) diff --git a/lib/shared/src/prompt/prompt-string.ts b/lib/shared/src/prompt/prompt-string.ts index 87ea3509a9ca..a83bfffca36c 100644 --- a/lib/shared/src/prompt/prompt-string.ts +++ b/lib/shared/src/prompt/prompt-string.ts @@ -413,6 +413,10 @@ export class PromptString { : undefined, } } + // TODO: Lift the results of the matches so they're also PromptStrings. + public match(regexp: RegExp): RegExpMatchArray | null { + return internal_toString(this).match(regexp) + } } type TemplateArgs = readonly (PromptString | '' | number)[] diff --git a/lib/shared/src/tracing/index.ts b/lib/shared/src/tracing/index.ts index 06522c43ca9b..62244e02e394 100644 --- a/lib/shared/src/tracing/index.ts +++ b/lib/shared/src/tracing/index.ts @@ -1,4 +1,11 @@ -import opentelemetry, { SpanStatusCode, context, propagation, type Span } from '@opentelemetry/api' +import opentelemetry, { + type Context, + ROOT_CONTEXT, + SpanStatusCode, + context, + propagation, + type Span, +} from '@opentelemetry/api' import type { BrowserOrNodeResponse } from '../sourcegraph-api/graphql/client' const INSTRUMENTATION_SCOPE_NAME = 'cody' @@ -89,3 +96,19 @@ export function recordErrorToSpan(span: Span, error: Error): Error { span.end() return error } + +// Extracts a context from a traceparent header for use in a wrapped function +// that is called with context.with. This is useful for propagating the trace +// context between webview and extension host. +export function extractContextFromTraceparent(traceparent?: string | undefined | null): Context { + const carrier = { traceparent } as Record + const getter = { + get(carrier: Record, key: string) { + return carrier[key] + }, + keys(carrier: Record) { + return Object.keys(carrier) + }, + } + return propagation.extract(ROOT_CONTEXT, carrier, getter) +} diff --git a/vscode/src/chat/chat-view/ChatController.ts b/vscode/src/chat/chat-view/ChatController.ts index 83f11a22f12e..38cc6699d239 100644 --- a/vscode/src/chat/chat-view/ChatController.ts +++ b/vscode/src/chat/chat-view/ChatController.ts @@ -8,6 +8,7 @@ import { clientCapabilities, currentSiteVersion, distinctUntilChanged, + extractContextFromTraceparent, firstResultFromOperation, forceHydration, isAbortError, @@ -79,7 +80,7 @@ import { import * as uuid from 'uuid' import * as vscode from 'vscode' -import type { Span } from '@opentelemetry/api' +import { type Span, context } from '@opentelemetry/api' import { captureException } from '@sentry/core' import { getTokenCounterUtils } from '@sourcegraph/cody-shared/src/token/counter' import type { TelemetryEventParameters } from '@sourcegraph/telemetry' @@ -294,6 +295,7 @@ export class ChatController implements vscode.Disposable, vscode.WebviewViewProv intent: message.intent, intentScores: message.intentScores, manuallySelectedIntent: message.manuallySelectedIntent, + traceparent: message.traceparent, }) break } @@ -325,7 +327,8 @@ export class ChatController implements vscode.Disposable, vscode.WebviewViewProv message.code, currentAuthStatus(), message.instruction, - message.fileName + message.fileName, + message.traceparent ) break case 'trace-export': @@ -639,6 +642,7 @@ export class ChatController implements vscode.Disposable, vscode.WebviewViewProv intent: detectedIntent, intentScores: detectedIntentScores, manuallySelectedIntent, + traceparent, }: { requestID: string inputText: PromptString @@ -650,46 +654,50 @@ export class ChatController implements vscode.Disposable, vscode.WebviewViewProv intent?: ChatMessage['intent'] | undefined | null intentScores?: { intent: string; score: number }[] | undefined | null manuallySelectedIntent?: boolean | undefined | null + traceparent?: string | undefined | null }): Promise { - return tracer.startActiveSpan('chat.handleUserMessage', async (span): Promise => { - outputChannelLogger.logDebug( - 'ChatController', - 'handleUserMessageSubmission', - `traceId: ${span.spanContext().traceId}` - ) - span.setAttribute('sampled', true) - - if (inputText.toString().match(/^\/reset$/)) { - span.addEvent('clearAndRestartSession') - span.end() - return this.clearAndRestartSession() - } - - this.chatBuilder.addHumanMessage({ - text: inputText, - editorState, - intent: detectedIntent, - }) - this.postViewTranscript({ speaker: 'assistant' }) + return context.with(extractContextFromTraceparent(traceparent), () => { + return tracer.startActiveSpan('chat.handleUserMessage', async (span): Promise => { + span.setAttribute('sampled', true) + span.setAttribute('continued', true) + outputChannelLogger.logDebug( + 'ChatController', + 'handleUserMessageSubmission', + `traceId: ${span.spanContext().traceId}` + ) - void this.saveSession() - signal.throwIfAborted() + if (inputText.match(/^\/reset$/)) { + span.addEvent('clearAndRestartSession') + span.end() + return this.clearAndRestartSession() + } - return this.sendChat( - { - requestID, - inputText, - mentions, + this.chatBuilder.addHumanMessage({ + text: inputText, editorState, - signal, - source, - command, intent: detectedIntent, - intentScores: detectedIntentScores, - manuallySelectedIntent, - }, - span - ) + }) + this.postViewTranscript({ speaker: 'assistant' }) + + await this.saveSession() + signal.throwIfAborted() + + return this.sendChat( + { + requestID, + inputText, + mentions, + editorState, + signal, + source, + command, + intent: detectedIntent, + intentScores: detectedIntentScores, + manuallySelectedIntent, + }, + span + ) + }) }) } diff --git a/vscode/src/chat/protocol.ts b/vscode/src/chat/protocol.ts index 045904325f6d..ae9b99b01744 100644 --- a/vscode/src/chat/protocol.ts +++ b/vscode/src/chat/protocol.ts @@ -96,6 +96,7 @@ export type WebviewMessage = code: string instruction?: string | undefined | null fileName?: string | undefined | null + traceparent?: string | undefined | null } | { command: 'trace-export' @@ -206,6 +207,7 @@ export interface WebviewSubmitMessage extends WebviewContextMessage { intent?: ChatMessage['intent'] | undefined | null intentScores?: { intent: string; score: number }[] | undefined | null manuallySelectedIntent?: boolean | undefined | null + traceparent?: string | undefined | null } interface WebviewEditMessage extends WebviewContextMessage { diff --git a/vscode/src/edit/manager.ts b/vscode/src/edit/manager.ts index 9c760a3b269a..c8126dec899a 100644 --- a/vscode/src/edit/manager.ts +++ b/vscode/src/edit/manager.ts @@ -3,12 +3,15 @@ import * as vscode from 'vscode' import { type ChatClient, ClientConfigSingleton, + DEFAULT_EVENT_SOURCE, PromptString, currentSiteVersion, + extractContextFromTraceparent, firstResultFromOperation, modelsService, ps, telemetryRecorder, + wrapInActiveSpan, } from '@sourcegraph/cody-shared' import type { GhostHintDecorator } from '../commands/GhostHintDecorator' @@ -17,7 +20,7 @@ import type { VSCodeEditor } from '../editor/vscode-editor' import { FixupController } from '../non-stop/FixupController' import type { FixupTask } from '../non-stop/FixupTask' -import { DEFAULT_EVENT_SOURCE } from '@sourcegraph/cody-shared' +import { context } from '@opentelemetry/api' import { isUriIgnoredByContextFilterWithNotification } from '../cody-ignore/context-filter' import type { ExtensionClient } from '../extension-client' import { ACTIVE_TASK_STATES } from '../non-stop/codelenses/constants' @@ -236,196 +239,203 @@ export class EditManager implements vscode.Disposable { return task } - public async smartApplyEdit(args: SmartApplyArguments = {}): Promise { - const { configuration, source = 'chat' } = args - if (!configuration) { - return - } - - const document = configuration.document - if (await isUriIgnoredByContextFilterWithNotification(document.uri, 'edit')) { - return - } - - const model = - configuration.model || (await firstResultFromOperation(modelsService.getDefaultEditModel())) - if (!model) { - throw new Error('No default edit model found. Please set one.') - } - - telemetryRecorder.recordEvent('cody.command.smart-apply', 'executed', { - billingMetadata: { - product: 'cody', - category: 'core', - }, - }) - - const editor = await vscode.window.showTextDocument(document.uri) - - if (args.configuration?.isNewFile) { - // We are creating a new file, this means we are only _adding_ new code and _inserting_ it into the document. - // We do not need to re-prompt the LLM for this, let's just add the code directly. - const task = await this.controller.createTask( - document, - configuration.instruction, - [], - new vscode.Range(0, 0, 0, 0), - 'add', - 'insert', - model, - source, - configuration.document.uri, - undefined, - {}, - configuration.id - ) - - const legacyMetadata = { - intent: task.intent, - mode: task.mode, - source: task.source, - } - const { metadata, privateMetadata } = splitSafeMetadata(legacyMetadata) - telemetryRecorder.recordEvent('cody.command.edit', 'executed', { - metadata, - privateMetadata: { - ...privateMetadata, - model: task.model, - }, - billingMetadata: { - product: 'cody', - category: 'core', - }, - }) - - const provider = this.getProviderForTask(task) - await provider.applyEdit(configuration.replacement) - return task - } - - // Apply some decorations to the editor, this showcases that Cody is working on the full file range - // of the document. We will narrow it down to a selection soon. - const documentRange = new vscode.Range(0, 0, document.lineCount, 0) - editor.setDecorations(SMART_APPLY_FILE_DECORATION, [documentRange]) - - // We need to extract the proposed code, provided by the LLM, so we can use it in future - // queries to ask the LLM to generate a selection, and then ultimately apply the edit. - const replacementCode = PromptString.unsafe_fromLLMResponse(configuration.replacement) - - const versions = await currentSiteVersion() - if (!versions) { - throw new Error('unable to determine site version') - } - - const selection = await getSmartApplySelection( - configuration.id, - configuration.instruction, - replacementCode, - configuration.document, - model, - this.options.chat, - versions.codyAPIVersion - ) - - // We finished prompting the LLM for the selection, we can now remove the "progress" decoration - // that indicated we where working on the full file. - editor.setDecorations(SMART_APPLY_FILE_DECORATION, []) - - if (!selection) { - // We couldn't figure out the selection, let's inform the user and return early. - // TODO: Should we add a "Copy" button to this error? Then the user can copy the code directly. - void vscode.window.showErrorMessage( - 'Unable to apply this change to the file. Please try applying this code manually' - ) - telemetryRecorder.recordEvent('cody.smart-apply.selection', 'not-found') - return - } - - telemetryRecorder.recordEvent('cody.smart-apply', 'selected', { - metadata: { - [selection.type]: 1, - }, - }) - - // Move focus to the determined selection - editor.revealRange(selection.range, vscode.TextEditorRevealType.InCenter) - - if (selection.range.isEmpty) { - let insertionRange = selection.range - - if ( - selection.type === 'insert' && - document.lineAt(document.lineCount - 1).text.trim().length !== 0 - ) { - // Inserting to the bottom of the file, but the last line is not empty - // Inject an additional new line for us to use as the insertion range. - await editor.edit( - editBuilder => { - editBuilder.insert(selection.range.start, '\n') + public async smartApplyEdit(args: SmartApplyArguments = {}): Promise { + return context.with(extractContextFromTraceparent(args.configuration?.traceparent), async () => { + await wrapInActiveSpan('edit.smart-apply', async span => { + span.setAttribute('sampled', true) + span.setAttribute('continued', true) + const { configuration, source = 'chat' } = args + if (!configuration) { + return + } + + const document = configuration.document + if (await isUriIgnoredByContextFilterWithNotification(document.uri, 'edit')) { + return + } + + const model = + configuration.model || + (await firstResultFromOperation(modelsService.getDefaultEditModel())) + if (!model) { + throw new Error('No default edit model found. Please set one.') + } + + telemetryRecorder.recordEvent('cody.command.smart-apply', 'executed', { + billingMetadata: { + product: 'cody', + category: 'core', }, - { undoStopAfter: false, undoStopBefore: false } + }) + + const editor = await vscode.window.showTextDocument(document.uri) + + if (args.configuration?.isNewFile) { + // We are creating a new file, this means we are only _adding_ new code and _inserting_ it into the document. + // We do not need to re-prompt the LLM for this, let's just add the code directly. + const task = await this.controller.createTask( + document, + configuration.instruction, + [], + new vscode.Range(0, 0, 0, 0), + 'add', + 'insert', + model, + source, + configuration.document.uri, + undefined, + {}, + configuration.id + ) + + const legacyMetadata = { + intent: task.intent, + mode: task.mode, + source: task.source, + } + const { metadata, privateMetadata } = splitSafeMetadata(legacyMetadata) + telemetryRecorder.recordEvent('cody.command.edit', 'executed', { + metadata, + privateMetadata: { + ...privateMetadata, + model: task.model, + }, + billingMetadata: { + product: 'cody', + category: 'core', + }, + }) + + const provider = this.getProviderForTask(task) + await provider.applyEdit(configuration.replacement) + return task + } + + // Apply some decorations to the editor, this showcases that Cody is working on the full file range + // of the document. We will narrow it down to a selection soon. + const documentRange = new vscode.Range(0, 0, document.lineCount, 0) + editor.setDecorations(SMART_APPLY_FILE_DECORATION, [documentRange]) + + // We need to extract the proposed code, provided by the LLM, so we can use it in future + // queries to ask the LLM to generate a selection, and then ultimately apply the edit. + const replacementCode = PromptString.unsafe_fromLLMResponse(configuration.replacement) + + const versions = await currentSiteVersion() + if (!versions) { + throw new Error('unable to determine site version') + } + + const selection = await getSmartApplySelection( + configuration.id, + configuration.instruction, + replacementCode, + configuration.document, + model, + this.options.chat, + versions.codyAPIVersion ) - // Update the range to reflect the new end of document - insertionRange = document.lineAt(document.lineCount - 1).range - } - - // We determined a selection, but it was empty. This means that we will be _adding_ new code - // and _inserting_ it into the document. We do not need to re-prompt the LLM for this, let's just - // add the code directly. - const task = await this.controller.createTask( - document, - configuration.instruction, - [], - insertionRange, - 'add', - 'insert', - model, - source, - configuration.document.uri, - undefined, - {}, - configuration.id - ) - - const legacyMetadata = { - intent: task.intent, - mode: task.mode, - source: task.source, - } - const { metadata, privateMetadata } = splitSafeMetadata(legacyMetadata) - telemetryRecorder.recordEvent('cody.command.edit', 'executed', { - metadata, - privateMetadata: { - ...privateMetadata, - model: task.model, - }, - billingMetadata: { - product: 'cody', - category: 'core', - }, + // We finished prompting the LLM for the selection, we can now remove the "progress" decoration + // that indicated we where working on the full file. + editor.setDecorations(SMART_APPLY_FILE_DECORATION, []) + + if (!selection) { + // We couldn't figure out the selection, let's inform the user and return early. + // TODO: Should we add a "Copy" button to this error? Then the user can copy the code directly. + void vscode.window.showErrorMessage( + 'Unable to apply this change to the file. Please try applying this code manually' + ) + telemetryRecorder.recordEvent('cody.smart-apply.selection', 'not-found') + return + } + + telemetryRecorder.recordEvent('cody.smart-apply', 'selected', { + metadata: { + [selection.type]: 1, + }, + }) + + // Move focus to the determined selection + editor.revealRange(selection.range, vscode.TextEditorRevealType.InCenter) + + if (selection.range.isEmpty) { + let insertionRange = selection.range + + if ( + selection.type === 'insert' && + document.lineAt(document.lineCount - 1).text.trim().length !== 0 + ) { + // Inserting to the bottom of the file, but the last line is not empty + // Inject an additional new line for us to use as the insertion range. + await editor.edit( + editBuilder => { + editBuilder.insert(selection.range.start, '\n') + }, + { undoStopAfter: false, undoStopBefore: false } + ) + + // Update the range to reflect the new end of document + insertionRange = document.lineAt(document.lineCount - 1).range + } + + // We determined a selection, but it was empty. This means that we will be _adding_ new code + // and _inserting_ it into the document. We do not need to re-prompt the LLM for this, let's just + // add the code directly. + const task = await this.controller.createTask( + document, + configuration.instruction, + [], + insertionRange, + 'add', + 'insert', + model, + source, + configuration.document.uri, + undefined, + {}, + configuration.id + ) + + const legacyMetadata = { + intent: task.intent, + mode: task.mode, + source: task.source, + } + const { metadata, privateMetadata } = splitSafeMetadata(legacyMetadata) + telemetryRecorder.recordEvent('cody.command.edit', 'executed', { + metadata, + privateMetadata: { + ...privateMetadata, + model: task.model, + }, + billingMetadata: { + product: 'cody', + category: 'core', + }, + }) + + const provider = this.getProviderForTask(task) + await provider.applyEdit('\n' + configuration.replacement) + return task + } + + // We have a selection to replace, we re-prompt the LLM to generate the changes to ensure that + // we can reliably apply this edit. + // Just using the replacement code from the response is not enough, as it may contain parts that are not suitable to apply, + // e.g. // ... + return this.executeEdit({ + configuration: { + id: configuration.id, + document: configuration.document, + range: selection.range, + mode: 'edit', + instruction: ps`Ensuring that you do not duplicate code that it outside of the selection, apply the following change:\n${replacementCode}`, + model, + intent: 'edit', + }, + source, + }) }) - - const provider = this.getProviderForTask(task) - await provider.applyEdit('\n' + configuration.replacement) - return task - } - - // We have a selection to replace, we re-prompt the LLM to generate the changes to ensure that - // we can reliably apply this edit. - // Just using the replacement code from the response is not enough, as it may contain parts that are not suitable to apply, - // e.g. // ... - return this.executeEdit({ - configuration: { - id: configuration.id, - document: configuration.document, - range: selection.range, - mode: 'edit', - instruction: ps`Ensuring that you do not duplicate code that it outside of the selection, apply the following change:\n${replacementCode}`, - model, - intent: 'edit', - }, - source, }) } diff --git a/vscode/src/edit/provider.ts b/vscode/src/edit/provider.ts index d7e8de80945f..0cc5405f867e 100644 --- a/vscode/src/edit/provider.ts +++ b/vscode/src/edit/provider.ts @@ -12,6 +12,7 @@ import { modelsService, posixFilePaths, telemetryRecorder, + tracer, uriBasename, wrapInActiveSpan, } from '@sourcegraph/cody-shared' @@ -55,6 +56,8 @@ export class EditProvider { public async startEdit(): Promise { return wrapInActiveSpan('command.edit.start', async span => { + span.setAttribute('sampled', true) + const editTimeToFirstTokenSpan = tracer.startSpan('cody.edit.provider.timeToFirstToken') this.config.controller.startTask(this.config.task) const model = this.config.task.model const contextWindow = modelsService.getContextWindowByID(model) @@ -154,12 +157,17 @@ export class EditProvider { ) let textConsumed = 0 + let firstTokenReceived = false for await (const message of stream) { switch (message.type) { case 'change': { if (textConsumed === 0 && responsePrefix) { void multiplexer.publish(responsePrefix) } + if (!firstTokenReceived && message.text.length > 1) { + editTimeToFirstTokenSpan.end() + firstTokenReceived = true + } const text = message.text.slice(textConsumed) textConsumed += text.length void multiplexer.publish(text) diff --git a/vscode/src/edit/smart-apply.ts b/vscode/src/edit/smart-apply.ts index 2534f49e7085..722b0d707120 100644 --- a/vscode/src/edit/smart-apply.ts +++ b/vscode/src/edit/smart-apply.ts @@ -11,6 +11,7 @@ export interface SmartApplyArguments { document: vscode.TextDocument model?: EditModel isNewFile?: boolean + traceparent: string | undefined | null } source?: EventSource } diff --git a/vscode/src/services/open-telemetry/CodyTraceExport.ts b/vscode/src/services/open-telemetry/CodyTraceExport.ts index 015d743a41d0..b0e0d482daab 100644 --- a/vscode/src/services/open-telemetry/CodyTraceExport.ts +++ b/vscode/src/services/open-telemetry/CodyTraceExport.ts @@ -50,6 +50,20 @@ export class CodyTraceExporter extends OTLPTraceExporter { for (const span of spans) { const rootSpan = getRootSpan(spanMap, span) if (rootSpan === null) { + // The child of the root is sampled but root is not and the span is continued + // This for the cases where the root span is actually present in the webview + // but not in the extension host. + const effectiveRootSpan = getEffectiveRootSpan(spanMap, span) + if ( + effectiveRootSpan && + isSampled(effectiveRootSpan) && + isContinued(effectiveRootSpan) + ) { + spansToExport.push(span) + // Since we pushed the spans, we don't need to queue them + continue + } + const spanId = span.spanContext().spanId if (!this.queuedSpans.has(spanId)) { // No root span was found yet, so let's queue this span for a @@ -58,7 +72,7 @@ export class CodyTraceExporter extends OTLPTraceExporter { this.queuedSpans.set(spanId, { span, enqueuedAt: now }) } } else { - if (isRootSampled(rootSpan)) { + if (isSampled(rootSpan)) { spansToExport.push(span) } // else: The span is dropped @@ -69,6 +83,40 @@ export class CodyTraceExporter extends OTLPTraceExporter { } } +// This function checks if a span is continued in the extension host where the parent span is present +// in the webview. +function isContinued(span: ReadableSpan): boolean { + return span.attributes.continued === true +} + +// This function attempts to find the "effective root span" for a given span. +// The effective root span is defined as the first ancestor span that is not found in the span map. +// If a parent span is not found, it assumes the current span is the effective root. +function getEffectiveRootSpan( + spanMap: Map, + span: ReadableSpan +): ReadableSpan | null { + let currentSpan = span + + while (currentSpan.parentSpanId) { + const parentSpan = spanMap.get(currentSpan.parentSpanId) + if (!parentSpan) { + // If the parent span is not found in the map, the current span is considered the effective root. + return currentSpan + } + currentSpan = parentSpan + } + + // If there is no parent span ID, the span is considered a root span. + return null +} + +function isSampled(span: ReadableSpan): boolean { + return span.attributes.sampled === true +} + +// This function finds the root span of a given span so that we can check eventually check if it is sampled. +// This is useful to put all the spans that are part of the same trace together. function getRootSpan(spanMap: Map, span: ReadableSpan): ReadableSpan | null { if (span.parentSpanId) { const parentSpan = spanMap.get(span.parentSpanId) @@ -79,7 +127,3 @@ function getRootSpan(spanMap: Map, span: ReadableSpan): Re } return span } - -function isRootSampled(rootSpan: ReadableSpan): boolean { - return rootSpan.attributes.sampled === true -} diff --git a/vscode/src/services/utils/codeblock-action-tracker.ts b/vscode/src/services/utils/codeblock-action-tracker.ts index e97a2c0ad7b8..ef913457c787 100644 --- a/vscode/src/services/utils/codeblock-action-tracker.ts +++ b/vscode/src/services/utils/codeblock-action-tracker.ts @@ -158,7 +158,8 @@ export async function handleSmartApply( code: string, authStatus: AuthStatus, instruction?: string | null, - fileUri?: string | null + fileUri?: string | null, + traceparent?: string | undefined | null ): Promise { const activeEditor = getEditor()?.active const workspaceUri = vscode.workspace.workspaceFolders?.[0].uri @@ -196,6 +197,7 @@ export async function handleSmartApply( model: getSmartApplyModel(authStatus), replacement: code, isNewFile, + traceparent, }, source: 'chat', }) diff --git a/vscode/webviews/Chat.tsx b/vscode/webviews/Chat.tsx index 78ae80c9d7b5..61767bf6eadb 100644 --- a/vscode/webviews/Chat.tsx +++ b/vscode/webviews/Chat.tsx @@ -20,7 +20,8 @@ import WelcomeFooter from './chat/components/WelcomeFooter' import { WelcomeMessage } from './chat/components/WelcomeMessage' import { ScrollDown } from './components/ScrollDown' import type { View } from './tabs' -import { useTelemetryRecorder } from './utils/telemetry' +import { SpanManager } from './utils/spanManager' +import { getTraceparentFromSpanContext, useTelemetryRecorder } from './utils/telemetry' import { useUserAccountInfo } from './utils/useConfig' interface ChatboxProps { chatEnabled: boolean @@ -135,6 +136,15 @@ export const Chat: React.FunctionComponent instruction?: PromptString, fileName?: string ): void => { + const spanManager = new SpanManager('cody-webview') + const span = spanManager.startSpan('smartApplySubmit', { + attributes: { + sampled: true, + 'smartApply.id': id, + }, + }) + const traceparent = getTraceparentFromSpanContext(span.spanContext()) + vscodeAPI.postMessage({ command: 'smartApplySubmit', id, @@ -142,7 +152,9 @@ export const Chat: React.FunctionComponent // remove the additional /n added by the text area at the end of the text code: text.replace(/\n$/, ''), fileName, + traceparent, }) + span.end() }, onAccept: (id: string) => { vscodeAPI.postMessage({ diff --git a/vscode/webviews/chat/Transcript.tsx b/vscode/webviews/chat/Transcript.tsx index 84549055f82e..21af9f8d5eaa 100644 --- a/vscode/webviews/chat/Transcript.tsx +++ b/vscode/webviews/chat/Transcript.tsx @@ -35,7 +35,7 @@ import type { ApiPostMessage } from '../Chat' import { Button } from '../components/shadcn/ui/button' import { getVSCodeAPI } from '../utils/VSCodeApi' import { SpanManager } from '../utils/spanManager' -import { useTelemetryRecorder } from '../utils/telemetry' +import { getTraceparentFromSpanContext, useTelemetryRecorder } from '../utils/telemetry' import { useExperimentalOneBox } from '../utils/useExperimentalOneBox' import type { CodeBlockActionsProps } from './ChatMessageContent/ChatMessageContent' import { @@ -293,6 +293,9 @@ const TranscriptInteraction: FC = memo(props => { const spanContext = trace.setSpan(context.active(), span) setActiveChatContext(spanContext) + const currentSpanContext = span.spanContext() + + const traceparent = getTraceparentFromSpanContext(currentSpanContext) // Serialize the editor value after starting the span const editorValue = humanEditorRef.current?.getSerializedValue() @@ -306,6 +309,7 @@ const TranscriptInteraction: FC = memo(props => { intent: intentFromSubmit || intentResults.current?.intent, intentScores: intentFromSubmit ? undefined : intentResults.current?.allScores, manuallySelectedIntent: !!intentFromSubmit, + traceparent, } if (action === 'edit') { @@ -433,6 +437,7 @@ const TranscriptInteraction: FC = memo(props => { }) renderSpan.current.end() } + renderSpan.current = undefined hasRecordedFirstToken.current = false @@ -714,11 +719,13 @@ function submitHumanMessage({ intent, intentScores, manuallySelectedIntent, + traceparent, }: { editorValue: SerializedPromptEditorValue intent?: ChatMessage['intent'] intentScores?: { intent: string; score: number }[] manuallySelectedIntent?: boolean + traceparent: string }): void { getVSCodeAPI().postMessage({ command: 'submit', @@ -728,6 +735,7 @@ function submitHumanMessage({ intent, intentScores, manuallySelectedIntent, + traceparent, }) focusLastHumanMessageEditor() } diff --git a/vscode/webviews/utils/spanManager.ts b/vscode/webviews/utils/spanManager.ts index 9841644e6f36..fc685c4b76de 100644 --- a/vscode/webviews/utils/spanManager.ts +++ b/vscode/webviews/utils/spanManager.ts @@ -79,9 +79,9 @@ export class SpanManager { }) } - startSpan(name: string, options?: SpanManagerOptions): Span | undefined { + startSpan(name: string, options?: SpanManagerOptions): Span { if (this.spans.has(name)) { - return this.spans.get(name) + return this.spans.get(name)! } // Use provided context or fall back to active context diff --git a/vscode/webviews/utils/telemetry.ts b/vscode/webviews/utils/telemetry.ts index 232d5f8723d4..c40d4bc0365c 100644 --- a/vscode/webviews/utils/telemetry.ts +++ b/vscode/webviews/utils/telemetry.ts @@ -1,5 +1,6 @@ import type { TelemetryRecorder } from '@sourcegraph/cody-shared' +import type { SpanContext } from '@opentelemetry/api' import { createContext, useContext } from 'react' import type { WebviewRecordEventParameters } from '../../src/chat/protocol' import type { ApiPostMessage } from '../Chat' @@ -37,3 +38,9 @@ export function useTelemetryRecorder(): TelemetryRecorder { } return telemetryRecorder } + +export function getTraceparentFromSpanContext(spanContext: SpanContext): string { + return `00-${spanContext.traceId}-${spanContext.spanId}-${spanContext.traceFlags + .toString(16) + .padStart(2, '0')}` +}