refactor(telemetry): add functionality to @withTelemetryContext

docs/telemetry.md

@@ -142,20 +142,23 @@ Finally, if `setupStep2()` was the thing that failed we would see a metric like:
 
 ## Adding a "Stack Trace" to your metric
 
-### Problem
+When errors are thrown we do not attach the stack trace in telemetry. We only know about the error itself, but
+not the path it took to get there. We sometimes need this stack trace to debug, and only have telemetry to get insight on what happened since we do not have access to logs.
+
+### Scenario
 
 Common example: _"I have a function, `thisFailsSometimes()` that is called in multiple places. The function sometimes fails, I know from telemetry, but I do not know if it is failing when it is a specific caller. If I knew the call stack/trace that it took to call my function that would help me debug."_
 
 ```typescript
-function outerA() {
+function runsSuccessfully() {
     thisFailsSometimes(1) // this succeeds
 }
 
-function outerB() {
+function thisThrows() {
     thisFailsSometimes(0) // this fails
 }
 
-function thisFailsSometimes(num: number) {
+function failsDependingOnInput(num: number) {
     return telemetry.my_Metric.run(() => {
         if (number === 0) {
             throw Error('Cannot be 0')
@@ -167,31 +170,61 @@ function thisFailsSometimes(num: number) {
 
 ### Solution
 
-Add a value to `function` in the options of a `run()`. This will result in a stack of functions identifiers that were previously called
-before `thisFailsSometimes()` was run. You can then retrieve the stack in the `run()` of your final metric using `getFunctionStack()`.
+On class methods, use the `@withTelemetryContext()` decorator to add context to the execution. Depending on the args set, it provides features like emitting the result, or adding it's context to errors.
+
+> NOTE: Decorators are currently only supported for methods and not functions
+
+```typescript
+class MyClass {
+    @withTelemetryContext({ name: 'runsSuccessfully', class: 'MyClass' })
+    public runsSuccessfully() {
+        failsDependingOnInput(1)
+    }
+
+    @withTelemetryContext({ name: 'thisThrows', class: 'MyClass', errorCtx: true })
+    public thisThrows() {
+        failsDependingOnInput(0)
+    }
+
+    @withTelemetryContext({ name: 'failsDependingOnInput' class: 'MyClass', emit: true, errorCtx: true})
+    private failsDependingOnInput(num: number) {
+        if (number === 0) {
+            throw Error('Cannot be 0')
+        }
+        ...
+    }
+}
+
+// Results in a metric: { source: 'MyClass#thisThrows,failsDependingOnInput', result: 'Failed' }
+// Results in an error that has context about the methods that lead up to it.
+new MyClass().thisThrows()
+```
+
+Separately if you must use a function, add a value to `function` in the options of a `run()`. This will result in a stack of functions identifiers that were previously called
+before `failsDependingOnInput()` was run. You can then retrieve the stack in the `run()` of your final metric using `getFunctionStack()`.
 
 ```typescript
