refactor(telemetry): Remove telemetry.foo.record (#5646)

## Problem
- Theres quite a few missues of telemetry.foo.record inside of the
codebase and it can be confusing for external contributors

## Solution
- Clean up the telemetry interface by only allowing record/increment on
spans

Author: jpinkney-aws

package-lock.json

@@ -39,7 +39,7 @@
         "generateNonCodeFiles": "npm run generateNonCodeFiles -w packages/ --if-present"
     },
     "devDependencies": {
-        "@aws-toolkits/telemetry": "^1.0.258",
+        "@aws-toolkits/telemetry": "^1.0.267",
         "@playwright/browser-chromium": "^1.43.1",
         "@types/vscode": "^1.68.0",
         "@types/vscode-webview": "^1.57.1",

package.json

@@ -11,7 +11,8 @@ import {
     maxRepoSizeBytes,
 } from 'aws-core-vscode/amazonqFeatureDev'
 import { assertTelemetry, createTestWorkspace } from 'aws-core-vscode/test'
-import { fs, AmazonqCreateUpload, Metric } from 'aws-core-vscode/shared'
+import { fs, AmazonqCreateUpload } from 'aws-core-vscode/shared'
+import { Span } from 'aws-core-vscode/telemetry'
 import sinon from 'sinon'
 
 describe('file utils', () => {
@@ -28,7 +29,7 @@ describe('file utils', () => {
             const telemetry = new TelemetryHelper()
             const result = await prepareRepoData([workspace.uri.fsPath], [workspace], telemetry, {
                 record: () => {},
-            } as unknown as Metric<AmazonqCreateUpload>)
+            } as unknown as Span<AmazonqCreateUpload>)
             assert.strictEqual(Buffer.isBuffer(result.zipFileBuffer), true)
             // checksum is not the same across different test executions because some unique random folder names are generated
             assert.strictEqual(result.zipFileChecksum.length, 44)
@@ -46,7 +47,7 @@ describe('file utils', () => {
             const telemetry = new TelemetryHelper()
             const result = await prepareRepoData([workspace.uri.fsPath], [workspace], telemetry, {
                 record: () => {},
-            } as unknown as Metric<AmazonqCreateUpload>)
+            } as unknown as Span<AmazonqCreateUpload>)
 
             assert.strictEqual(Buffer.isBuffer(result.zipFileBuffer), true)
             // checksum is not the same across different test executions because some unique random folder names are generated
@@ -65,7 +66,7 @@ describe('file utils', () => {
                 () =>
                     prepareRepoData([workspace.uri.fsPath], [workspace], telemetry, {
                         record: () => {},
-                    } as unknown as Metric<AmazonqCreateUpload>),
+                    } as unknown as Span<AmazonqCreateUpload>),
                 ContentLengthError
             )
         })

packages/amazonq/test/unit/amazonqFeatureDev/util/files.test.ts

@@ -14,7 +14,7 @@ import { maxFileSizeBytes } from '../limits'
 import { createHash } from 'crypto'
 import { CurrentWsFolders } from '../types'
 import { ToolkitError } from '../../shared/errors'
-import { AmazonqCreateUpload, Metric, telemetry as amznTelemetry } from '../../shared/telemetry/telemetry'
+import { AmazonqCreateUpload, Span, telemetry as amznTelemetry } from '../../shared/telemetry/telemetry'
 import { TelemetryHelper } from './telemetryHelper'
 import { maxRepoSizeBytes } from '../constants'
 import { isCodeFile } from '../../shared/filetypes'
