devEnv: Fix Dev Env timeout + logs (#4980)

nkomonen-amazon · web-flow · commit 4512e2052ec2 · 2024-05-16T13:18:57.000-04:00
* fix: add more logs to Dev Env heartbeat

We want more logs to help determine if the dev env
heartbeat is working as expected

Signed-off-by: Nikolas Komonen &lt;nkomonen@amazon.com&gt;

* refactor: fix circular dependency

The command that shows the login vue was having circular
dependency issues.

This moves that logic in to its own file to fix this.

Signed-off-by: Nikolas Komonen &lt;nkomonen@amazon.com&gt;

* devEnv: try start activity heartbeat on connection change

Problem:

We only tried to start the activity heartbeat on extension startup. So if
that failed due to something like an invalid connection, the heartbeat mechanism
would not start

Solution:

Try and start the heartbeat mechanism when the CC connection changes as well.
But if it is already running then we just skip it.

Signed-off-by: Nikolas Komonen &lt;nkomonen@amazon.com&gt;

---------

Signed-off-by: Nikolas Komonen &lt;nkomonen@amazon.com&gt;
diff --git a/packages/core/src/auth/activation.ts b/packages/core/src/auth/activation.ts
@@ -11,8 +11,8 @@ import { getLogger } from '../shared/logger'
 import { ExtensionUse, initAuthCommands } from './utils'
 import { isCloud9 } from '../shared/extensionUtilities'
 import { isInDevEnv } from '../shared/vscode/env'
-import { registerCommands, getShowManageConnections } from './ui/vue/show'
 import { isWeb } from '../shared/extensionGlobals'
+import { getShowManageConnections, registerCommands } from '../login/command'
 
 export async function initialize(
     extensionContext: vscode.ExtensionContext,
diff --git a/packages/core/src/auth/ui/vue/show.ts b/packages/core/src/auth/ui/vue/show.ts
@@ -5,7 +5,7 @@
  * This module sets up the necessary components
  * for the webview to be shown.
  */
-import globals, { isWeb } from '../../../shared/extensionGlobals'
+import globals from '../../../shared/extensionGlobals'
 import { getIdeProperties, isCloud9 } from '../../../shared/extensionUtilities'
 import { VueWebview } from '../../../webviews/main'
 import * as vscode from 'vscode'
@@ -35,26 +35,24 @@ import {
     hasSso,
     BuilderIdKind,
     findSsoConnections,
-    authCommands,
 } from '../../utils'
 import { Region } from '../../../shared/regions/endpoints'
 import { CancellationError } from '../../../shared/utilities/timeoutUtils'
 import { validateSsoUrl, validateSsoUrlFormat } from '../../sso/validation'
-import { awsIdSignIn, showCodeWhispererConnectionPrompt } from '../../../codewhisperer/util/showSsoPrompt'
-import { AuthError, ServiceItemId, isServiceItemId, authFormTelemetryMapping, userCancelled } from './types'
+import { awsIdSignIn } from '../../../codewhisperer/util/showSsoPrompt'
+import { AuthError, ServiceItemId, authFormTelemetryMapping, userCancelled } from './types'
 import { connectToEnterpriseSso } from '../../../codewhisperer/util/getStartUrl'
 import { trustedDomainCancellation } from '../../sso/model'
 import { CredentialSourceId, Result, telemetry } from '../../../shared/telemetry/telemetry'
 import { AuthFormId } from './authForms/types'
 import { handleWebviewError } from '../../../webviews/server'
-import { Commands, RegisteredCommand, VsCodeCommandArg, placeholder } from '../../../shared/vscode/commands2'
+import { placeholder } from '../../../shared/vscode/commands2'
 import { ClassToInterfaceType } from '../../../shared/utilities/tsUtils'
 import { debounce } from 'lodash'
 import { submitFeedback } from '../../../feedback/vue/submitFeedback'
 import { InvalidGrantException } from '@aws-sdk/client-sso-oidc'
 import { ExtStartUpSources } from '../../../shared/telemetry'
-import { CommonAuthWebview } from '../../../login/webview/vue/backend'
-import { AuthSource, AuthSources } from '../../../login/webview/util'
+import { AuthSource } from '../../../login/webview/util'
 import { focusAmazonQPanel } from '../../../codewhispererChat/commands/registerCommands'
 
 // This file has some used functions, but most of it should be removed soon. We have a new
@@ -714,47 +712,6 @@ const Panel = VueWebview.compilePanel(AuthWebview)
 let activePanel: InstanceType<typeof Panel> | undefined
 let subscriptions: vscode.Disposable[] | undefined
 
-let showManageConnections: RegisteredCommand<any> | undefined
-export function getShowManageConnections(): RegisteredCommand<any> {
-    if (!showManageConnections) {
-        throw new Error('showManageConnections not registered')
-    }
-    return showManageConnections
-}
-
-export function registerCommands(context: vscode.ExtensionContext, prefix: string) {
-    showManageConnections = Commands.register(
-        { id: `aws.${prefix}.auth.manageConnections`, compositeKey: { 1: 'source' } },
-        async (_: VsCodeCommandArg, source: AuthSource, serviceToShow?: ServiceItemId) => {
-            if (_ !== placeholder) {
-                source = AuthSources.vscodeComponent
-            }
-
-            // The auth webview page does not make sense to use in C9,
-            // so show the auth quick pick instead.
-            if (isCloud9('any') || isWeb()) {
-                if (source.toLowerCase().includes('codewhisperer')) {
-                    // Show CW specific quick pick for CW connections
-                    return showCodeWhispererConnectionPrompt()
-                }
-                return authCommands().addConnection.execute()
-            }
-
-            if (!isServiceItemId(serviceToShow)) {
-                serviceToShow = undefined
-            }
-
-            // TODO: hack
-            if (prefix === 'toolkit') {
-                CommonAuthWebview.authSource = source
-                await vscode.commands.executeCommand('aws.explorer.setLoginService', serviceToShow)
-                await vscode.commands.executeCommand('setContext', 'aws.explorer.showAuthView', true)
-                await vscode.commands.executeCommand('aws.toolkit.AmazonCommonAuth.focus')
-            }
-        }
-    )
-}
-
 //todo: delete?
 export async function showAuthWebview(
     ctx: vscode.ExtensionContext,
diff --git a/packages/core/src/codecatalyst/activation.ts b/packages/core/src/codecatalyst/activation.ts
@@ -12,7 +12,7 @@ import { CodeCatalystCommands } from './commands'
 import { GitExtension } from '../shared/extensions/git'
 import { CodeCatalystAuthenticationProvider } from './auth'
 import { registerDevfileWatcher, updateDevfileCommand } from './devfile'
-import { DevEnvClient, DevEnvActivity } from '../shared/clients/devenvClient'
+import { DevEnvClient } from '../shared/clients/devenvClient'
 import { watchRestartingDevEnvs } from './reconnect'
 import { ToolkitPromptSettings } from '../shared/settings'
 import { dontShow } from '../shared/localizedText'
@@ -22,9 +22,9 @@ import { createClient, getCodeCatalystConfig } from '../shared/clients/codecatal
 import { isDevenvVscode } from './utils'
 import { codeCatalystConnectCommand, getThisDevEnv } from './model'
 import { getLogger } from '../shared/logger/logger'
-import { InactivityMessage, shouldTrackUserActivity } from './devEnv'
-import { getShowManageConnections } from '../auth/ui/vue/show'
+import { DevEnvActivityStarter } from './devEnv'
 import { learnMoreCommand, onboardCommand, reauth } from './explorer'
+import { getShowManageConnections } from '../login/command'
 
 const localize = nls.loadMessageBundle()
 
@@ -123,17 +123,10 @@ export async function activate(ctx: ExtContext): Promise<void> {
                 }
             })
         }