-function outerA() {
-    telemetry.my_Metric.run(() => thisFailsSometimes(1), { functionId: { name: 'outerA' }})
+function runsSuccessfully() {
+    telemetry.my_Metric.run(() => failsDependingOnInput(1), { functionId: { name: 'runsSuccessfully' }})
 }
 
-function outerB() {
-    telemetry.my_Metric.run(() => thisFailsSometimes(0), { functionId: { source: 'outerB' }})
+function thisThrows() {
+    telemetry.my_Metric.run(() => failsDependingOnInput(0), { functionId: { source: 'thisThrows' }})
 }
 
-function thisFailsSometimes(num: number) {
+function failsDependingOnInput(num: number) {
     return telemetry.my_Metric.run(() => {
         telemetry.record({ theCallStack: asStringifiedStack(telemetry.getFunctionStack())})
         if (number === 0) {
             throw Error('Cannot be 0')
         }
         ...
-    }, { functionId: { name: 'thisFailsSometimes' }})
+    }, { functionId: { name: 'failsDependingOnInput' }})
 }
 
-// Results in a metric: { theCallStack: 'outerB:thisFailsSometimes', result: 'Failed' }
-// { theCallStack: 'outerB:thisFailsSometimes' } implies 'outerB' was run first, then 'thisFailsSometimes'. See docstrings for more info.
-outerB()
+// Results in a metric: { theCallStack: 'thisThrows:failsDependingOnInput', result: 'Failed' }
+// { theCallStack: 'thisThrows:failsDependingOnInput' } implies 'thisThrows' was run first, then 'failsDependingOnInput'. See docstrings for more info.
+thisThrows()
 ```
 
 ### Important Notes
@@ -216,25 +249,6 @@ outerB()
     c() // result: 'a:c', note that 'b' is not included
     ```
 
--   If you are using `run()` with a class method, you can also add the class to the entry for more context
-
-    ```typescript
-    class A {
-        a() {
-            return telemetry.my_Metric.run(() => this.b(), { functionId: { name: 'a', class: 'A' } })
-        }
-
-        b() {
-            return telemetry.my_Metric.run(() => asStringifiedStack(telemetry.getFunctionStack()), {
-                functionId: { name: 'b', class: 'A' },
-            })
-        }
-    }
-
-    const inst = new A()
-    inst.a() // 'A#a,b'
-    ```
-
 -   If you do not want your `run()` to emit telemetry, set `emit: false` in the options
 
     ```typescript

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

@@ -26,9 +26,10 @@ import { isValidationExemptMetric } from './exemptMetrics'
 import { isAmazonQ, isCloud9, isSageMaker } from '../../shared/extensionUtilities'
 import { isUuid, randomUUID } from '../crypto'
 import { ClassToInterfaceType } from '../utilities/tsUtils'
-import { FunctionEntry, type TelemetryTracer } from './spans'
+import { asStringifiedStack, FunctionEntry } from './spans'
 import { telemetry } from './telemetry'
 import { v5 as uuidV5 } from 'uuid'
+import { ToolkitError } from '../errors'
 
 const legacySettingsTelemetryValueDisable = 'Disable'
 const legacySettingsTelemetryValueEnable = 'Enable'
@@ -365,9 +366,10 @@ export function getOperatingSystem(): OperatingSystem {
     }
 }
 
+type TelemetryContextArgs = FunctionEntry & { emit?: boolean; errorCtx?: boolean }
 /**
  * Decorator that simply wraps the method with a non-emitting telemetry `run()`, automatically
- * `record()`ing the provided function id for later use by {@link TelemetryTracer.getFunctionStack()}
+ * `record()`ing the provided function id for later use by TelemetryTracer.getFunctionStack()
  *
  * This saves us from needing to wrap the entire function:
  *
@@ -393,25 +395,60 @@ export function getOperatingSystem(): OperatingSystem {
  *     }
  * }
  * ```
+ *
+ * @param opts.name The name of the function
+ * @param opts.class The class name of the function
+ * @param opts.emit Whether or not to emit the telemetry event (default: false)
+ * @param opts.errorCtx Whether or not to add the error context to the error (default: false)
  */
-export function withTelemetryContext(functionId: FunctionEntry) {
+export function withTelemetryContext(opts: TelemetryContextArgs) {
+    const shouldErrorCtx = opts.errorCtx !== undefined ? opts.errorCtx : false
     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(
-                () => {
-                    // DEVELOPERS: Set a breakpoint here and step in to it to debug the original function
-                    return originalMethod.call(this, ...args)
+                (span) => {
+                    try {
+                        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) => {
+                                if (shouldErrorCtx) {
+                                    throw addContextToError(e, opts)
+                                }
+                                throw e
+                            }) as Return
+                        }
+                        return result
+                    } catch (e) {
+                        if (shouldErrorCtx) {
+                            throw addContextToError(e, opts)
+                        }
+                        throw e
+                    }
                 },
                 {
-                    emit: false,
-                    functionId: functionId,
+                    emit: opts.emit !== undefined ? opts.emit : false,
+                    functionId: { name: opts.name, class: opts.class },
                 }
             )
         }
         return decoratedMethod
     }
     return decorator
+
+    function addContextToError(e: unknown, functionId: FunctionEntry) {
+        return ToolkitError.chain(e, `ctx: ${functionId.name}`, {
+            code: functionId.class,
+        })
+    }
 }

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

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

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

@@ -4,7 +4,7 @@
  */
 
 import assert from 'assert'
