refactor(telemetry): @withTelemetryContext can emit with source

Now the @withTelemetryContext decorator can have a method emit when called
and will also add the `source` field automatically.

By default it does not emit.

Signed-off-by: nkomonen-amazon <[email protected]>

Author: nkomonen-amazon

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

@@ -26,7 +26,7 @@ import { isValidationExemptMetric } from './exemptMetrics'
 import { isAmazonQ, isCloud9, isSageMaker } from '../../shared/extensionUtilities'
 import { randomUUID } from '../crypto'
 import { ClassToInterfaceType } from '../utilities/tsUtils'
-import { FunctionEntry } from './spans'
+import { asStringifiedStack, FunctionEntry } from './spans'
 import { telemetry } from './telemetry'
 import { v5 as uuidV5 } from 'uuid'
 import { ToolkitError } from '../errors'
@@ -369,30 +369,37 @@ export function getOperatingSystem(): OperatingSystem {
  * }
  * ```
  */
-export function withTelemetryContext(functionId: FunctionEntry) {
+export function withTelemetryContext(opts: FunctionEntry & { emit?: boolean }) {
     function decorator<This, Args extends any[], Return>(
         originalMethod: (this: This, ...args: Args) => Return,
         _context: ClassMethodDecoratorContext // we dont need this currently but it keeps the compiler happy
     ) {
         function decoratedMethod(this: This, ...args: Args): Return {
             return telemetry.function_call.run(
-                () => {
+                (span) => {
                     try {
-                        // DEVELOPERS: Set a breakpoint here and step in to it to debug the original function
+                        span.record({
+                            functionName: opts.name,
+                            className: opts.class,
+                            source: asStringifiedStack(telemetry.getFunctionStack()),
+                        })
+
+                        // DEVELOPERS: Set a breakpoint here and step in and debug the original function
                         const result = originalMethod.call(this, ...args)
+
                         if (result instanceof Promise) {
                             return result.catch((e) => {
-                                throw addContextToError(e, functionId)
+                                throw addContextToError(e, opts)
                             }) as Return
                         }
                         return result
                     } catch (e) {
-                        throw addContextToError(e, functionId)
+                        throw addContextToError(e, opts)
                     }
                 },
                 {
-                    emit: false,
-                    functionId: functionId,
+                    emit: opts.emit !== undefined ? opts.emit : false,
+                    functionId: opts,
                 }
             )
         }
@@ -401,7 +408,7 @@ export function withTelemetryContext(functionId: FunctionEntry) {
     return decorator
 
     function addContextToError(e: unknown, functionId: FunctionEntry) {
-        return ToolkitError.chain(e, `ctx: ${functionId.class ? functionId.class + '#' : ''}${functionId.name}`, {
+        return ToolkitError.chain(e, `ctx: ${functionId.name}`, {
             code: functionId.class,
         })
     }

packages/core/src/shared/telemetry/vscodeTelemetry.json

@@ -1116,6 +1116,10 @@
                 {
                     "type": "className",
                     "required": false
+                },
+                {
+                    "type": "source",
+                    "required": false
                 }
             ],
             "passive": true

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

@@ -576,6 +576,29 @@ describe('TelemetryTracer', function () {
                 )
             })
 
+            class DecoratorEmits {
+                @withTelemetryContext({ name: 'outerFunc', class: 'DecoratorEmits' })
+                outerFunc() {
+                    this.decoratorEmits()
+                }
+                @withTelemetryContext({ name: 'decoratorEmits', class: 'DecoratorEmits', emit: true })
+                decoratorEmits() {
+                    return true
+                }
+            }
+
+            it(`will emit the method call when emit is true`, function () {
+                const inst = new DecoratorEmits()
+                inst.outerFunc()
+                assertTelemetry('function_call', [
+                    {
+                        functionName: 'decoratorEmits',
+                        className: 'DecoratorEmits',
+                        source: 'DecoratorEmits#outerFunc,decoratorEmits',
+                    },
+                ])
+            })
+
             class TestEmit {
                 @withTelemetryContext({ name: 'doesNotEmit', class: 'TestEmit' })
                 doesNotEmit() {