telemetry(amazonq): measure e2e latency for amazon q requests (#5566)

## Problem:
- We can't measure e2e how long a chat request takes

## Solution:
- Plumb the trace id through to the webview -> vscode -> webview. This
allows us to know exactly what is happening to a message and how long
everything took to get there
- chat e2e latency can be linked with the other telemetry events via the
traceId

Author: jpinkney-aws

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

@@ -0,0 +1,93 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { globals } from '../../shared'
+import { telemetry } from '../../shared/telemetry'
+import { Event, uiEventRecorder } from '../util/eventRecorder'
+
+export class AmazonQChatMessageDuration {
+    /**
+     * Record the initial requests in the chat message flow
+     */
+    static startChatMessageTelemetry(msg: { traceId: string; startTime: number; trigger?: string }) {
+        const { traceId, startTime, trigger } = msg
+
+        uiEventRecorder.set(traceId, {
+            events: {
+                chatMessageSent: startTime,
+            },
+        })
+        uiEventRecorder.set(traceId, {
+            events: {
+                editorReceivedMessage: globals.clock.Date.now(),
+            },
+        })
+        if (trigger) {
+            uiEventRecorder.set(traceId, {
+                trigger,
+            })
+        }
+    }
+
+    /**
+     * Stop listening to all incoming events and emit what we've found
+     */
+    static stopChatMessageTelemetry(msg: { traceId: string }) {
+        const { traceId } = msg
+
+        // We can't figure out what trace this event was associated with
+        if (!traceId) {
+            return
+        }
+
+        uiEventRecorder.set(traceId, {
+            events: {
+                messageDisplayed: globals.clock.Date.now(),
+            },
+        })
+
+        const metrics = uiEventRecorder.get(traceId)
+
+        // 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
+            })
+
+        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
+
+        /**
+         * 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 timeDifference = currentEvent.duration - previousEvent.duration
+
+            timings.set(currentEvent.event as Event, timeDifference)
+        }
+
+        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,
+        })
+
+        uiEventRecorder.delete(traceId)
+    }
+}

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

@@ -0,0 +1,29 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { RecordMap } from '../../shared/utilities/map'
+
+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
+
+/**
+ * For a given traceID, map an event to a time
+ *
+ * This is used to correlated disjoint events that are happening in different
+ * parts of Q Chat.
+ *
+ * It allows us to tracks time intervals between key events:
+ *  - when VSCode received the message
+ *  - 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
+ */
+export const uiEventRecorder = new RecordMap<{
+    trigger: string
+    events: Partial<Record<Event, number>>
+}>()

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

@@ -10,6 +10,7 @@ import { TabType } from '../ui/storages/tabsStorage'
 import { getLogger } from '../../../shared/logger'
 import { amazonqMark } from '../../../shared/performance/marks'
 import { telemetry } from '../../../shared/telemetry'
+import { AmazonQChatMessageDuration } from '../../messages/chatMessageDuration'
 
 export function dispatchWebViewMessagesToApps(
     webview: Webview,
@@ -33,6 +34,14 @@ export function dispatchWebViewMessagesToApps(
             return
         }
 
+        if (msg.type === 'startChatMessageTelemetry') {
+            AmazonQChatMessageDuration.startChatMessageTelemetry(msg)
+            return
+        } else if (msg.type === 'stopChatMessageTelemetry') {
+            AmazonQChatMessageDuration.stopChatMessageTelemetry(msg)
+            return
+        }
+
         if (msg.type === 'error') {
             const event = msg.event === 'webview_load' ? telemetry.webview_load : telemetry.webview_error
             event.emit({

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

@@ -8,9 +8,11 @@ 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
 }
 
@@ -188,6 +190,7 @@ export class Connector {
                 command: 'chat-prompt',
                 chatMessage: payload.chatMessage,
                 chatCommand: payload.chatCommand,
+                traceId: payload.traceId,
                 tabType: 'cwc',
             })
         })
@@ -258,13 +261,14 @@ export class Connector {
                       }
                     : undefined
 
-            const answer: ChatItem = {
+            const answer: TracedChatItem = {
                 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
@@ -291,7 +295,7 @@ export class Connector {
             return
         }
         if (messageData.messageType === ChatItemType.ANSWER) {
-            const answer: ChatItem = {
+            const answer: TracedChatItem = {
                 type: messageData.messageType,
                 body: undefined,
                 relatedContent: undefined,
@@ -304,6 +308,7 @@ export class Connector {
                               options: messageData.followUps,
                           }
                         : undefined,
+                traceId: messageData.traceId,
             }
             this.onChatAnswerReceived(messageData.tabID, answer)
 

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

@@ -26,18 +26,23 @@ 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: ChatItem) => void
+    onChatAnswerReceived?: (tabID: string, message: TracedChatItem) => 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: ChatItem, command?: string) => string | undefined
+    onCWCContextCommandMessage: (message: TracedChatItem, command?: string) => string | undefined
     onOpenSettingsMessage: (tabID: string) => void
     onError: (tabID: string, message: string, title: string) => void
     onWarning: (tabID: string, message: string, title: string) => void

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

@@ -2,7 +2,7 @@
  * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
-import { Connector } from './connector'
+import { Connector, TracedChatItem } from './connector'
 import { ChatItem, ChatItemType, MynahIcons, MynahUI, MynahUIDataModel, NotificationType } from '@aws/mynah-ui'
 import { ChatPrompt } from '@aws/mynah-ui/dist/static'
 import { TabsStorage, TabType } from './storages/tabsStorage'
@@ -189,7 +189,7 @@ export const createMynahUI = (ideApi: any, amazonQEnabled: boolean) => {
                 } as ChatItem)
             }
         },
