telemetry(amazonq): webview didLoadModule

Author: Will-ShaoHua

packages/core/src/login/webview/vue/amazonq/backend_amazonq.ts

@@ -23,6 +23,8 @@ import { ToolkitError } from '../../../../shared/errors'
 import { withTelemetryContext } from '../../../../shared/telemetry/util'
 import { builderIdStartUrl } from '../../../../auth/sso/constants'
 import { RegionProfile } from '../../../../codewhisperer/models/model'
+import { randomUUID } from '../../../../shared/crypto'
+import globals from '../../../../shared/extensionGlobals'
 import { telemetry } from '../../../../shared/telemetry/telemetry'
 import { ProfileSwitchIntent } from '../../../../codewhisperer/region/regionProfileManager'
 
@@ -162,6 +164,13 @@ export class AmazonQLoginWebview extends CommonAuthWebview {
             return
         } else if (featureAuthStates.amazonQ === 'pendingProfileSelection') {
             this.authState = 'PENDING_PROFILE_SELECTION'
+            // possible that user starts with "profile selection" state therefore the timeout for auth flow should be disposed otherwise will emit failure
+            this.loadMetadata?.loadTimeout?.dispose()
+            this.loadMetadata = {
+                traceId: randomUUID(),
+                loadTimeout: undefined,
+                start: globals.clock.Date.now(),
+            }
             return
         }
         this.authState = 'LOGIN'

packages/core/src/login/webview/vue/backend.ts

@@ -69,7 +69,7 @@ export abstract class CommonAuthWebview extends VueWebview {
      * @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) {
+    public setUiReady(state: 'login' | 'reauth' | 'selectProfile', errorMessage?: string) {
         if (errorMessage) {
             this.setLoadFailure(state, errorMessage)
         } else {

packages/core/src/login/webview/vue/regionProfileSelector.vue

@@ -79,7 +79,13 @@
                     <button id="signout" v-on:click="signout">Sign Out</button>
                 </template>
                 <template v-else>
-                    <button class="continue-button" v-on:click="onClickContinue()">Continue</button>
+                    <button
+                        class="continue-button"
+                        id="profile-selection-continue-button"
+                        v-on:click="onClickContinue()"
+                    >
+                        Continue
+                    </button>
                 </template>
             </div>
         </div>
@@ -162,6 +168,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 'profile-selection-continue-button'
+}
 </script>
 <style scoped>
 @import './base.css';

packages/core/src/login/webview/vue/root.vue

@@ -23,7 +23,7 @@ configure app to AMAZONQ if for Amazon Q login
 import { PropType, defineComponent } from 'vue'
 import Login, { getReadyElementId as getLoginReadyElementId } from './login.vue'
 import Reauthenticate, { getReadyElementId as getReauthReadyElementId } from './reauthenticate.vue'
-import RegionProfileSelector from './regionProfileSelector.vue'
+import RegionProfileSelector, { getReadyElementId as getSelectProfileReadyElementId } from './regionProfileSelector.vue'
 import { AuthFlowState, FeatureId } from './types'
 import { WebviewClientFactory } from '../../../webviews/client'
 import { CommonAuthWebview } from './backend'
@@ -69,6 +69,11 @@ export default defineComponent({
             })
         }
     },
+    async updated() {
+        if (didLoad) {
+            handleLoaded()
+        }
+    },
     methods: {
         async refreshAuthState() {
             await client.refreshAuthState()
@@ -78,6 +83,9 @@ export default defineComponent({
             if (this.authFlowState === 'LOGIN') {
                 ;(window as any).uiState = 'login'
                 ;(window as any).uiReadyElementId = getLoginReadyElementId()
+            } else if (this.authFlowState === 'PENDING_PROFILE_SELECTION') {
+                ;(window as any).uiState = 'selectProfile'
+                ;(window as any).uiReadyElementId = getSelectProfileReadyElementId()
             } else if (this.authFlowState && this.authFlowState !== undefined) {
                 ;(window as any).uiState = 'reauth'
                 ;(window as any).uiReadyElementId = getReauthReadyElementId()
@@ -90,12 +98,14 @@ export default defineComponent({
 
 // ---- 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
+const didPageSetReady: { auth: boolean; selectProfile: boolean } = { auth: false, selectProfile: false }
 
 // 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) {
+    const uiState = (window as any).uiState
+    const page: 'auth' | 'selectProfile' = uiState === 'login' || uiState === 'reauth' ? 'auth' : 'selectProfile'
+    if (didPageSetReady[page]) {
         return
     }
 
@@ -104,7 +114,9 @@ window.onerror = function (message) {
 document.addEventListener(
     'error',
     (e) => {
-        if (didSetReady) {
+        const uiState = (window as any).uiState
+        const page: 'auth' | 'selectProfile' = uiState === 'login' || uiState === 'reauth' ? 'auth' : 'selectProfile'
+        if (didPageSetReady[page]) {
             return
         }
 
@@ -119,7 +131,9 @@ window.addEventListener('load', () => {
 })
 const handleLoaded = () => {
     // in case some unexpected behavior triggers this flow again, skip since we already emitted for this instance
-    if (didSetReady) {
+    const uiState = (window as any).uiState
+    const page: 'auth' | 'selectProfile' = uiState === 'login' || uiState === 'reauth' ? 'auth' : 'selectProfile'
+    if (didPageSetReady[page]) {
         return
     }
 
@@ -131,9 +145,11 @@ const handleLoaded = () => {
         setUiReady((window as any).uiState)
     }
 }
-const setUiReady = (state: 'login' | 'reauth', errorMessage?: string) => {
+const setUiReady = (state: 'login' | 'reauth' | 'selectProfile', errorMessage?: string) => {
     client.setUiReady(state, errorMessage)
-    didSetReady = true
+
+    const page = state === 'selectProfile' ? 'selectProfile' : 'auth'
+    didPageSetReady[page] = true
 }
 // ---- END ----
 </script>