Merge master into feature/emr

Author: aws-toolkit-automation

docs/telemetry.md

@@ -293,3 +293,38 @@ traceId: 'aaaaa-aaaaa-aaaaa-aaaaa-aaaaa'
 allowing us to look up `traceId=aaaaa-aaaaa-aaaaa-aaaaa-aaaaa` in our telemetry instance and find all the related events.
 
 For more information visit the OpenTelemetry documentation on traces: https://opentelemetry.io/docs/concepts/signals/traces/
+
+### Manual Trace ID Instrumentation
+
+In certain scenarios you may need to manually instrument disjoint flows to track how a `traceId` propagates through them. e.g.
+
+1. Measuring the time it takes for a message to travel from Amazon Q chat, through VS Code, and back to the customer.
+2. Determining the duration for Amazon Q inline to display a message to the user.
+
+In these cases, where there isn't a direct hierarchy of function calls, manual instrumentation of the `traceId` is necessary.
+
+#### Implementation Options
+
+#### 1. When not currently running in a span
+
+If you're not within an active span and you know the `traceId` you want to use:
+
+```javascript
+telemetry.withTraceId(() => {
+    // Code to be executed within this trace
+}, 'myTraceId')
+```
+
+This method wraps the provided function with the specified traceId
+
+#### 2. When currently running in a span
+
+If you're already executing within a span (e.g., vscode_executeCommand) and you know the traceId you want to use:
+
+```javascript
+telemetry.record({
+    traceId: 'myTraceId',
+})
+```
+
+This approach records the traceId for the current span and all future spans within the same execution context.

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

@@ -353,22 +353,7 @@ 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 {
-        const initTraceId = (callback: () => T): T => {
-            /**
-             * Generate a new traceId if one doesn't exist.
-             * This ensures the traceId is created before the span,
-             * allowing it to propagate to all child telemetry metrics.
-             */
-            if (!this.attributes?.traceId) {
-                return this.runRoot(() => {
-                    this.record({ traceId: randomUUID() })
-                    return callback()
-                })
-            }
-            return callback()
-        }
-
-        return initTraceId(() => {
+        return this.withTraceId(() => {
             const span = this.createSpan(name, options).start()
             const frame = this.switchContext(span)
 
@@ -396,7 +381,34 @@ export class TelemetryTracer extends TelemetryBase {
                 span.stop(e)
                 throw e
             }
-        })
+        }, this.attributes?.traceId ?? randomUUID())
+    }
+
+    /**
+     * **Use {@link run} for most scenarios. Only use this method when you have a specific, pre-existing trace ID that you need to instrument.**
+     *
+     * Associates a known trace ID with subsequent telemetry events in the provided callback,
+     * enabling correlation of events from multiple disjoint sources (e.g., webview, VSCode, partner team code).
+     *
+     * The difference between this and using telemetry.record({traceId: 'foo'}) directly is that this will create an active
+     * span if one doesn't already exist. If you already know you're operating in a span (like vscode_executeCommand) then use
+     * telemetry.record directly.
+     *
+     * Records traceId iff this metric is not already associated with a trace
+     */
+    withTraceId<T>(callback: () => T, traceId: string): T {
+        /**
+         * Generate a new traceId if one doesn't exist.
+         * This ensures the traceId is created before the span,
+         * allowing it to propagate to all child telemetry metrics.
+         */
+        if (!this.attributes?.traceId) {
+            return this.runRoot(() => {
+                this.record({ traceId })
+                return callback()
+            })
+        }
+        return callback()
     }
 
     /**

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

@@ -630,4 +630,56 @@ describe('TelemetryTracer', function () {
             })
         })
     })
+
+    describe('withTraceId()', function () {
+        const expectedId = 'foo-foo-foo-foo-foo'
+
+        beforeEach(() => {
+            sandbox.stub(crypto, 'randomUUID').returns(expectedId)
+        })
+
+        it('uses trace id', function () {
+            telemetry.withTraceId(() => {
+                telemetry.vscode_viewLogs.emit({
+                    result: 'Succeeded',
+                })
+            }, expectedId)
+
+            const viewLogsEvent = getMetrics('vscode_viewLogs')[0]
+            assert.deepStrictEqual(viewLogsEvent.traceId, expectedId)
+        })
+
+        it('does not use traceId when in a span', function () {
+            telemetry.vscode_executeCommand.run(() => {
+                telemetry.withTraceId(() => {
+                    telemetry.vscode_viewLogs.emit({
+                        result: 'Succeeded',
+                    })
+                }, 'testTraceId')
+            })
+
+            const executeCommandEvent = getMetrics('vscode_executeCommand')[0]
+            assert.deepStrictEqual(executeCommandEvent.traceId, expectedId)
+
+            const viewLogsEvent = getMetrics('vscode_viewLogs')[0]
+            assert.deepStrictEqual(viewLogsEvent.traceId, expectedId)
+        })
+
+        it('passes traceId through regular functions', function () {
+            telemetry.withTraceId(() => {
+                function foo() {
+                    function fi() {
+                        telemetry.vscode_viewLogs.emit({
+                            result: 'Succeeded',
+                        })
+                    }
+                    fi()
+                }
+                foo()
+            }, expectedId)
+
+            const viewLogsEvent = getMetrics('vscode_viewLogs')[0]
+            assert.deepStrictEqual(viewLogsEvent.traceId, expectedId)
+        })
+    })
 })