Autocomplete: Improve sampling code and prepare for Honeycomb export

lib/shared/src/experimentation/FeatureFlagProvider.ts

@@ -68,6 +68,10 @@ export class FeatureFlagProvider {
         return this.featureFlags[endpoint]?.[flagName]
     }
 
+    public getExposedExperiments(endpoint: string = this.apiClient.endpoint): Record<string, boolean> {
+        return this.featureFlags[endpoint] || {}
+    }
+
     public async evaluateFeatureFlag(flagName: FeatureFlag): Promise<boolean> {
         const endpoint = this.apiClient.endpoint
         if (process.env.BENCHMARK_DISABLE_FEATURE_FLAGS) {

lib/shared/src/sourcegraph-api/graphql/client.ts

@@ -769,7 +769,7 @@ export class SourcegraphGraphQLAPIClient {
         const queryName = query.match(QUERY_TO_NAME_REGEXP)?.[1]
 
         const url = buildGraphQLUrl({ request: query, baseUrl: this.config.serverEndpoint })
-        return wrapInActiveSpan(`graphql.fetch${queryName ? `.${queryName}` : ''}`, () =>
+        return (
             fetch(url, {
                 method: 'POST',
                 body: JSON.stringify({ query, variables }),
@@ -800,7 +800,7 @@ export class SourcegraphGraphQLAPIClient {
 
         const queryName = query.match(QUERY_TO_NAME_REGEXP)?.[1]
 
-        return wrapInActiveSpan(`graphql.dotcom.fetch${queryName ? `.${queryName}` : ''}`, () =>
+        return wrapInActiveSpan(`graphql.fetch${queryName ? `.${queryName}` : ''}`, () =>
             fetch(url, {
                 method: 'POST',
                 body: JSON.stringify({ query, variables }),

lib/shared/src/tracing/index.ts

@@ -1,4 +1,10 @@
-import opentelemetry, { SpanStatusCode, context, propagation, type Exception } from '@opentelemetry/api'
+import opentelemetry, {
+    SpanStatusCode,
+    context,
+    propagation,
+    type Exception,
+    Span,
+} from '@opentelemetry/api'
 
 const INSTRUMENTATION_SCOPE_NAME = 'cody'
 const INSTRUMENTATION_SCOPE_VERSION = '0.1'
@@ -17,31 +23,32 @@ export function getActiveTraceAndSpanId(): { traceId: string; spanId: string } |
     return undefined
 }
 
-export function wrapInActiveSpan<R>(name: string, fn: () => R): R {
+export function wrapInActiveSpan<R>(name: string, fn: (span: Span) => R): R {
     return tracer.startActiveSpan(name, (span): R => {
         const handleSuccess = (response: R): R => {
             span.setStatus({ code: SpanStatusCode.OK })
+            span.end()
             return response
         }
 
-        const catchError = (error: unknown): never => {
+        const handleError = (error: unknown): never => {
             span.recordException(error as Exception)
             span.setStatus({ code: SpanStatusCode.ERROR })
+            span.end()
             throw error
         }
 
         try {
-            const response = fn()
+            const response = fn(span)
 
-            if (response instanceof Promise) {
-                return response.then(handleSuccess, catchError) as R
+            if (typeof response === 'object' && response !== null && 'then' in response) {
+                // @ts-ignore Response seems to be a Thenable
+                return response.then(handleSuccess, handleError) as R
             }
 
             return handleSuccess(response)
         } catch (error) {
-            return catchError(error)
-        } finally {
-            span.end()
+            return handleError(error)
         }
     })
 }

vscode/package.json

@@ -38,7 +38,12 @@
         "test:unit:tree-sitter-queries": "vitest ./src/tree-sitter/query-tests/*.test.ts",
         "github-changelog": "ts-node-transpile-only ./scripts/github-changelog.ts"
     },
-    "categories": ["Programming Languages", "Machine Learning", "Snippets", "Education"],
+    "categories": [
+        "Programming Languages",
+        "Machine Learning",
+        "Snippets",
+        "Education"
+    ],
     "keywords": [
         "ai",
         "openai",
@@ -89,7 +94,11 @@
     },
     "main": "./dist/extension.node.js",
     "browser": "./dist/extension.web.js",
-    "activationEvents": ["onLanguage", "onStartupFinished", "onWebviewPanel:cody.chatPanel"],
+    "activationEvents": [
+        "onLanguage",
+        "onStartupFinished",
+        "onWebviewPanel:cody.chatPanel"
+    ],
     "contributes": {
         "walkthroughs": [
             {
@@ -744,7 +753,12 @@
                 "cody.useContext": {
                     "order": 99,
                     "type": "string",
-                    "enum": ["embeddings", "keyword", "blended", "none"],
+                    "enum": [
+                        "embeddings",
+                        "keyword",
+                        "blended",
+                        "none"
+                    ],
                     "default": "blended",
                     "markdownDescription": "Controls which context providers Cody uses for chat, commands and inline edits. Use 'blended' for best results. For debugging other context sources, 'embeddings' will use an embeddings-based index if available. 'keyword' will use a search-based index. 'none' will not use embeddings or search-based indexes."
                 },
@@ -808,7 +822,9 @@
                     "order": 6,
                     "type": "string",
                     "markdownDescription": "A custom instruction to be included at the start of all chat messages. (E.g., \"Answer all my questions in Spanish.\")",
-                    "examples": ["Answer all my questions in Spanish."]
+                    "examples": [
+                        "Answer all my questions in Spanish."
+                    ]
                 },
                 "cody.codeActions.enabled": {
                     "order": 11,
@@ -866,7 +882,10 @@
                 "cody.telemetry.level": {
                     "order": 99,
                     "type": "string",
-                    "enum": ["all", "off"],
+                    "enum": [
+                        "all",
+                        "off"
+                    ],
                     "enumDescriptions": [
                         "Sends usage data and errors.",
                         "Disables all extension telemetry."
@@ -877,7 +896,13 @@
                 "cody.autocomplete.advanced.provider": {
                     "type": "string",
                     "default": null,
-                    "enum": [null, "anthropic", "fireworks", "unstable-openai", "unstable-ollama"],
+                    "enum": [
+                        null,
+                        "anthropic",
+                        "fireworks",
+                        "unstable-openai",
+                        "unstable-ollama"
+                    ],
                     "markdownDescription": "The provider used for code autocomplete. Most providers other than `anthropic` require the `cody.autocomplete.advanced.serverEndpoint` and `cody.autocomplete.advanced.accessToken` settings to also be set. Check the Cody output channel for error messages if autocomplete is not working as expected."
                 },
                 "cody.autocomplete.advanced.serverEndpoint": {
@@ -916,7 +941,10 @@
                 },
                 "cody.experimental.foldingRanges": {
                     "type": "string",
-                    "enum": ["lsp", "indentation-based"],
+                    "enum": [
+                        "lsp",
+                        "indentation-based"
+                    ],
                     "enumDescriptions": [
                         "Use folding ranges that are enabled by default in VS Code, and are usually powered by LSP",
                         "Use custom implementation of folding ranges that is indentation based. This is the implementation that is used by other Cody clients like the JetBrains plugin"
@@ -947,7 +975,11 @@
                 "cody.autocomplete.experimental.graphContext": {
                     "type": "string",
                     "default": null,
-                    "enum": [null, "bfg", "bfg-mixed"],
+                    "enum": [
+                        null,
+                        "bfg",
+                        "bfg-mixed"
+                    ],
                     "markdownDescription": "Use the code graph to retrieve context for autocomplete requests."
                 },
                 "cody.autocomplete.experimental.ollamaOptions": {
@@ -1040,6 +1072,7 @@
     "dependencies": {
         "@anthropic-ai/sdk": "^0.4.2",
         "@opentelemetry/api": "^1.7.0",
+        "@opentelemetry/auto-instrumentations-node": "^0.41.1",
         "@opentelemetry/core": "^1.18.1",
         "@opentelemetry/exporter-trace-otlp-http": "^0.45.1",
         "@opentelemetry/instrumentation-http": "^0.45.1",

vscode/src/completions/inline-completion-item-provider.ts

@@ -191,7 +191,7 @@ export class InlineCompletionItemProvider
             return null
         }
 
-        return wrapInActiveSpan('autocomplete.provideInlineCompletionItems', async () => {
+        return wrapInActiveSpan('autocomplete.provideInlineCompletionItems', async span => {
             // Update the last request
             const lastCompletionRequest = this.lastCompletionRequest
             const completionRequest: CompletionRequest = {
@@ -209,10 +209,7 @@ export class InlineCompletionItemProvider
                     // the error banner for autocomplete config turned off
                     throw new Error('AutocompleteConfigTurnedOff')
                 }
-            } catch (error) {
-                this.onError(error as Error)
-                throw error
-            }
+            } catch (error) {}
             const start = performance.now()
 
             if (!this.lastCompletionRequestTimestamp) {
@@ -406,7 +403,8 @@ export class InlineCompletionItemProvider
                     docContext,
                     position,
                     visibleItems,
-                    context
+                    context,
+                    span
                 )
 
                 // Store the log ID for each completion item so that we can later map to the selected
@@ -514,14 +512,14 @@ export class InlineCompletionItemProvider
      * same name, it's prefixed with `unstable_` to avoid a clash when the new API goes GA.
      */
     public unstable_handleDidShowCompletionItem(
-        completionOrItemId: Pick<AutocompleteItem, 'logId' | 'analyticsItem'> | CompletionItemID
+        completionOrItemId: Pick<AutocompleteItem, 'logId' | 'analyticsItem' | 'span'> | CompletionItemID
     ): void {
         const completion = suggestedAutocompleteItemsCache.get(completionOrItemId)
         if (!completion) {
             return
         }
 
-        CompletionLogger.suggested(completion.logId)
+        CompletionLogger.suggested(completion.logId, completion.span)
     }
 
     /**

vscode/src/completions/logger.ts

@@ -2,7 +2,12 @@ import { LRUCache } from 'lru-cache'
 import * as uuid from 'uuid'
 import * as vscode from 'vscode'
 
-import { isNetworkError, type BillingCategory, type BillingProduct } from '@sourcegraph/cody-shared'
+import {
+    isNetworkError,
+    type BillingCategory,
+    type BillingProduct,
+    featureFlagProvider,
+} from '@sourcegraph/cody-shared'
 import type { KnownString, TelemetryEventParameters } from '@sourcegraph/telemetry'
 
 import { getConfiguration } from '../configuration'
@@ -19,6 +24,7 @@ import * as statistics from './statistics'
 import type { InlineCompletionItemWithAnalytics } from './text-processing/process-inline-completions'
 import { lines } from './text-processing/utils'
 import type { InlineCompletionItem } from './types'
+import { Span, trace } from '@opentelemetry/api'
 
 // A completion ID is a unique identifier for a specific completion text displayed at a specific
 // point in the document. A single completion can be suggested multiple times.
@@ -505,7 +511,7 @@ export function loaded(
 //
 // For statistics logging we start a timeout matching the READ_TIMEOUT_MS so we can increment the
 // suggested completion count as soon as we count it as such.
-export function suggested(id: CompletionLogID): void {
+export function suggested(id: CompletionLogID, span?: Span): void {
     const event = activeSuggestionRequests.get(id)
     if (!event) {
         return
@@ -519,6 +525,12 @@ export function suggested(id: CompletionLogID): void {
     if (!event.suggestedAt) {
         event.suggestedAt = performance.now()
 
+        // Add exposed experiments at the very end to make sure we include experiments that the user
+        // is being exposed to while the completion was generated
+        span?.setAttributes(featureFlagProvider.getExposedExperiments())
+        span?.setAttributes(getSharedParams(event) as any)
+        span?.addEvent('suggested')
+
         setTimeout(() => {
             const event = activeSuggestionRequests.get(id)
             if (!event) {

vscode/src/completions/providers/fireworks.ts

@@ -129,8 +129,8 @@ class FireworksProvider extends Provider {
             this.options.fastPath &&
             config.accessToken &&
             // Require the upstream to be dotcom
-            this.authStatus.isDotCom &&
-            // The fast path client only supports Node.js style response streams
+            // this.authStatus.isDotCom &&
+            // The fast path client only suppor ts Node.js style response streams
             isNode
                 ? dotcomTokenToGatewayToken(config.accessToken)
                 : undefined

vscode/src/completions/suggested-autocomplete-items-cache.ts

@@ -6,6 +6,7 @@ import type { DocumentContext } from './get-current-doc-context'
 import type { CompletionItemID, CompletionLogID } from './logger'
 import type { RequestParams } from './request-manager'
 import type { InlineCompletionItemWithAnalytics } from './text-processing/process-inline-completions'
+import type { Span } from '@opentelemetry/api'
 
 interface AutocompleteItemParams {
     insertText: string | vscode.SnippetString
@@ -15,6 +16,7 @@ interface AutocompleteItemParams {
     requestParams: RequestParams
     completionItem: InlineCompletionItemWithAnalytics
     command?: vscode.Command
+    span?: Span
 }
 
 export class AutocompleteItem extends vscode.InlineCompletionItem {
@@ -52,8 +54,14 @@ export class AutocompleteItem extends vscode.InlineCompletionItem {
      */
     public analyticsItem: InlineCompletionItemWithAnalytics
 
+    /**
+     * Eventual Open Telemetry span associated with the completion request
+     */
+    public span: Span | undefined
+
     constructor(params: AutocompleteItemParams) {
-        const { insertText, logId, range, trackedRange, requestParams, completionItem, command } = params
+        const { insertText, logId, range, trackedRange, requestParams, completionItem, command, span } =
+            params
 
         super(insertText, range, command)
 
@@ -62,6 +70,7 @@ export class AutocompleteItem extends vscode.InlineCompletionItem {
         this.trackedRange = trackedRange
         this.requestParams = requestParams
         this.analyticsItem = completionItem
+        this.span = span
     }
 }
 
@@ -103,7 +112,8 @@ export function analyticsItemToAutocompleteItem(
     docContext: DocumentContext,
     position: vscode.Position,
     items: InlineCompletionItemWithAnalytics[],
-    context: vscode.InlineCompletionContext
+    context: vscode.InlineCompletionContext,
+    span: Span
 ): AutocompleteItem[] {
     return items.map(item => {
         const { insertText, range } = item
@@ -144,6 +154,7 @@ export function analyticsItemToAutocompleteItem(
             requestParams,
             completionItem: item,
             command,
+            span,
         })
 
         command.arguments[0].codyCompletion = autocompleteItem

vscode/src/services/open-telemetry/OpenTelemetryService.node.ts

@@ -5,6 +5,7 @@ import { Resource } from '@opentelemetry/resources'
 import { NodeSDK } from '@opentelemetry/sdk-node'
 import { BatchSpanProcessor } from '@opentelemetry/sdk-trace-base'
 import { SemanticResourceAttributes } from '@opentelemetry/semantic-conventions'
+import { getNodeAutoInstrumentations } from '@opentelemetry/auto-instrumentations-node'
 
 import {
     FeatureFlag,
@@ -62,13 +63,11 @@ export class OpenTelemetryService {
                 [SemanticResourceAttributes.SERVICE_VERSION]: version,
             }),
             instrumentations: [new HttpInstrumentation()],
-            traceExporter: new OTLPTraceExporter({
-                url: traceUrl,
-            }),
+            traceExporter: new OTLPTraceExporter({ url: traceUrl }),
 
-            ...(process.env.NODE_ENV === 'development' && {
-                spanProcessor: new BatchSpanProcessor(new ConsoleBatchSpanExporter()),
-            }),
+            // ...(process.env.NODE_ENV === 'development' && {
+            //     spanProcessor: new BatchSpanProcessor(new ConsoleBatchSpanExporter()),
+            // }),
         })
         this.sdk.start()
     }