From 965969c3fad1d58346761f6fb96272854355b110 Mon Sep 17 00:00:00 2001
From: nkomonen-amazon <nkomonen@amazon.com>
Date: Fri, 14 Mar 2025 13:25:32 -0400
Subject: [PATCH 1/2] telemetry(webview): Emit toolkit_X module telemetry on
 auth webview

Problem:

With our telemetry, we do not know when the frontend webview UI has actually loaded

Solution:

Emit certain metrics during the webview loading process to get a better idea of if the
webview UI successfully completed its initial load.

- toolkit_willOpenModule, indicates intent to render a webview
- toolkit_didLoadModule, indicates the final result of loading the webview
  - On Success it it just a success result. We know a success when the frontend send a successful
    message to the backend. It knows this by ensuring there were no errors and that a certain HTML element
    can be found, then once the page finishes its initial load it will send a success message to the backend.
  - On Failure, what happens is a timer times out after 10 seconds.

Signed-off-by: nkomonen-amazon <nkomonen@amazon.com>
---
 .../webview/messages/messageDispatcher.ts     |  25 ++--
 packages/core/src/amazonq/webview/ui/main.ts  |   2 +-
 .../webview/vue/amazonq/backend_amazonq.ts    |   1 +
 .../core/src/login/webview/vue/backend.ts     |  23 ++--
 packages/core/src/login/webview/vue/login.vue |  15 ++-
 .../src/login/webview/vue/reauthenticate.vue  |  19 ++-
 packages/core/src/login/webview/vue/root.vue  |  72 ++++++++++-
 .../src/shared/telemetry/vscodeTelemetry.json |  10 --
 packages/core/src/webviews/main.ts            | 113 ++++++++++++++++++
 9 files changed, 243 insertions(+), 37 deletions(-)

diff --git a/packages/core/src/amazonq/webview/messages/messageDispatcher.ts b/packages/core/src/amazonq/webview/messages/messageDispatcher.ts
index dc3bc4c77c8..f17de54e416 100644
--- a/packages/core/src/amazonq/webview/messages/messageDispatcher.ts
+++ b/packages/core/src/amazonq/webview/messages/messageDispatcher.ts
@@ -16,6 +16,8 @@ import globals from '../../../shared/extensionGlobals'
 import { openUrl } from '../../../shared/utilities/vsCodeUtils'
 import { DefaultAmazonQAppInitContext } from '../../apps/initContext'
 
