Merge master into feature/emr

Author: aws-toolkit-automation

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

@@ -24,6 +24,7 @@ import {
 } from '../errors'
 import { entries, NumericKeys } from '../utilities/tsUtils'
 import { PerformanceTracker } from '../performance/performance'
+import { randomUUID } from '../crypto'
 
 const AsyncLocalStorage: typeof AsyncLocalStorageClass =
     require('async_hooks').AsyncLocalStorage ??
@@ -145,6 +146,7 @@ export class TelemetrySpan<T extends MetricBase = MetricBase> {
         passive: true,
         requiredMetadata: [],
     }
+    private readonly metricId: string
 
     /**
      * These fields appear on the base metric instead of the 'metadata' and
@@ -166,6 +168,10 @@ export class TelemetrySpan<T extends MetricBase = MetricBase> {
                 this.definition.trackPerformance && (options?.emit ?? false) // only track the performance if we are also emitting
             ),
         }
+
+        this.metricId = randomUUID()
+        // forced to cast to any since apparently even though <T extends MetricBase>, Partial<T> doesn't guarentee that metricId is available
+        this.record({ metricId: this.metricId } as any)
     }
 
     public get startTime(): Date | undefined {
@@ -261,6 +267,10 @@ export class TelemetrySpan<T extends MetricBase = MetricBase> {
         }
     }
 
+    getMetricId() {
+        return this.metricId
+    }
+
     /**
      * Creates a copy of the span with an uninitialized start time.
      */