-        onChatAnswerReceived: (tabID: string, item: ChatItem) => {
+        onChatAnswerReceived: (tabID: string, item: TracedChatItem) => {
             if (item.type === ChatItemType.ANSWER_PART || item.type === ChatItemType.CODE_RESULT) {
                 mynahUI.updateLastChatAnswer(tabID, {
                     ...(item.messageId !== undefined ? { messageId: item.messageId } : {}),
@@ -234,6 +234,20 @@ export const createMynahUI = (ideApi: any, amazonQEnabled: boolean) => {
                     promptInputDisabledState: tabsStorage.isTabDead(tabID),
                 })
                 tabsStorage.updateTabStatus(tabID, 'free')
+
+                if (item.traceId) {
+                    /**
+                     * We've received an answer for a traceId and this message has
+                     * completed its round trip. Send that information back to
+                     * VSCode so we can emit a round trip event
+                     **/
+                    ideApi.postMessage({
+                        type: 'stopChatMessageTelemetry',
+                        tabID,
+                        traceId: item.traceId,
+                        tabType: tabsStorage.getTab(tabID)?.type,
+                    })
+                }
             }
         },
         onMessageReceived: (tabID: string, messageData: MynahUIDataModel) => {
@@ -362,6 +376,7 @@ export const createMynahUI = (ideApi: any, amazonQEnabled: boolean) => {
         },
         onTabRemove: connector.onTabRemove,
         onTabChange: connector.onTabChange,
+        // TODO: update mynah-ui this type doesn't seem correct https://github.com/aws/mynah-ui/blob/3777a39eb534a91fd6b99d6cf421ce78ee5c7526/src/main.ts#L372
         onChatPrompt: (tabID: string, prompt: ChatPrompt, eventId: string | undefined) => {
             if ((prompt.prompt ?? '') === '' && (prompt.command ?? '') === '') {
                 return
@@ -374,6 +389,7 @@ export const createMynahUI = (ideApi: any, amazonQEnabled: boolean) => {
             } else if (tabsStorage.getTab(tabID)?.type === 'gumby') {
                 connector.requestAnswer(tabID, {
                     chatMessage: prompt.prompt ?? '',
+                    traceId: eventId as string,
                 })
                 return
             }
@@ -383,7 +399,20 @@ export const createMynahUI = (ideApi: any, amazonQEnabled: boolean) => {
                 return
             }
 
-            textMessageHandler.handle(prompt, tabID)
+            /**
+             * When a user presses "enter" send an event that indicates
+             * we should start tracking the round trip time for this message
+             **/
+            ideApi.postMessage({
+                type: 'startChatMessageTelemetry',
+                trigger: 'onChatPrompt',
+                tabID,
+                traceId: eventId,
+                tabType: tabsStorage.getTab(tabID)?.type,
+                startTime: Date.now(),
+            })
+
+            textMessageHandler.handle(prompt, tabID, eventId as string)
         },
         onVote: connector.onChatItemVoted,
         onInBodyButtonClicked: (tabId, messageId, action, eventId) => {

packages/core/src/amazonq/webview/ui/messages/handler.ts

@@ -24,7 +24,7 @@ export class TextMessageHandler {
         this.tabsStorage = props.tabsStorage
     }
 
-    public handle(chatPrompt: ChatPrompt, tabID: string) {
+    public handle(chatPrompt: ChatPrompt, tabID: string, eventID: string) {
         this.tabsStorage.updateTabTypeFromUnknown(tabID, 'cwc')
         this.tabsStorage.resetTabTimer(tabID)
         this.connector.onUpdateTabType(tabID)
@@ -44,6 +44,7 @@ export class TextMessageHandler {
             .requestGenerativeAIAnswer(tabID, {
                 chatMessage: chatPrompt.prompt ?? '',
                 chatCommand: chatPrompt.command,
+                traceId: eventID,
             })
             .then(() => {})
     }

packages/core/src/codewhispererChat/controllers/chat/controller.ts

@@ -49,6 +49,9 @@ import { CodeWhispererSettings } from '../../../codewhisperer/util/codewhisperer
 import { getSelectedCustomization } from '../../../codewhisperer/util/customizationUtil'
 import { FeatureConfigProvider } from '../../../shared/featureConfig'
 import { getHttpStatusCode, AwsClientResponseError } from '../../../shared/errors'
+import { uiEventRecorder } from '../../../amazonq/util/eventRecorder'
+import { globals } from '../../../shared'
+import { telemetry } from '../../../shared/telemetry'
 
 export interface ChatControllerMessagePublishers {
     readonly processPromptChatMessage: MessagePublisher<PromptMessage>
@@ -122,7 +125,21 @@ export class ChatController {
         })
 
         this.chatControllerMessageListeners.processPromptChatMessage.onMessage((data) => {
-            return this.processPromptChatMessage(data)
+            if (data.traceId) {
+                uiEventRecorder.set(data.traceId, {
+                    events: {
+                        featureReceivedMessage: globals.clock.Date.now(),
+                    },
+                })
+            }
+            /**
+             * traceId is only instrumented for chat-prompt but not for things
+             * like follow-up-was-clicked. In those cases we fallback to a different
+             * uuid
+             **/
+            return telemetry.withTraceId(() => {
+                return this.processPromptChatMessage(data)
+            }, data.traceId ?? randomUUID())
         })
 
         this.chatControllerMessageListeners.processTabCreatedMessage.onMessage((data) => {
@@ -486,6 +503,7 @@ export class ChatController {
                         codeQuery: context?.focusAreaContext?.names,
                         userIntent: this.userIntentRecognizer.getFromPromptChatMessage(message),
                         customization: getSelectedCustomization(),
+                        traceId: message.traceId,
                     },
                     triggerID
                 )

packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts

@@ -293,6 +293,7 @@ export class Messenger {
                             relatedSuggestions: undefined,
                             triggerID,
                             messageID,
+                            traceId: triggerPayload.traceId,
                         },
                         tabID
                     )

packages/core/src/codewhispererChat/controllers/chat/model.ts

@@ -72,6 +72,7 @@ export type ChatPromptCommandType =
 export interface PromptMessage {
     message: string | undefined
     messageId: string
+    traceId?: string
     command: ChatPromptCommandType | undefined
     userIntent: UserIntent | undefined
     tabID: string
@@ -145,6 +146,7 @@ export interface TriggerPayload {
     readonly customization: Customization
     relevantTextDocuments?: RelevantTextDocument[]
     useRelevantDocuments?: boolean
+    traceId?: string
 }
 
 export interface InsertedCode {