@@ -28,7 +28,7 @@ export async function prepareRepoData(
     repoRootPaths: string[],
     workspaceFolders: CurrentWsFolders,
     telemetry: TelemetryHelper,
-    span: Metric<AmazonqCreateUpload>
+    span: Span<AmazonqCreateUpload>
 ) {
     try {
         const files = await collectFiles(repoRootPaths, workspaceFolders, true, maxRepoSizeBytes)

packages/core/src/amazonqFeatureDev/util/files.ts

@@ -5,7 +5,7 @@
 
 import { globals } from '../../shared'
 import { getLogger } from '../../shared/logger/logger'
-import { AmazonqApproachInvoke, AmazonqCodeGenerationInvoke, Metric } from '../../shared/telemetry/telemetry'
+import { AmazonqApproachInvoke, AmazonqCodeGenerationInvoke, Span } from '../../shared/telemetry/telemetry'
 import { LLMResponseType } from '../types'
 
 export class TelemetryHelper {
@@ -32,7 +32,7 @@ export class TelemetryHelper {
     }
 
     public recordUserApproachTelemetry(
-        span: Metric<AmazonqApproachInvoke>,
+        span: Span<AmazonqApproachInvoke>,
         amazonqConversationId: string,
         responseType: LLMResponseType
     ) {
@@ -48,7 +48,7 @@ export class TelemetryHelper {
         span.record(event)
     }
 
-    public recordUserCodeGenerationTelemetry(span: Metric<AmazonqCodeGenerationInvoke>, amazonqConversationId: string) {
+    public recordUserCodeGenerationTelemetry(span: Span<AmazonqCodeGenerationInvoke>, amazonqConversationId: string) {
         const event = {
             amazonqConversationId,
             amazonqGenerateCodeIteration: this.generateCodeIteration,

packages/core/src/amazonqFeatureDev/util/telemetryHelper.ts

@@ -12,7 +12,7 @@ import { isNonNullable } from '../shared/utilities/tsUtils'
 import { ToolIdStateKey } from '../shared/globalState'
 import { Connection, getTelemetryMetadataForConn, SsoConnection, StatefulConnection } from './connection'
 import { indent } from '../shared/utilities/textUtilities'
-import { AuthStatus, telemetry } from '../shared/telemetry/telemetry'
+import { AuthModifyConnection, AuthStatus, Span, telemetry } from '../shared/telemetry/telemetry'
 import { asStringifiedStack } from '../shared/telemetry/spans'
 import { withTelemetryContext } from '../shared/telemetry/util'
 import { isNetworkError } from '../shared/errors'
@@ -307,7 +307,7 @@ export class SecondaryAuth<T extends Connection = Connection> {
                     connectionState: 'undefined',
                 })
                 await this.auth.tryAutoConnect()
-                this.#savedConnection = await this._loadSavedConnection()
+                this.#savedConnection = await this._loadSavedConnection(span)
                 this.#onDidChangeActiveConnection.fire(this.activeConnection)
 
                 const conn = this.#savedConnection
@@ -328,7 +328,7 @@ export class SecondaryAuth<T extends Connection = Connection> {
     /**
      * Provides telemetry if called by restoreConnection() (or another auth_modifyConnection context)
      */
-    private async _loadSavedConnection() {
+    private async _loadSavedConnection(span: Span<AuthModifyConnection>) {
         const id = this.getStateConnectionId()
         if (id === undefined) {
             return
@@ -349,7 +349,7 @@ export class SecondaryAuth<T extends Connection = Connection> {
             let connectionState = this.auth.getConnectionState(conn)
 
             // This function is expected to be called in the context of restoreConnection()
-            telemetry.auth_modifyConnection.record({
+            span.record({
                 connectionState,
                 authStatus: getAuthStatus(connectionState),
             })
@@ -358,7 +358,7 @@ export class SecondaryAuth<T extends Connection = Connection> {
                 await this.auth.refreshConnectionState(conn)
 
                 connectionState = this.auth.getConnectionState(conn)
-                telemetry.auth_modifyConnection.record({
+                span.record({
                     connectionState,
                     authStatus: getAuthStatus(connectionState),
                 })
@@ -367,7 +367,7 @@ export class SecondaryAuth<T extends Connection = Connection> {
                 // If updating the state fails, then we should delegate downstream to handle getting the proper state.
                 getLogger().error('loadSavedConnection: Failed to refresh connection state: %s', err)
                 if (isNetworkError(err) && connectionState === 'valid') {
-                    telemetry.auth_modifyConnection.record({
+                    span.record({
                         authStatus: 'connectedWithNetworkError',
                     })
                 }

packages/core/src/auth/secondaryAuth.ts

@@ -4,5 +4,5 @@
  */
 
 export { activate } from './activation'
-export { telemetry, AuthStatus } from './telemetry'
+export { telemetry, AuthStatus, Span } from './telemetry'
 export { ExtStartUpSources } from './util'

packages/core/src/shared/telemetry/index.ts

@@ -13,6 +13,7 @@ import {
     MetricDefinition,
     MetricName,
     MetricShapes,
+    Span,
     TelemetryBase,
 } from './telemetry.gen'
 import {
@@ -132,7 +133,7 @@ export type SpanOptions = {
  *
  * See also: docs/telemetry.md
  */
-export class TelemetrySpan<T extends MetricBase = MetricBase> {
+export class TelemetrySpan<T extends MetricBase = MetricBase> implements Span<T> {
     #startTime?: Date
     #options: SpanOptions
 
@@ -317,6 +318,9 @@ export class TelemetryTracer extends TelemetryBase {
      * All changes made to {@link attributes} (via {@link record}) during the execution are
      * reverted after the execution completes.
      *
+     * Runs can be nested within each other. This allows for creating hierarchical spans,
+     * where child spans inherit the context of their parent spans.
+     *
      * This method automatically handles traceId generation and propagation:
      * - If no traceId exists in the current context, a new one is generated.
      * - The traceId is attached to all telemetry events created within this span.
@@ -325,19 +329,13 @@ export class TelemetryTracer extends TelemetryBase {
      *
      * See docs/telemetry.md
      */
-    public run<T, U extends MetricName>(name: U, fn: (span: Metric<MetricShapes[U]>) => T, options?: SpanOptions): T {
+    public run<T, U extends MetricName>(name: U, fn: (span: Span<MetricShapes[U]>) => T, options?: SpanOptions): T {
         return this.withTraceId(() => {
             const span = this.createSpan(name, options).start()
             const frame = this.switchContext(span)
 
             try {
-                //
-                // TODO: Since updating to `@types/node@16`, typescript flags this code with error:
-                //
-                //      Error: npm ERR! src/shared/telemetry/spans.ts(255,57): error TS2345: Argument of type
-                //      'TelemetrySpan<MetricBase>' is not assignable to parameter of type 'Metric<MetricShapes[U]>'.
-                //
-                const result = this.#context.run(frame, fn, span as any)
+                const result = this.#context.run(frame, fn, span)
 
                 if (result instanceof Promise) {
                     return result
@@ -444,9 +442,7 @@ export class TelemetryTracer extends TelemetryBase {
         return {
             name,
             emit: (data) => getSpan().emit(data),
-            record: (data) => getSpan().record(data),
             run: (fn, options?: SpanOptions) => this.run(name as MetricName, fn, options),
-            increment: (data) => getSpan().increment(data),
         }
     }
 

packages/core/src/shared/telemetry/spans.ts

@@ -23,7 +23,10 @@ export const telemetry = new TelemetryTracer()
  */
 declare module './telemetry.gen' {
     interface Metric<T extends MetricBase = MetricBase> {
+        run<U>(fn: (span: Span<T>) => U, options?: SpanOptions): U
+    }
+
+    interface Span<T extends MetricBase = MetricBase> {
         increment(data: { [P in NumericKeys<T>]+?: number }): void
-        run<U>(fn: (span: this) => U, options?: SpanOptions): U
     }
 }

packages/core/src/shared/telemetry/telemetry.ts

@@ -8,7 +8,7 @@ import { toTitleCase } from '../utilities/textUtilities'
 import { getLogger, NullLogger } from '../logger/logger'
 import { FunctionKeys, Functions, getFunctions } from '../utilities/classUtils'
 import { TreeItemContent, TreeNode } from '../treeview/resourceTreeDataProvider'
-import { telemetry, MetricName, VscodeExecuteCommand, Metric } from '../telemetry/telemetry'
+import { telemetry, MetricName, VscodeExecuteCommand, Metric, Span } from '../telemetry/telemetry'
 import globals from '../extensionGlobals'
 import { ToolkitError } from '../errors'
 import crypto from 'crypto'
@@ -498,7 +498,7 @@ function getInstrumenter(
     return <T extends Callback>(fn: T, ...args: Parameters<T>) =>
         span.run(
             (span) => {
-                ;(span as Metric<VscodeExecuteCommand>).record({
+                ;(span as Span<VscodeExecuteCommand>).record({
                     command: id.id,
                     debounceCount,
                     ...fields,