+const qChatModuleName = 'amazonqChat'
+
 export function dispatchWebViewMessagesToApps(
     webview: Webview,
     webViewToAppsMessagePublishers: Map<TabType, MessagePublisher<any>>
@@ -29,8 +31,8 @@ export function dispatchWebViewMessagesToApps(
                  * 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',
+                telemetry.toolkit_didLoadModule.emit({
+                    module: qChatModuleName,
                     duration: performance.measure(amazonqMark.uiReady, amazonqMark.open).duration,
                     result: 'Succeeded',
                 })
@@ -86,12 +88,19 @@ export function dispatchWebViewMessagesToApps(
         }
 
         if (msg.type === 'error') {
-            const event = msg.event === 'webview_load' ? telemetry.webview_load : telemetry.webview_error
-            event.emit({
-                webviewName: 'amazonqChat',
-                result: 'Failed',
-                reasonDesc: msg.errorMessage,
-            })
+            if (msg.event === 'toolkit_didLoadModule') {
+                telemetry.toolkit_didLoadModule.emit({
+                    module: qChatModuleName,
+                    result: 'Failed',
+                    reasonDesc: msg.errorMessage,
+                })
+            } else {
+                telemetry.webview_error.emit({
+                    webviewName: qChatModuleName,
+                    result: 'Failed',
+                    reasonDesc: msg.errorMessage,
+                })
+            }
             return
         }
 
diff --git a/packages/core/src/amazonq/webview/ui/main.ts b/packages/core/src/amazonq/webview/ui/main.ts
index d7285d81ba5..5ae03840f8f 100644
--- a/packages/core/src/amazonq/webview/ui/main.ts
+++ b/packages/core/src/amazonq/webview/ui/main.ts
@@ -63,7 +63,7 @@ export const createMynahUI = (
         const { error, message } = e
         ideApi.postMessage({
             type: 'error',
-            event: connector.isUIReady ? 'webview_error' : 'webview_load',
+            event: connector.isUIReady ? 'webview_error' : 'toolkit_didLoadModule',
             errorMessage: error ? error.toString() : message,
         })
     })
diff --git a/packages/core/src/login/webview/vue/amazonq/backend_amazonq.ts b/packages/core/src/login/webview/vue/amazonq/backend_amazonq.ts
index 83c511ad9f6..7132e91afec 100644
--- a/packages/core/src/login/webview/vue/amazonq/backend_amazonq.ts
+++ b/packages/core/src/login/webview/vue/amazonq/backend_amazonq.ts
@@ -27,6 +27,7 @@ const className = 'AmazonQLoginWebview'
 export class AmazonQLoginWebview extends CommonAuthWebview {
     public override id: string = 'aws.amazonq.AmazonCommonAuth'
     public static sourcePath: string = 'vue/src/login/webview/vue/amazonq/index.js'
+    public override supportsLoadTelemetry: boolean = true
 
     override onActiveConnectionModified = new vscode.EventEmitter<void>()
 
diff --git a/packages/core/src/login/webview/vue/backend.ts b/packages/core/src/login/webview/vue/backend.ts
index 9bc6c5ae339..43d86feedf9 100644
--- a/packages/core/src/login/webview/vue/backend.ts
+++ b/packages/core/src/login/webview/vue/backend.ts
@@ -62,18 +62,25 @@ export abstract class CommonAuthWebview extends VueWebview {
     }
 
     private didCall: { login: boolean; reauth: boolean } = { login: false, reauth: false }
-    public setUiReady(state: 'login' | 'reauth') {
-        // Prevent telemetry spam, since showing/hiding chat triggers this each time.
-        // So only emit once.
+    /**
+     * Called when the UI load process is completed, regardless of success or failure
+     *
+     * @param errorMessage IF an error is caught on the frontend, this is the message. It will result in a failure metric.
+     *                     Otherwise we assume success.
+     */
+    public setUiReady(state: 'login' | 'reauth', errorMessage?: string) {
+        // Only emit once to prevent telemetry spam, since showing/hiding chat triggers this each time.
+        // TODO: Research how to not trigger this on every show/hide
         if (this.didCall[state]) {
             return
         }
 
-        telemetry.webview_load.emit({
-            passive: true,
-            webviewName: state,
-            result: 'Succeeded',
-        })
+        if (errorMessage) {
+            this.setLoadFailure(state, errorMessage)
+        } else {
+            this.setDidLoad(state)
+        }
+
         this.didCall[state] = true
     }
 
diff --git a/packages/core/src/login/webview/vue/login.vue b/packages/core/src/login/webview/vue/login.vue
index d21bf911802..831a1dfcb22 100644
--- a/packages/core/src/login/webview/vue/login.vue
+++ b/packages/core/src/login/webview/vue/login.vue
@@ -145,6 +145,7 @@
                 ></SelectableItem>
                 <button
                     class="continue-button"
+                    id="connection-selection-continue-button"
                     :disabled="selectedLoginOption === 0"
                     v-on:click="handleContinueClick()"
                 >
@@ -368,8 +369,6 @@ export default defineComponent({
         // Pre-select the first available login option
         await this.preselectLoginOption()
         await this.handleUrlInput() // validate the default startUrl
-
-        await client.setUiReady('login')
     },
     methods: {
         toggleItemSelection(itemId: number) {
@@ -590,6 +589,18 @@ export default defineComponent({
         },
     },
 })
+
+/**
+ * The ID of the element we will use to determine that the UI has completed its initial load.
+ *
+ * This makes assumptions that we will be in a certain state of the UI (eg showing a form vs. a loading bar).
+ * So if the UI flow changes, this may need to be updated.
+ */
+export function getReadyElementId() {
+    // On every initial load, we ASSUME that the user will always be in the connection selection state,
+    // which is why we specifically look for this button.
+    return 'connection-selection-continue-button'
+}
 </script>
 
 <style>
diff --git a/packages/core/src/login/webview/vue/reauthenticate.vue b/packages/core/src/login/webview/vue/reauthenticate.vue
index 3fc810f36b7..dbd0c16f6cb 100644
--- a/packages/core/src/login/webview/vue/reauthenticate.vue
+++ b/packages/core/src/login/webview/vue/reauthenticate.vue
@@ -61,7 +61,7 @@ and the final results are retrieved by the frontend. For this Component to updat
                 </div>
 
                 <div>
-                    <button id="reauthenticate" v-on:click="reauthenticate">Re-authenticate</button>
+                    <button id="reauthenticate-button" v-on:click="reauthenticate">Re-authenticate</button>
                     <div v-if="errorMessage" id="error-message" style="color: red">{{ errorMessage }}</div>
                 </div>
 
@@ -127,9 +127,6 @@ export default defineComponent({
 
         this.doShow = true
     },
-    async mounted() {
-        await client.setUiReady('reauth')
-    },
     methods: {
         async reauthenticate() {
             client.emitUiClick('auth_reauthenticate')
@@ -148,6 +145,16 @@ export default defineComponent({
         },
     },
 })
+
+/**
+ * The ID of the element we will use to determine that the UI has completed its initial load.
+ *
+ * This makes assumptions that we will be in a certain state of the UI (eg showing a form vs. a loading bar).
+ * So if the UI flow changes, this may need to be updated.
+ */
+export function getReadyElementId() {
+    return 'reauthenticate-button'
+}
 </script>
 <style>
 @import './base.css';
@@ -199,7 +206,7 @@ export default defineComponent({
     flex-direction: column;
 }
 
-button#reauthenticate {
+button#reauthenticate-button {
     cursor: pointer;
     background-color: var(--vscode-button-background);
     color: white;
@@ -231,7 +238,7 @@ button#cancel {
     cursor: pointer;
 }
 
-body.vscode-high-contrast:not(body.vscode-high-contrast-light) button#reauthenticate {
+body.vscode-high-contrast:not(body.vscode-high-contrast-light) button#reauthenticate-button {
     background-color: white;
     color: black;
 }
diff --git a/packages/core/src/login/webview/vue/root.vue b/packages/core/src/login/webview/vue/root.vue
index efc34881c9b..2dfe55cf5ca 100644
--- a/packages/core/src/login/webview/vue/root.vue
+++ b/packages/core/src/login/webview/vue/root.vue
@@ -16,8 +16,8 @@ configure app to AMAZONQ if for Amazon Q login
 </template>
 <script lang="ts">
 import { PropType, defineComponent } from 'vue'
-import Login from './login.vue'
-import Reauthenticate from './reauthenticate.vue'
+import Login, { getReadyElementId as getLoginReadyElementId } from './login.vue'
+import Reauthenticate, { getReadyElementId as getReauthReadyElementId } from './reauthenticate.vue'
 import { AuthFlowState, FeatureId } from './types'
 import { WebviewClientFactory } from '../../../webviews/client'
 import { CommonAuthWebview } from './backend'
@@ -51,15 +51,83 @@ export default defineComponent({
         })
 
         await this.refreshAuthState()
+
+        // We were recieving the 'load' event before refreshAuthState() resolved (I'm assuming behavior w/ Vue + browser loading not blocking),
+        // so post refreshAuhState() if we detect we already loaded, then execute immediately since the event already happened.
+        if (didLoad) {
+            handleLoaded()
+        } else {
+            window.addEventListener('load', () => {
+                handleLoaded()
+            })
+        }
     },
     methods: {
         async refreshAuthState() {
             await client.refreshAuthState()
             this.authFlowState = await client.getAuthState()
+
+            // Used for telemetry purposes
+            if (this.authFlowState === 'LOGIN') {
+                ;(window as any).uiState = 'login'
+                ;(window as any).uiReadyElementId = getLoginReadyElementId()
+            } else if (this.authFlowState && this.authFlowState !== undefined) {
+                ;(window as any).uiState = 'reauth'
+                ;(window as any).uiReadyElementId = getReauthReadyElementId()
+            }
+
             this.refreshKey += 1
         },
     },
 })
+
+// ---- START ---- The following handles the process of indicating the UI has loaded successfully.
+// TODO: Move this in to a reusable class for other webviews, it feels a bit messy here
+let didSetReady = false
+
+// Setup error handlers to report 
+window.onerror = function (message) {
+    if (didSetReady) {
+        return
+    }
+
+    setUiReady((window as any).uiState, message.toString())
+}
+document.addEventListener(
+    'error',
+    (e) => {
+        if (didSetReady) {
+            return
+        }
+
+        setUiReady((window as any).uiState, e.message)
+    },
+    true
+)
+
+let didLoad = false
+window.addEventListener('load', () => {
+    didLoad = true
+})
+const handleLoaded = () => {
+    // in case some unexpected behavior triggers this flow again, skip since we already emitted for this instance
+    if (didSetReady) {
+        return
+    }
+
+    const foundElement = !!document.getElementById((window as any).uiReadyElementId)
+    if (!foundElement) {
+        setUiReady((window as any).uiState, `Could not find element: ${(window as any).uiReadyElementId}`)
+    } else {
+        // Successful load!    
+        setUiReady((window as any).uiState)
+    }
+}
+const setUiReady = (state: 'login' | 'reauth', errorMessage?: string) => {
+    client.setUiReady(state, errorMessage)
+    didSetReady = true
+}
+// ---- END ----
 </script>
 <style>
 body {
diff --git a/packages/core/src/shared/telemetry/vscodeTelemetry.json b/packages/core/src/shared/telemetry/vscodeTelemetry.json
index bba9a67a6ec..6dd6d5a3120 100644
--- a/packages/core/src/shared/telemetry/vscodeTelemetry.json
+++ b/packages/core/src/shared/telemetry/vscodeTelemetry.json
@@ -985,16 +985,6 @@
             ],
             "passive": true
         },
