refactor(telemetry): add context to errors

The withTelemetryContext() decorator adds context to telemetry
but not thrown errors.

Solution:

Adding this decorator to a method will add context to any thrown exceptions.
This is helpful in telemetry as it will provide information about the caller.

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

Author: nkomonen-amazon

docs/telemetry.md

@@ -142,7 +142,10 @@ 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."_
 
@@ -194,6 +197,37 @@ function thisFailsSometimes(num: number) {
 outerB()
 ```
 
+Additionalyl the `@withTelemetryContext()` decorator can be added to methods to do the same as above, but with a cleaner syntax.
+
+```typescript
+class MyClass {
+    @withTelemetryContext({ name: 'outerA', class: 'MyClass' })
+    public outerA() {
+        thisFailsSometimes(1)
+    }
+
+    @withTelemetryContext({ name: 'outerB', class: 'MyClass' })
+    public outerB() {
+        thisFailsSometimes(0)
+    }
+
+    private thisFailsSometimes(num: number) {
+        return telemetry.my_Metric.run(() => {
+            telemetry.record({ theCallStack: asStringifiedStack(telemetry.getFunctionStack())})
+            if (number === 0) {
+                throw Error('Cannot be 0')
+            }
+            ...
+        }, { functionId: { name: 'thisFailsSometimes' }})
+    }
+
+}
+
+
+// Results in a metric: { theCallStack: 'MyClass#outerB,thisFailsSometimes', result: 'Failed' }
+outerB()
+```
+
 ### Important Notes
 
 -   If a nested function does not use a `run()` then it will not be part of the call stack.
@@ -261,6 +295,8 @@ outerB()
     }
     ```
 
+-   Using `@withTelemetryContext` will wrap errors with the functionId properties to add more context
+
 ## Tracing Telemetry Events
 
 All telemetry events include a traceId in addition to other attributes. Traceids allow for improved tracking and correlation of related events across a single operation or user flow.

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

@@ -26,9 +26,10 @@ import { isValidationExemptMetric } from './exemptMetrics'
 import { isAmazonQ, isCloud9, isSageMaker } from '../../shared/extensionUtilities'
 import { randomUUID } from '../crypto'
 import { ClassToInterfaceType } from '../utilities/tsUtils'
-import { FunctionEntry, type TelemetryTracer } from './spans'
+import { FunctionEntry } from './spans'
 import { telemetry } from './telemetry'
 import { v5 as uuidV5 } from 'uuid'
+import { ToolkitError } from '../errors'
 
 const legacySettingsTelemetryValueDisable = 'Disable'
 const legacySettingsTelemetryValueEnable = 'Enable'
@@ -341,7 +342,7 @@ export function getOperatingSystem(): OperatingSystem {
 
 /**
  * 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:
  *
@@ -376,8 +377,18 @@ export function withTelemetryContext(functionId: FunctionEntry) {
         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)
+                    try {
+                        // DEVELOPERS: Set a breakpoint here and step in to it to debug the original function
+                        const result = originalMethod.call(this, ...args)
+                        if (result instanceof Promise) {
+                            return result.catch((e) => {
+                                throw addContextToError(e, functionId)
+                            }) as Return
+                        }
+                        return result
+                    } catch (e) {
+                        throw addContextToError(e, functionId)
+                    }
                 },
                 {
                     emit: false,
@@ -388,4 +399,10 @@ export function withTelemetryContext(functionId: FunctionEntry) {
         return decoratedMethod
     }
     return decorator
+
+    function addContextToError(e: unknown, functionId: FunctionEntry) {
+        return ToolkitError.chain(e, `ctx: ${functionId.class ? functionId.class + '#' : ''}${functionId.name}`, {
+            code: functionId.class,
+        })
+    }
 }

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'
@@ -588,6 +588,64 @@ describe('TelemetryTracer', function () {
                 inst.doesNotEmit()
                 assertTelemetry('function_call', [])
             })
+
+            class TestThrows {
+                @withTelemetryContext({ name: 'throwsError', class: 'TestThrows' })
+                throwsError() {
+                    throw arbitraryError
+                }
+
+                @withTelemetryContext({ name: 'throwsError' })
+                throwsErrorButNoClass() {
+                    throw arbitraryError
+                }
+
+                @withTelemetryContext({ name: 'throwsError' })
+                async throwsAsyncError() {
+                    throw arbitraryError
+                }
+            }
+            const arbitraryError = new Error('arbitrary error')
+
+            it(`withTelemetryContext wraps errors with function id context`, async function () {
+                const inst = new TestThrows()
+                assert.throws(
+                    () => inst.throwsError(),
+                    (e) => {
+                        if (!(e instanceof ToolkitError)) {
+                            return false
+                        }
+                        const id = getErrorId(e)
+                        const message = e.message
+                        const cause = e.cause
+                        return id === 'TestThrows' && message === 'ctx: throwsError' && cause === arbitraryError
+                    }
+                )
+                assert.throws(
+                    () => inst.throwsErrorButNoClass(),
+                    (e) => {
+                        if (!(e instanceof ToolkitError)) {
+                            return false
+                        }
+                        const id = getErrorId(e)
+                        const message = e.message
+                        const cause = e.cause
+                        return id === 'Error' && message === 'ctx: throwsError' && cause === arbitraryError
+                    }
+                )
+                await assert.rejects(
+                    () => inst.throwsAsyncError(),
+                    (e) => {
+                        if (!(e instanceof ToolkitError)) {
+                            return false
+                        }
+                        const id = getErrorId(e)
+                        const message = e.message
+                        const cause = e.cause
+                        return id === 'Error' && message === 'ctx: throwsError' && cause === arbitraryError
+                    }
+                )
+            })
         })
     })