@@ -335,33 +345,50 @@ export class TelemetryTracer extends TelemetryBase {
      * reverted after the execution completes.
      */
     public run<T, U extends MetricName>(name: U, fn: (span: Metric<MetricShapes[U]>) => T, options?: SpanOptions): T {
-        const span = this.createSpan(name, options).start()
-        const frame = this.switchContext(span)
+        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(() => {
+            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)
+            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)
+
+                if (result instanceof Promise) {
+                    return result
+                        .then((v) => (span.stop(), v))
+                        .catch((e) => {
+                            span.stop(e)
+                            throw e
+                        }) as unknown as T
+                }
 
-            if (result instanceof Promise) {
+                span.stop()
                 return result
-                    .then((v) => (span.stop(), v))
-                    .catch((e) => {
-                        span.stop(e)
-                        throw e
-                    }) as unknown as T
+            } catch (e) {
+                span.stop(e)
+                throw e
             }
-
-            span.stop()
-            return result
-        } catch (e) {
-            span.stop(e)
-            throw e
-        }
+        })
     }
 
     /**
@@ -437,7 +464,7 @@ export class TelemetryTracer extends TelemetryBase {
     private createSpan(name: string, options?: SpanOptions): TelemetrySpan {
         const span = new TelemetrySpan(name, options).record(this.attributes ?? {})
         if (this.activeSpan && this.activeSpan.name !== rootSpanName) {
-            return span.record({ parentMetric: this.activeSpan.name } satisfies { parentMetric: string } as any)
+            return span.record({ parentId: this.activeSpan.getMetricId() })
         }
 
         return span

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

@@ -315,6 +315,11 @@
             "name": "totalFiles",
             "type": "int",
             "description": "The total number of files being sent to Amazon Q"
+        },
+        {
+            "name": "traceId",
+            "type": "string",
+            "description": "The unique identifier of a trace"
         }
     ],
     "metrics": [
@@ -1178,6 +1183,21 @@
                     "required": true
                 }
             ]
+        },
+        {
+            "name": "trace_event",
+            "description": "Represents the root trace spanning multiple operations",
+            "passive": true,
+            "metadata": [
+                {
+                    "type": "traceId",
+                    "required": true
+                },
+                {
+                    "type": "name",
+                    "required": true
+                }
+            ]
         }
     ]
 }

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

@@ -15,6 +15,7 @@ import { withTelemetryContext } from '../../../shared/telemetry/util'
 import { SinonSandbox } from 'sinon'
 import sinon from 'sinon'
 import { stubPerformance } from '../../utilities/performance'
+import * as crypto from '../../../shared/crypto'
 
 describe('TelemetrySpan', function () {
     let clock: ReturnType<typeof installFakeClock>
@@ -291,7 +292,17 @@ describe('TelemetryTracer', function () {
         })
 
         describe('nested run()', function () {
+            let uuidStub: sinon.SinonStub
             const nestedName = 'nested_metric' as MetricName
+            const testId = 'foo-foo-foo-foo-foo'
+            const flowName = 'testTraceFlow'
+
+            beforeEach(() => {
+                uuidStub = sandbox.stub(crypto, 'randomUUID')
+
+                // in the first call we set the trace id in subsequent calls we get the span ids
+                uuidStub.returns(testId)
+            })
 
             it('can record metadata in nested spans', function () {
                 tracer.run(metricName, (span1) => {
@@ -310,18 +321,18 @@ describe('TelemetryTracer', function () {
             it('removes spans when exiting an execution context', function () {
                 tracer.run(metricName, () => {
                     tracer.run(nestedName, () => {
-                        assert.strictEqual(tracer.spans.length, 2)
+                        assert.strictEqual(tracer.spans.length, 3)
                     })
 
-                    assert.strictEqual(tracer.spans.length, 1)
+                    assert.strictEqual(tracer.spans.length, 2)
                 })
             })
 
             it('adds spans during a nested execution, closing them when after', function () {
                 tracer.run(metricName, () => {
-                    tracer.run(nestedName, () => assert.strictEqual(tracer.spans.length, 2))
-                    tracer.run(nestedName, () => assert.strictEqual(tracer.spans.length, 2))
-                    assert.strictEqual(tracer.spans.length, 1)
+                    tracer.run(nestedName, () => assert.strictEqual(tracer.spans.length, 3))
+                    tracer.run(nestedName, () => assert.strictEqual(tracer.spans.length, 3))
+                    assert.strictEqual(tracer.spans.length, 2)
                 })
 
                 assert.strictEqual(tracer.spans.length, 0)
@@ -330,16 +341,102 @@ describe('TelemetryTracer', function () {
             it('supports nesting the same event name', function () {
                 tracer.run(metricName, () => {
                     tracer.run(metricName, () => {
-                        assert.strictEqual(tracer.spans.length, 2)
-                        assert.ok(tracer.spans.every((s) => s.name === metricName))
+                        assert.strictEqual(tracer.spans.length, 3)
+                        assert.ok(tracer.spans.filter((m) => m.name !== 'root').every((s) => s.name === metricName))
                     })
                 })
             })
 
-            it('attaches the parent event name to the child span', function () {
+            it('attaches the parent id to the child span', function () {
                 tracer.run(metricName, () => tracer.run(nestedName, () => {}))
                 assertTelemetry(metricName, { result: 'Succeeded' })
-                assertTelemetry(nestedName, { result: 'Succeeded', parentMetric: metricName } as any)
+                assertTelemetry(nestedName, { result: 'Succeeded', parentId: testId } as any)
+            })
+
+            it('should set trace id', function () {
+                telemetry.trace_event.run((span) => {
+                    span.record({ name: flowName })
+                    assert.deepStrictEqual(telemetry.activeSpan?.getMetricId(), testId)
+                })
+                const event = getMetrics('trace_event')
+                assert.deepStrictEqual(event[0].traceId, testId)
+                assert.deepStrictEqual(event[0].name, 'testTraceFlow')
+            })
+
+            it('trace id is propogated to children', function () {
+                const metricsIds = {
+                    trace_event: {
+                        metricId: 'traceEvent',
+                        traceId: testId,
+                        parentId: undefined,
+                    },
+                    amazonq_startConversation: {
+                        metricId: 'amazonq_startConversation',
+                        traceId: testId,
+                        parentId: 'traceEvent',
+                    },
+                    amazonq_addMessage: {
+                        metricId: 'amazonq_addMessage',
+                        traceId: testId,
+                        parentId: 'amazonq_startConversation',
+                    },
+                    vscode_executeCommand: {
+                        metricId: 'vscode_executeCommand',
+                        traceId: testId,
+                        parentId: 'traceEvent',
+                    },
+                    amazonq_enterFocusConversation: {
+                        metricId: 'amazonq_enterFocusConversation',
+                        traceId: testId,
+                        parentId: 'vscode_executeCommand',
+                    },
+                    amazonq_exitFocusConversation: {
+                        metricId: 'amazonq_exitFocusConversation',
+                        traceId: testId,
+                        parentId: 'amazonq_enterFocusConversation',
+                    },
+                    amazonq_closeChat: {
+                        metricId: 'amazonq_closeChat',
+                        traceId: testId,
+                        parentId: 'traceEvent',
+                    },
+                }
+
+                /**
+                 * randomUUID calls:
+                 * The first is called on the root event that never gets emitted
+                 * The second is called when generating the traceId
+                 * The rest are called when generating the spanIds
+                 */
+                uuidStub.onCall(1).returns(testId)
+                let index = 2
+                for (const v of Object.values(metricsIds)) {
+                    uuidStub.onCall(index).returns(v.metricId)
+                    index++
+                }
+
+                telemetry.trace_event.run(() => {
+                    telemetry.amazonq_startConversation.run(() => {
+                        telemetry.amazonq_addMessage.run(() => {})
+                    })
+                    telemetry.vscode_executeCommand.run(() => {
+                        telemetry.amazonq_enterFocusConversation.run(() => {
+                            telemetry.amazonq_exitFocusConversation.run(() => {})
+                        })
+                    })
+                    telemetry.amazonq_closeChat.emit({
+                        result: 'Succeeded',
+                    })
+                })
+
+                const spanEntries = Object.entries(metricsIds)
+                for (let x = 0; x < spanEntries.length; x++) {
+                    const [metricName, { metricId, traceId, parentId }] = spanEntries[x]
+                    const metric = getMetrics(metricName as keyof MetricShapes)[0] as any
+                    assert.deepStrictEqual(metric.traceId, traceId)
+                    assert.deepStrictEqual(metric.metricId, metricId)
+                    assert.deepStrictEqual(metric.parentId, parentId)
+                }
             })
         })