-        {
-            "name": "webview_load",
-            "description": "Represents a webview load event",
-            "metadata": [
-                {
-                    "type": "webviewName",
-                    "required": true
-                }
-            ]
-        },
         {
             "name": "webview_error",
             "description": "Represents an error that occurs in a webview",
diff --git a/packages/core/src/webviews/main.ts b/packages/core/src/webviews/main.ts
index 1fa18d3c460..2b78eb42df8 100644
--- a/packages/core/src/webviews/main.ts
+++ b/packages/core/src/webviews/main.ts
@@ -7,6 +7,10 @@ import * as vscode from 'vscode'
 import { Protocol, registerWebviewServer } from './server'
 import { getIdeProperties } from '../shared/extensionUtilities'
 import { getFunctions } from '../shared/utilities/classUtils'
+import { telemetry } from '../shared/telemetry/telemetry'
+import { randomUUID } from '../shared/crypto'
+import { Timeout } from '../shared/utilities/timeoutUtils'
+import globals from '../shared/extensionGlobals'
 
 interface WebviewParams {
     /**
@@ -151,6 +155,35 @@ export abstract class VueWebview {
      */
     public abstract readonly id: string
 
+    /**
+     * Implementations must override this field to `true` if the webview can support
+     * {@link telemetry.toolkit_didLoadModule}.
+     * A webview that supports this will call {@link setDidLoad}
+     * to confirm the UI has successfully loaded.
+     */
+    public supportsLoadTelemetry: boolean = false
+    /**
+     * The metadata used to connect {@link telemetry.toolkit_willOpenModule} and {@link telemetry.toolkit_didLoadModule}.
+     * Always clear this field after load is completed, as it should not be shared with future open/loads.
+     */
+    public loadMetadata:
+        | {
+              /**
+               * A unique identifier used to connect the opening/loaded metrics of a webview
+               */
+              traceId: string
+              /**
+               * When the webview is doing its initial render/load, this times out if it takes too long
+               * to get a "success" message (we assume something went wrong)
+               */
+              loadTimeout: Timeout | undefined
+              /**
+               * The time in milliseconds, when the module was triggered to start loading
+               */
+              start: number
+          }
+        | undefined = undefined
+
     /**
      * The relative location, from the repository root, to the frontend entrypoint.
      */
@@ -189,6 +222,7 @@ export abstract class VueWebview {
 
     protected dispose(): void {
         this.disposed = true
+        this.loadMetadata?.loadTimeout?.dispose()
         this.onDidDisposeEmitter.fire()
     }
 
@@ -225,12 +259,51 @@ export abstract class VueWebview {
             }
 
             public async setup(webview: vscode.Webview) {
+                this.setupTelemetry()
+
                 const server = registerWebviewServer(webview, this.instance.protocol, this.instance.id)
                 this.instance.onDidDispose(() => {
                     server.dispose()
                 })
             }
 
+            /**
+             * Setup telemetry events that report on the initial loading of a webview
+             */
+            private setupTelemetry() {
+                const traceId = randomUUID()
+                // Notify intent to open a module, this does not mean it successfully opened
+                telemetry.toolkit_willOpenModule.emit({
+                    module: this.instance.id,
+                    result: 'Succeeded',
+                    traceId: traceId,
+                })
+
+                // setup for when the webview is (possibly) loaded
+                if (this.instance.supportsLoadTelemetry) {
+                    // Arbitrary assumption that UI will take no longer than 10 seconds to load
+                    const loadTimeout = new Timeout(10_000)
+                    const start = globals.clock.Date.now()
+                    this.instance.loadMetadata = {
+                        traceId,
+                        loadTimeout,
+                        start,
+                    }
+
+                    // webview frontend did not successfuly respond quick enough, so we assume loading failed
+                    loadTimeout.token.onCancellationRequested(() => {
+                        telemetry.toolkit_didLoadModule.emit({
+                            module: this.instance.id,
+                            result: 'Failed',
+                            traceId: traceId,
+                            reason: 'LoadTimedOut',
+                            reasonDesc: 'Likely due to an error in the frontend',
+                        })
+                        this.instance.loadMetadata = undefined
+                    })
+                }
+            }
+
             public async show(params: Omit<WebviewPanelParams, 'id' | 'webviewJs'>): Promise<vscode.WebviewPanel> {
                 if (this.panel) {
                     this.panel.reveal(params.viewColumn, false)
@@ -317,6 +390,46 @@ export abstract class VueWebview {
             }
         }
     }
+
+    /**
+     * Call this after the webview has successfully loaded
+     */
+    protected setDidLoad(module: string) {
+        this.loadMetadata?.loadTimeout?.dispose()
+
+        // Represents time from intent to open, to confirmation of a successful load
+        const duration = globals.clock.Date.now() - this.loadMetadata!.start
+
+        telemetry.toolkit_didLoadModule.emit({
+            passive: true,
+            module,
+            result: 'Succeeded',
+            traceId: this.loadMetadata?.traceId,
+            duration,
+        })
+
+        this.loadMetadata = undefined
+    }
+
+    /**
+     * Call this if we catch an error in the frontend, and are able to forward it to the backed
+     *
+     * @param message Error message from the frontend
+     */
+    public setLoadFailure(module: string, message: string) {
+        this.loadMetadata?.loadTimeout?.dispose()
+
+        telemetry.toolkit_didLoadModule.emit({
+            passive: true,
+            module,
+            result: 'Failed',
+            traceId: this.loadMetadata?.traceId,
+            reason: 'CaughtFrontendError',
+            reasonDesc: message,
+        })
+
+        this.loadMetadata = undefined
+    }
 }
 
 type FilteredKeys<T> = { [P in keyof T]: unknown extends T[P] ? never : P }[keyof T]

From da0b2f96efa406f793e668034f1177dd56f0a77c Mon Sep 17 00:00:00 2001
From: nkomonen-amazon <nkomonen@amazon.com>
Date: Fri, 14 Mar 2025 17:36:18 -0400
Subject: [PATCH 2/2] fix race condition

Setting .html starts the loading of the UI, but setup() sets up the message listeners in the backend
for messages from the UI.

We had setup() come after, and it worked, but if I added a small sleep() before setup() was run
it would result in a failure due to messages not being handled due to handlers not being setup in time.

As a solution this just moves the handler setup before we set the new UI.

Signed-off-by: nkomonen-amazon <nkomonen@amazon.com>
---
 packages/core/src/login/webview/commonAuthViewProvider.ts | 3 ++-
 packages/core/src/login/webview/vue/root.vue              | 5 +++--
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/packages/core/src/login/webview/commonAuthViewProvider.ts b/packages/core/src/login/webview/commonAuthViewProvider.ts
index 12caf3c2ad5..d9d6f3b4e30 100644
--- a/packages/core/src/login/webview/commonAuthViewProvider.ts
+++ b/packages/core/src/login/webview/commonAuthViewProvider.ts
@@ -135,9 +135,10 @@ export class CommonAuthViewProvider implements WebviewViewProvider {
             enableCommandUris: true,
             localResourceRoots: [dist, resources],
         }
-        webviewView.webview.html = this._getHtmlForWebview(this.extensionContext.extensionUri, webviewView.webview)
         // register the webview server
         await this.webView?.setup(webviewView.webview)
+
+        webviewView.webview.html = this._getHtmlForWebview(this.extensionContext.extensionUri, webviewView.webview)
     }
 
     private _getHtmlForWebview(extensionURI: Uri, webview: vscode.Webview) {
diff --git a/packages/core/src/login/webview/vue/root.vue b/packages/core/src/login/webview/vue/root.vue
index 2dfe55cf5ca..d50634048c9 100644
--- a/packages/core/src/login/webview/vue/root.vue
+++ b/packages/core/src/login/webview/vue/root.vue
@@ -85,7 +85,8 @@ export default defineComponent({
 // TODO: Move this in to a reusable class for other webviews, it feels a bit messy here
 let didSetReady = false
 
-// Setup error handlers to report 
+// Setup error handlers to report. This may not actually be able to catch certain errors that we'd expect,
+// so this may have to be revisited.
 window.onerror = function (message) {
     if (didSetReady) {
         return
@@ -119,7 +120,7 @@ const handleLoaded = () => {
     if (!foundElement) {
         setUiReady((window as any).uiState, `Could not find element: ${(window as any).uiReadyElementId}`)
     } else {
-        // Successful load!    
+        // Successful load!
         setUiReady((window as any).uiState)
     }
 }