-
-        const maxInactivityMinutes = thisDevenv.summary.inactivityTimeoutMinutes
-        const devEnvClient = thisDevenv.devenvClient
-        const devEnvActivity = await DevEnvActivity.instanceIfActivityTrackingEnabled(devEnvClient)
-        if (shouldTrackUserActivity(maxInactivityMinutes) && devEnvActivity) {
-            const inactivityMessage = new InactivityMessage()
-            await inactivityMessage.setupMessage(maxInactivityMinutes, devEnvActivity)
-
-            ctx.extensionContext.subscriptions.push(inactivityMessage, devEnvActivity)
-        }
     }
+
+    // This must always be called on activation
+    DevEnvActivityStarter.register(authProvider)
 }
 
 async function showReadmeFileOnFirstLoad(workspaceState: vscode.ExtensionContext['workspaceState']): Promise<void> {
diff --git a/packages/core/src/codecatalyst/commands.ts b/packages/core/src/codecatalyst/commands.ts
@@ -25,8 +25,8 @@ import { AccountStatus } from '../shared/telemetry/telemetryClient'
 import { CreateDevEnvironmentRequest, UpdateDevEnvironmentRequest } from 'aws-sdk/clients/codecatalyst'
 import { Auth } from '../auth/auth'
 import { SsoConnection } from '../auth/connection'
-import { getShowManageConnections } from '../auth/ui/vue/show'
 import { isInDevEnv, isRemoteWorkspace } from '../shared/vscode/env'
+import { getShowManageConnections } from '../login/command'
 
 /** "List CodeCatalyst Commands" command. */
 export async function listCommands(): Promise<void> {
diff --git a/packages/core/src/codecatalyst/devEnv.ts b/packages/core/src/codecatalyst/devEnv.ts
@@ -7,10 +7,123 @@ import { DevEnvironment } from '../shared/clients/codecatalystClient'
 import { DevEnvActivity } from '../shared/clients/devenvClient'
 import globals from '../shared/extensionGlobals'
 import * as vscode from 'vscode'
-import { Timeout } from '../shared/utilities/timeoutUtils'
+import { Timeout, waitUntil } from '../shared/utilities/timeoutUtils'
 import { showMessageWithCancel } from '../shared/utilities/messages'
 import { isCloud9 } from '../shared/extensionUtilities'
 import { getLogger } from '../shared/logger'
+import { CodeCatalystAuthenticationProvider } from './auth'
+import { getThisDevEnv } from './model'
+import { isInDevEnv } from '../shared/vscode/env'
+import { shared } from '../shared/utilities/functionUtils'
+
+/**
+ * This class exists due to complexities with starting the {@link DevEnvActivity} Hearbeat mechanism.
+ * All logic to start it is contained within this class.
+ */
+export class DevEnvActivityStarter {
+    /**
+     * Trys to start the Dev Env Activity Heartbeat mechanism.
+     * This mechanism keeps the Dev Env from timing out.
+     */
+    public static register(authProvider: CodeCatalystAuthenticationProvider) {
+        if (!isInDevEnv()) {
+            getLogger().debug('codecatalyst: not in a devenv, not registering DevEnvActivityHeartbeatStarter')
+            return
+        }
+
+        DevEnvActivityStarter.authProvider = authProvider
+        void DevEnvActivityStarter.instance.tryStartDevEnvActivityHeartbeatWithRetry()
+    }
+
+    /** If true, the Activity Heartbeat is running */
+    private didStart: boolean = false
+
+    protected constructor(private readonly authProvider: CodeCatalystAuthenticationProvider) {}
+
+    private static _instance: DevEnvActivityStarter | undefined = undefined
+    private static get instance(): DevEnvActivityStarter {
+        return (DevEnvActivityStarter._instance ??= new DevEnvActivityStarter(DevEnvActivityStarter.getAuthProvider()))
+    }
+
+    /**
+     * Keeps executing {@link DevEnvActivityStarter.tryStartDevEnvActivityHeartbeat} on an interval until it succeeds.
+     * After a certain amount of time this will stop retrying.
+     */
+    private async tryStartDevEnvActivityHeartbeatWithRetry() {
+        return waitUntil(
+            async () => {
+                await this.tryStartDevEnvActivityHeartbeat()
+                return this.didStart
+            },
+            { interval: 20_000, timeout: 60_000 * 5 }
+        )
+    }
+
+    /** Trys to start the Activity Heartbeat mechanism */
+    private tryStartDevEnvActivityHeartbeat = shared(async () => {
+        if (this.didStart) {
+            return
+        }
+
+        const thisDevenv = (await getThisDevEnv(this.authProvider))?.unwrapOrElse(err => {
+            getLogger().warn('codecatalyst: failed to get current Dev Enviroment: %s', err)
+            return undefined
+        })
+
+        if (!thisDevenv) {
+            const connection = this.authProvider.activeConnection
+            if (connection) {
+                void vscode.window
+                    .showErrorMessage(
+                        'CodeCatalyst: Reauthenticate your connection or the Dev Environment will time out.',
+                        'Reauthenticate'
+                    )
+                    .then(async res => {
+                        if (res !== 'Reauthenticate') {
+                            return
+                        }
+                        await this.authProvider.auth.reauthenticate(connection)
+                        void this.tryStartDevEnvActivityHeartbeat()
+                    })
+
+                getLogger().warn('codecatalyst: dev env needs reauth to not time out')
+            } else {
+                getLogger().warn(`codecatalyst: dev env needs a connection to not time out`)
+            }
+            return
+        }
+
+        const maxInactivityMinutes = thisDevenv.summary.inactivityTimeoutMinutes
+        if (!shouldTrackUserActivity(maxInactivityMinutes)) {
+            getLogger().debug(
+                `codecatalyst: not tracking user inactivity due to inactivity minutes being: ${maxInactivityMinutes}`
+            )
+            return
+        }
+
+        const devEnvActivity = await DevEnvActivity.instanceIfActivityTrackingEnabled(thisDevenv.devenvClient)
+        if (!devEnvActivity) {
+            getLogger().debug(`codecatalyst: not tracking user inactivity since activity api is not enabled`)
+            return
+        }
+
+        // Everything is good, we can start the activity heartbeat now
+        devEnvActivity.setUpdateActivityOnIdeActivity(true)
+        const inactivityMessage = new InactivityMessage()
+        await inactivityMessage.setupMessage(maxInactivityMinutes, devEnvActivity)
+
+        globals.context.subscriptions.push(inactivityMessage, devEnvActivity)
+        this.didStart = true
+    })
+
+    private static authProvider: CodeCatalystAuthenticationProvider | undefined = undefined
+    private static getAuthProvider() {
+        if (!this.authProvider) {
+            throw new Error('DevEnvActivityHeartbeatStarter authProvider is not set')
+        }
+        return this.authProvider
+    }
+}
 
 /** If we should be sending the dev env activity timestamps to track user activity */
 export function shouldTrackUserActivity(maxInactivityMinutes: DevEnvironment['inactivityTimeoutMinutes']): boolean {
diff --git a/packages/core/src/codecatalyst/explorer.ts b/packages/core/src/codecatalyst/explorer.ts
@@ -16,7 +16,7 @@ import * as codecatalyst from './model'
 import { getLogger } from '../shared/logger'
 import { Connection } from '../auth/connection'
 import { openUrl } from '../shared/utilities/vsCodeUtils'
-import { getShowManageConnections } from '../auth/ui/vue/show'
+import { getShowManageConnections } from '../login/command'
 
 export const learnMoreCommand = Commands.declare('aws.learnMore', () => async (docsUrl: vscode.Uri) => {
     return openUrl(docsUrl)
diff --git a/packages/core/src/login/command.ts b/packages/core/src/login/command.ts
@@ -0,0 +1,54 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+import vscode from 'vscode'
+import { Commands, RegisteredCommand, VsCodeCommandArg, placeholder } from '../shared/vscode/commands2'
+import { ServiceItemId, isServiceItemId } from '../auth/ui/vue/types'
+import { authCommands } from '../auth/utils'
+import { showCodeWhispererConnectionPrompt } from '../codewhisperer/util/showSsoPrompt'
+import { AuthSource, AuthSources } from './webview/util'
+import { isCloud9 } from '../shared/extensionUtilities'
+import { isWeb } from '../shared/extensionGlobals'
+import { CommonAuthWebview } from './webview/vue/backend'
+
+let showManageConnections: RegisteredCommand<any> | undefined
+export function getShowManageConnections(): RegisteredCommand<any> {
+    if (!showManageConnections) {
+        throw new Error('showManageConnections not registered')
+    }
+    return showManageConnections
+}
+
+export function registerCommands(context: vscode.ExtensionContext, prefix: string) {
+    showManageConnections = Commands.register(
+        { id: `aws.${prefix}.auth.manageConnections`, compositeKey: { 1: 'source' } },
+        async (_: VsCodeCommandArg, source: AuthSource, serviceToShow?: ServiceItemId) => {
+            if (_ !== placeholder) {
+                source = AuthSources.vscodeComponent
+            }
+
+            // The auth webview page does not make sense to use in C9,
+            // so show the auth quick pick instead.
+            if (isCloud9('any') || isWeb()) {
+                if (source.toLowerCase().includes('codewhisperer')) {
+                    // Show CW specific quick pick for CW connections
+                    return showCodeWhispererConnectionPrompt()
+                }
+                return authCommands().addConnection.execute()
+            }
+
+            if (!isServiceItemId(serviceToShow)) {
+                serviceToShow = undefined
+            }
+
+            // TODO: hack
+            if (prefix === 'toolkit') {
+                CommonAuthWebview.authSource = source
+                await vscode.commands.executeCommand('aws.explorer.setLoginService', serviceToShow)
+                await vscode.commands.executeCommand('setContext', 'aws.explorer.showAuthView', true)
+                await vscode.commands.executeCommand('aws.toolkit.AmazonCommonAuth.focus')
+            }
+        }
+    )
+}
diff --git a/packages/core/src/shared/clients/devenvClient.ts b/packages/core/src/shared/clients/devenvClient.ts
diff --git a/packages/core/src/test/shared/clients/devenvClient.test.ts b/packages/core/src/test/shared/clients/devenvClient.test.ts
