telemetry(amazonq): Add client latency metrics for amazon q chat (#5705)

## Problem:
- It's not clear what's meant by cwsprChatTimeToFirstChunk,
cwsprChatTimeBetweenChunks, cwsprChatFullResponseLatency
- We don't have any client latency metrics

## Solution:
- Clarify the descriptions of cwsprChatTimeToFirstChunk,
cwsprChatTimeBetweenChunks, cwsprChatFullResponseLatency
- Implement cwsprTimeToFirstDisplay, cwsprChatTimeBetweenDisplay,
cwsprChatFullDisplayLatency metrics

## Other changes:
- Instead of instrumenting everything via a traceId instead use the
tabId as the unique identifier and store a traceId when we get it inside
of the uiEventRecorder. That way we can associate events from a tab to
the trace without needing to pass the trace around the entire
application.
- Lock stopChatMessageTelemetry on the tabId just in case something
calls that twice for the same event

## Risks:
- All flows from start of chat message -> end of chat message need to be
instrumented otherwise an addChatMessage event wont be sent.
- RIght now:
   - Right click context menu works
   - Normal chat message work
   - Follow up button works
   - Are there any other flows that we need to know about?

Author: jpinkney-aws

packages/core/src/amazonq/messages/chatMessageDuration.ts

@@ -3,91 +3,111 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+import AsyncLock from 'async-lock'
 import { globals } from '../../shared'
 import { telemetry } from '../../shared/telemetry'
 import { Event, uiEventRecorder } from '../util/eventRecorder'
+import { CWCTelemetryHelper } from '../../codewhispererChat/controllers/chat/telemetryHelper'
+import { TabType } from '../webview/ui/storages/tabsStorage'
 
 export class AmazonQChatMessageDuration {
+    private static _asyncLock = new AsyncLock()
+    private static getAsyncLock() {
+        if (!AmazonQChatMessageDuration._asyncLock) {
+            AmazonQChatMessageDuration._asyncLock = new AsyncLock()
+        }
+        return AmazonQChatMessageDuration._asyncLock
+    }
+
     /**
      * Record the initial requests in the chat message flow
      */
-    static startChatMessageTelemetry(msg: { traceId: string; startTime: number; trigger?: string }) {
-        const { traceId, startTime, trigger } = msg
+    static startChatMessageTelemetry(msg: { traceId: string; startTime: number; tabID: string; trigger?: string }) {
+        const { traceId, startTime, tabID, trigger } = msg
 
-        uiEventRecorder.set(traceId, {
+        uiEventRecorder.set(tabID, {
+            traceId,
             events: {
                 chatMessageSent: startTime,
-            },
-        })
-        uiEventRecorder.set(traceId, {
-            events: {
                 editorReceivedMessage: globals.clock.Date.now(),
             },
         })
         if (trigger) {
-            uiEventRecorder.set(traceId, {
+            uiEventRecorder.set(tabID, {
                 trigger,
             })
         }
+        CWCTelemetryHelper.instance.setDisplayTimeForChunks(tabID, startTime)
     }
 
     /**
      * Stop listening to all incoming events and emit what we've found
      */
-    static stopChatMessageTelemetry(msg: { traceId: string }) {
-        const { traceId } = msg
+    static stopChatMessageTelemetry(msg: { tabID: string; time: number; tabType: TabType }) {
+        const { tabID, time, tabType } = msg
 
         // We can't figure out what trace this event was associated with
-        if (!traceId) {
+        if (!tabID || tabType !== 'cwc') {
             return
         }
 
-        uiEventRecorder.set(traceId, {
-            events: {
-                messageDisplayed: globals.clock.Date.now(),
-            },
-        })
-
-        const metrics = uiEventRecorder.get(traceId)
+        // Lock the tab id just in case another event tries to trigger this
+        void AmazonQChatMessageDuration.getAsyncLock().acquire(tabID, () => {
+            const metrics = uiEventRecorder.get(tabID)
+            if (!metrics) {
+                return
+            }
 
-        // get events sorted by the time they were created
-        const events = Object.entries(metrics.events)
-            .map((x) => ({
-                event: x[0],
-                duration: x[1],
-            }))
-            .sort((a, b) => {
-                return a.duration - b.duration
+            uiEventRecorder.set(tabID, {
+                events: {
+                    messageDisplayed: time,
+                },
             })
 
-        const chatMessageSentTime = events[events.length - 1].duration
-        // Get the total duration by subtracting when the message was displayed and when the chat message was first sent
-        const totalDuration = events[events.length - 1].duration - events[0].duration
+            const displayTime = metrics.events.messageDisplayed
+            const sentTime = metrics.events.chatMessageSent
+            if (!displayTime || !sentTime) {
+                return
+            }
 
-        /**
-         * Find the time it took to get between two metric events
-         */
-        const timings = new Map<Event, number>()
-        for (let i = 1; i < events.length; i++) {
-            const currentEvent = events[i]
-            const previousEvent = events[i - 1]
+            const totalDuration = displayTime - sentTime
 
-            const timeDifference = currentEvent.duration - previousEvent.duration
+            function durationFrom(start: Event, end: Event) {
+                const startEvent = metrics.events[start]
+                const endEvent = metrics.events[end]
+                if (!startEvent || !endEvent) {
+                    return -1
+                }
+                return endEvent - startEvent
+            }
 
-            timings.set(currentEvent.event as Event, timeDifference)
-        }
+            // TODO: handle onContextCommand round trip time
+            if (metrics.trigger !== 'onContextCommand') {
+                telemetry.amazonq_chatRoundTrip.emit({
+                    amazonqChatMessageSentTime: metrics.events.chatMessageSent ?? -1,
+                    amazonqEditorReceivedMessageMs: durationFrom('chatMessageSent', 'editorReceivedMessage') ?? -1,
+                    amazonqFeatureReceivedMessageMs:
+                        durationFrom('editorReceivedMessage', 'featureReceivedMessage') ?? -1,
+                    amazonqMessageDisplayedMs: durationFrom('featureReceivedMessage', 'messageDisplayed') ?? -1,
+                    source: metrics.trigger,
+                    duration: totalDuration,
+                    result: 'Succeeded',
+                    traceId: metrics.traceId,
+                })
+            }
 
-        telemetry.amazonq_chatRoundTrip.emit({
-            amazonqChatMessageSentTime: chatMessageSentTime,
-            amazonqEditorReceivedMessageMs: timings.get('editorReceivedMessage') ?? -1,
-            amazonqFeatureReceivedMessageMs: timings.get('featureReceivedMessage') ?? -1,
-            amazonqMessageDisplayedMs: timings.get('messageDisplayed') ?? -1,
-            source: metrics.trigger,
-            duration: totalDuration,
-            result: 'Succeeded',
-            traceId,
+            CWCTelemetryHelper.instance.emitAddMessage(tabID, totalDuration, metrics.events.chatMessageSent)
+
+            uiEventRecorder.delete(tabID)
         })
+    }
+
+    static updateChatMessageTelemetry(msg: { tabID: string; time: number; tabType: TabType }) {
+        const { tabID, time, tabType } = msg
+        if (!tabID || tabType !== 'cwc') {
+            return
+        }
 
-        uiEventRecorder.delete(traceId)
+        CWCTelemetryHelper.instance.setDisplayTimeForChunks(tabID, time)
     }
 }

packages/core/src/amazonq/util/eventRecorder.ts

@@ -9,10 +9,10 @@ export type Event =
     | 'chatMessageSent' // initial on chat prompt event in the ui
     | 'editorReceivedMessage' // message gets from the chat prompt to VSCode
     | 'featureReceivedMessage' // message gets redirected from VSCode -> Partner team features implementation
-    | 'messageDisplayed' // message gets received in the UI
+    | 'messageDisplayed' // message gets shown in the UI
 
 /**
- * For a given traceID, map an event to a time
+ * For a given tabId, map an event to a time
  *
  * This is used to correlated disjoint events that are happening in different
  * parts of Q Chat.
@@ -22,8 +22,14 @@ export type Event =
  *  - when the feature starts processing the message
  *  - final message rendering
  * and emit those as a final result, rather than having to emit each event individually
+ *
+ * Event timings are generated using Date.now() instead of performance.now() for cross-context consistency.
+ * performance.now() provides timestamps relative to the context's time origin (when the webview or VS Code was opened),
+ * which can lead to inconsistent measurements between the webview and vscode.
+ * Date.now() is more consistent across both contexts
  */
 export const uiEventRecorder = new RecordMap<{
     trigger: string
+    traceId: string
     events: Partial<Record<Event, number>>
 }>()

packages/core/src/amazonq/webview/messages/messageDispatcher.ts

@@ -17,29 +17,35 @@ export function dispatchWebViewMessagesToApps(
     webViewToAppsMessagePublishers: Map<TabType, MessagePublisher<any>>
 ) {
     webview.onDidReceiveMessage((msg) => {
-        if (msg.command === 'ui-is-ready') {
-            /**
-             * ui-is-ready isn't associated to any tab so just record the telemetry event and continue.
-             * This would be equivalent of the duration between "user clicked open q" and "ui has become available"
-             * NOTE: Amazon Q UI is only loaded ONCE. The state is saved between each hide/show of the webview.
-             */
+        switch (msg.command) {
+            case 'ui-is-ready': {
+                /**
+                 * ui-is-ready isn't associated to any tab so just record the telemetry event and continue.
+                 * This would be equivalent of the duration between "user clicked open q" and "ui has become available"
+                 * NOTE: Amazon Q UI is only loaded ONCE. The state is saved between each hide/show of the webview.
+                 */
 
-            telemetry.webview_load.emit({
-                webviewName: 'amazonq',
-                duration: performance.measure(amazonqMark.uiReady, amazonqMark.open).duration,
-                result: 'Succeeded',
-            })
-            performance.clearMarks(amazonqMark.uiReady)
-            performance.clearMarks(amazonqMark.open)
-            return
-        }
-
-        if (msg.type === 'startChatMessageTelemetry') {
-            AmazonQChatMessageDuration.startChatMessageTelemetry(msg)
-            return
-        } else if (msg.type === 'stopChatMessageTelemetry') {
-            AmazonQChatMessageDuration.stopChatMessageTelemetry(msg)
-            return
+                telemetry.webview_load.emit({
+                    webviewName: 'amazonq',
+                    duration: performance.measure(amazonqMark.uiReady, amazonqMark.open).duration,
+                    result: 'Succeeded',
+                })
+                performance.clearMarks(amazonqMark.uiReady)
+                performance.clearMarks(amazonqMark.open)
+                return
+            }
+            case 'start-chat-message-telemetry': {
+                AmazonQChatMessageDuration.startChatMessageTelemetry(msg)
+                return
+            }
+            case 'update-chat-message-telemetry': {
+                AmazonQChatMessageDuration.updateChatMessageTelemetry(msg)
+                return
+            }
+            case 'stop-chat-message-telemetry': {
+                AmazonQChatMessageDuration.stopChatMessageTelemetry(msg)
+                return
+            }
         }
 
         if (msg.type === 'error') {

packages/core/src/amazonq/webview/ui/apps/cwChatConnector.ts

@@ -8,11 +8,9 @@ import { ExtensionMessage } from '../commands'
 import { CodeReference } from './amazonqCommonsConnector'
 import { TabOpenType, TabsStorage } from '../storages/tabsStorage'
 import { FollowUpGenerator } from '../followUps/generator'
-import { TracedChatItem } from '../connector'
 
 interface ChatPayload {
     chatMessage: string
-    traceId?: string
     chatCommand?: string
 }
 
@@ -74,6 +72,17 @@ export class Connector {
     }
 
     followUpClicked = (tabID: string, messageId: string, followUp: ChatItemAction): void => {
+        /**
+         * We've pressed on a followup button and should start watching that round trip telemetry
+         */
+        this.sendMessageToExtension({
+            command: 'start-chat-message-telemetry',
+            trigger: 'followUpClicked',
+            tabID,
+            traceId: messageId,
+            tabType: 'cwc',
+            startTime: Date.now(),
+        })
         this.sendMessageToExtension({
             command: 'follow-up-was-clicked',
             followUp,
@@ -183,17 +192,29 @@ export class Connector {
         })
     }
 
-    requestGenerativeAIAnswer = (tabID: string, payload: ChatPayload): Promise<any> =>
-        new Promise((resolve, reject) => {
+    requestGenerativeAIAnswer = (tabID: string, messageId: string, payload: ChatPayload): Promise<any> => {
+        /**
+         * When a user presses "enter" send an event that indicates
+         * we should start tracking the round trip time for this message
+         **/
+        this.sendMessageToExtension({
+            command: 'start-chat-message-telemetry',
+            trigger: 'onChatPrompt',
+            tabID,
+            traceId: messageId,
+            tabType: 'cwc',
+            startTime: Date.now(),
+        })
+        return new Promise((resolve, reject) => {
             this.sendMessageToExtension({
                 tabID: tabID,
                 command: 'chat-prompt',
                 chatMessage: payload.chatMessage,
                 chatCommand: payload.chatCommand,
-                traceId: payload.traceId,
                 tabType: 'cwc',
             })
         })
+    }
 
     clearChat = (tabID: string): void => {
         this.sendMessageToExtension({
@@ -261,14 +282,13 @@ export class Connector {
                       }
                     : undefined
 
-            const answer: TracedChatItem = {
+            const answer: ChatItem = {
                 type: messageData.messageType,
                 messageId: messageData.messageID ?? messageData.triggerID,
                 body: messageData.message,
                 followUp: followUps,
                 canBeVoted: true,
                 codeReference: messageData.codeReference,
-                traceId: messageData.traceId,
             }
 
             // If it is not there we will not set it
@@ -295,7 +315,7 @@ export class Connector {
             return
         }
         if (messageData.messageType === ChatItemType.ANSWER) {
-            const answer: TracedChatItem = {
+            const answer: ChatItem = {
                 type: messageData.messageType,
                 body: undefined,
                 relatedContent: undefined,
@@ -308,7 +328,6 @@ export class Connector {
                               options: messageData.followUps,
                           }
                         : undefined,
-                traceId: messageData.traceId,
             }
             this.onChatAnswerReceived(messageData.tabID, answer)
 

packages/core/src/amazonq/webview/ui/commands.ts

@@ -31,5 +31,7 @@ type MessageCommand =
     | 'file-click'
     | 'form-action-click'
     | 'open-settings'
+    | 'start-chat-message-telemetry'
+    | 'stop-chat-message-telemetry'
 
 export type ExtensionMessage = Record<string, any> & { command: MessageCommand }

packages/core/src/amazonq/webview/ui/connector.ts

@@ -26,23 +26,18 @@ export interface CodeReference {
 
 export interface ChatPayload {
     chatMessage: string
-    traceId?: string // TODO: instrumented for cwc, not for gumby/featuredev. Remove the ? once we support all features
     chatCommand?: string
 }
 
-export interface TracedChatItem extends ChatItem {
-    traceId?: string
-}
-
 export interface ConnectorProps {
     sendMessageToExtension: (message: ExtensionMessage) => void
     onMessageReceived?: (tabID: string, messageData: any, needToShowAPIDocsTab: boolean) => void
     onChatAnswerUpdated?: (tabID: string, message: ChatItem) => void
-    onChatAnswerReceived?: (tabID: string, message: TracedChatItem) => void
+    onChatAnswerReceived?: (tabID: string, message: ChatItem) => void
     onWelcomeFollowUpClicked: (tabID: string, welcomeFollowUpType: WelcomeFollowupType) => void
     onAsyncEventProgress: (tabID: string, inProgress: boolean, message: string | undefined) => void
     onQuickHandlerCommand: (tabID: string, command?: string, eventId?: string) => void
-    onCWCContextCommandMessage: (message: TracedChatItem, command?: string) => string | undefined
+    onCWCContextCommandMessage: (message: ChatItem, command?: string) => string | undefined
     onOpenSettingsMessage: (tabID: string) => void
     onError: (tabID: string, message: string, title: string) => void
     onWarning: (tabID: string, message: string, title: string) => void
@@ -120,18 +115,18 @@ export class Connector {
         }
     }
 
-    requestGenerativeAIAnswer = (tabID: string, payload: ChatPayload): Promise<any> =>
+    requestGenerativeAIAnswer = (tabID: string, messageId: string, payload: ChatPayload): Promise<any> =>
         new Promise((resolve, reject) => {
             if (this.isUIReady) {
                 switch (this.tabsStorage.getTab(tabID)?.type) {
                     case 'featuredev':
                         return this.featureDevChatConnector.requestGenerativeAIAnswer(tabID, payload)
                     default:
-                        return this.cwChatConnector.requestGenerativeAIAnswer(tabID, payload)
+                        return this.cwChatConnector.requestGenerativeAIAnswer(tabID, messageId, payload)
                 }
             } else {
                 return setTimeout(() => {
-                    return this.requestGenerativeAIAnswer(tabID, payload)
+                    return this.requestGenerativeAIAnswer(tabID, messageId, payload)
                 }, 2000)
             }
         })