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 #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

<img width="1483" alt="image"
src="https://github.com/user-attachments/assets/1db63541-46c4-4055-80ef-0a9dab1e165e">


## Changelog

<!-- OPTIONAL; info at
https://www.notion.so/sourcegraph/Writing-a-changelog-entry-dd997f411d524caabf0d8d38a24a878c
-->

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)[]

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<string, any>
+    const getter = {
+        get(carrier: Record<string, any>, key: string) {
+            return carrier[key]
+        },
+        keys(carrier: Record<string, any>) {
+            return Object.keys(carrier)
+        },
+    }
+    return propagation.extract(ROOT_CONTEXT, carrier, getter)
+}

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<void> {
-        return tracer.startActiveSpan('chat.handleUserMessage', async (span): Promise<void> => {
-            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<void> => {
+                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
+                )
+            })
         })
     }
 

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 {