-import { ToolkitError } from '../../../shared/errors'
+import { getErrorId, ToolkitError } from '../../../shared/errors'
 import { asStringifiedStack, FunctionEntry, TelemetrySpan, TelemetryTracer } from '../../../shared/telemetry/spans'
 import { MetricName, MetricShapes, telemetry } from '../../../shared/telemetry/telemetry'
 import { assertTelemetry, getMetrics, installFakeClock } from '../../testUtil'
@@ -575,19 +575,112 @@ describe('TelemetryTracer', function () {
                     'A#a1,a2:B#b1,b2'
                 )
             })
+        })
+    })
 
-            class TestEmit {
-                @withTelemetryContext({ name: 'doesNotEmit', class: 'TestEmit' })
-                doesNotEmit() {
-                    return
-                }
+    describe('withTelemetryContext', async function () {
+        class TestEmit {
+            @withTelemetryContext({ name: 'doesNotEmit', class: 'TestEmit' })
+            doesNotEmit() {
+                return
             }
 
-            it(`withTelemetryContext does not emit an event on its own`, function () {
-                const inst = new TestEmit()
-                inst.doesNotEmit()
-                assertTelemetry('function_call', [])
-            })
+            @withTelemetryContext({ name: 'doesEmit', class: 'TestEmit', emit: false })
+            doesEmit() {
+                return this.doesEmitNested()
+            }
+
+            @withTelemetryContext({ name: 'doesEmitNested', class: 'TestEmit', emit: true })
+            doesEmitNested() {
+                return
+            }
+        }
+
+        it(`does NOT emit an event if not explicitly set`, function () {
+            const inst = new TestEmit()
+            inst.doesNotEmit()
+            assertTelemetry('function_call', [])
+        })
+
+        it(`does emit an event on its own when explicitly set`, function () {
+            const inst = new TestEmit()
+            inst.doesEmit()
+            assertTelemetry('function_call', [
+                {
+                    functionName: 'doesEmitNested',
+                    className: 'TestEmit',
+                    source: 'TestEmit#doesEmit,doesEmitNested',
+                },
+            ])
+        })
+
+        class TestThrows {
+            @withTelemetryContext({ name: 'throwsError', class: 'TestThrows', errorCtx: true })
+            throwsError() {
+                throw arbitraryError
+            }
+
+            @withTelemetryContext({ name: 'throwsError', errorCtx: true })
+            throwsErrorButNoClass() {
+                throw arbitraryError
+            }
+
+            @withTelemetryContext({ name: 'throwsError', errorCtx: true })
+            async throwsAsyncError() {
+                throw arbitraryError
+            }
+
+            @withTelemetryContext({ name: 'throwsErrorButNoCtx', class: 'TestThrows' })
+            throwsErrorButNoCtx() {
+                throw arbitraryError
+            }
+        }
+        const arbitraryError = new Error('arbitrary error')
+
+        it(`wraps errors with function id context`, async function () {
+            const inst = new TestThrows()
+            assert.throws(
+                () => inst.throwsError(),
+                (e) => {
+                    return (
+                        e instanceof ToolkitError &&
+                        getErrorId(e) === 'TestThrows' &&
+                        e.message === 'ctx: throwsError' &&
+                        e.cause === arbitraryError
+                    )
+                }
+            )
+            assert.throws(
+                () => inst.throwsErrorButNoClass(),
+                (e) => {
+                    return (
+                        e instanceof ToolkitError &&
+                        getErrorId(e) === 'Error' &&
+                        e.message === 'ctx: throwsError' &&
+                        e.cause === arbitraryError
+                    )
+                }
+            )
+            await assert.rejects(
+                () => inst.throwsAsyncError(),
+                (e) => {
+                    return (
+                        e instanceof ToolkitError &&
+                        getErrorId(e) === 'Error' &&
+                        e.message === 'ctx: throwsError' &&
+                        e.cause === arbitraryError
+                    )
+                }
+            )
+        })
+
+        it('does not add error context by default', async function () {
+            const inst = new TestThrows()
+
+            assert.throws(
+                () => inst.throwsErrorButNoCtx(),
+                (e) => e === arbitraryError
+            )
         })
     })