Merge aws#5293 globalState abstraction

Author: justinmk3

packages/amazonq/src/extensionCommon.ts

@@ -3,35 +3,38 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import * as vscode from 'vscode'
-import * as semver from 'semver'
-import { join } from 'path'
+import { AuthUtils, CredentialsStore, LoginManager, SsoConnection, initializeAuth } from 'aws-core-vscode/auth'
 import {
+    AuthUtil,
     activate as activateCodeWhisperer,
     shutdown as shutdownCodeWhisperer,
-    AuthUtil,
 } from 'aws-core-vscode/codewhispererCommon'
+import { makeEndpointsProvider, registerGenericCommands } from 'aws-core-vscode/extensionCommon'
+import { CommonAuthWebview } from 'aws-core-vscode/login'
 import {
+    DefaultAWSClientBuilder,
+    DefaultAwsContext,
     ExtContext,
-    initialize,
+    RegionProvider,
+    Settings,
     activateLogger,
     activateTelemetry,
-    Settings,
-    DefaultAwsContext,
-    initializeComputeRegion,
-    DefaultAWSClientBuilder,
-    globals,
-    RegionProvider,
+    env,
+    errors,
+    fs,
     getLogger,
     getMachineId,
+    globals,
+    initialize,
+    initializeComputeRegion,
     messages,
+    setContext,
 } from 'aws-core-vscode/shared'
-import { fs, errors, setContext } from 'aws-core-vscode/shared'
-import { initializeAuth, CredentialsStore, LoginManager, AuthUtils, SsoConnection } from 'aws-core-vscode/auth'
-import { CommonAuthWebview } from 'aws-core-vscode/login'
+import { ExtStartUpSources, telemetry } from 'aws-core-vscode/telemetry'
 import { VSCODE_EXTENSION_ID } from 'aws-core-vscode/utils'
-import { telemetry, ExtStartUpSources } from 'aws-core-vscode/telemetry'
-import { makeEndpointsProvider, registerGenericCommands } from 'aws-core-vscode/extensionCommon'
+import { join } from 'path'
+import * as semver from 'semver'
+import * as vscode from 'vscode'
 import { registerCommands } from './commands'
 
 export const amazonQContextPrefix = 'amazonq'
@@ -44,7 +47,7 @@ export async function activateAmazonQCommon(context: vscode.ExtensionContext, is
     const homeDirLogs = await fs.init(context, homeDir => {
         void messages.showViewLogsMessage(`Invalid home directory (check $HOME): "${homeDir}"`)
     })
-    errors.init(fs.getUsername())
+    errors.init(fs.getUsername(), env.isAutomation())
     await initializeComputeRegion()
 
     globals.contextPrefix = 'amazonq.' //todo: disconnect from above line

packages/amazonq/test/unit/codewhisperer/service/telemetry.test.ts

@@ -26,7 +26,7 @@ import {
     vsCodeCursorUpdateDelay,
     AuthUtil,
 } from 'aws-core-vscode/codewhisperer'
-import { sleep, waitUntil, getMinVscodeVersion, CodewhispererUserTriggerDecision } from 'aws-core-vscode/shared'
+import { sleep, waitUntil, env, CodewhispererUserTriggerDecision } from 'aws-core-vscode/shared'
 import { resetCodeWhispererGlobalVariables } from 'aws-core-vscode/test'
 
 type CodeWhispererResponse = ListRecommendationsResponse & {
@@ -288,7 +288,7 @@ describe.skip('CodeWhisperer telemetry', async function () {
         // as per vscode official repo PR https://github.com/microsoft/vscode/commit/cb0e59c56677181b570b110167d13efb4ba7677d#diff-84b7f4a5ab7c383d86e2d40e2c704d255dc1e187a29386c036023a4696196556R19
         // navigation commands seem to be introduced since 1.78.0
         function shouldRun() {
-            const version = getMinVscodeVersion()
+            const version = env.getMinVscodeVersion()
             if (semver.gte(version, '1.78.0')) {
                 throw new Error('Minimum VSCode version is greater than 1.78.0, this check should be removed')
             }

packages/core/src/auth/sso/ssoAccessTokenProvider.ts

@@ -107,7 +107,7 @@ export abstract class SsoAccessTokenProvider {
         const access = await this.runFlow()
         const identity = this.tokenCacheKey
         await this.cache.token.save(identity, access)
-        await setSessionCreationDate(this.tokenCacheKey, new Date())
+        await globals.globalState.setSsoSessionCreationDate(this.tokenCacheKey, new globals.clock.Date())
 
         return { ...access.token, identity }
     }
@@ -309,25 +309,13 @@ async function pollForTokenWithProgress<T extends { requestId?: string }>(
     )
 }
 
-const sessionCreationDateKey = '#sessionCreationDates'
-async function setSessionCreationDate(id: string, date: Date, memento = globals.context.globalState) {
-    try {
-        await memento.update(sessionCreationDateKey, {
-            ...memento.get(sessionCreationDateKey),
-            [id]: date.getTime(),
-        })
-    } catch (err) {
-        getLogger().verbose('auth: failed to set session creation date: %s', err)
-    }
-}
-
-function getSessionCreationDate(id: string, memento = globals.context.globalState): number | undefined {
-    return memento.get(sessionCreationDateKey, {} as Record<string, number>)[id]
-}
-
-function getSessionDuration(id: string, memento = globals.context.globalState) {
-    const creationDate = getSessionCreationDate(id, memento)
-
+/**
+ * Gets SSO session creation timestamp for the given session `id`.
+ *
+ * @param id Session id
+ */
+function getSessionDuration(id: string) {
+    const creationDate = globals.globalState.getSsoSessionCreationDate(id)
     return creationDate !== undefined ? Date.now() - creationDate : undefined
 }
 

packages/core/src/awsService/redshift/activation.ts

@@ -10,14 +10,14 @@ import { RedshiftNotebookController } from './notebook/redshiftNotebookControlle
 import { CellStatusBarItemProvider } from './notebook/cellStatusBarItemProvider'
 import { Commands } from '../../shared/vscode/commands2'
 import { NotebookConnectionWizard, RedshiftNodeConnectionWizard } from './wizards/connectionWizard'
-import { ConnectionParams, ConnectionType } from './models/models'
+import { deleteConnection, ConnectionParams, ConnectionType } from './models/models'
 import { DefaultRedshiftClient } from '../../shared/clients/redshiftClient'
 import { localize } from '../../shared/utilities/vsCodeUtils'
 import { RedshiftWarehouseNode } from './explorer/redshiftWarehouseNode'
 import { ToolkitError } from '../../shared/errors'
-import { deleteConnection, updateConnectionParamsState } from './explorer/redshiftState'
 import { showViewLogsMessage } from '../../shared/utilities/messages'
 import { showConnectionMessage } from './messageUtils'
+import globals from '../../shared/extensionGlobals'
 
 export async function activate(ctx: ExtContext): Promise<void> {
     if ('NotebookEdit' in vscode) {
@@ -123,7 +123,10 @@ function getEditConnectionHandler() {
                     connectionParams.secret = secretArnFetched
                 }
                 redshiftWarehouseNode.setConnectionParams(connectionParams)
-                await updateConnectionParamsState(redshiftWarehouseNode.arn, redshiftWarehouseNode.connectionParams)
+                await globals.globalState.saveRedshiftConnection(
+                    redshiftWarehouseNode.arn,
+                    redshiftWarehouseNode.connectionParams
+                )
                 await vscode.commands.executeCommand('aws.refreshAwsExplorerNode', redshiftWarehouseNode)
             }
         } catch (error) {
@@ -135,7 +138,7 @@ function getEditConnectionHandler() {
 function getDeleteConnectionHandler() {
     return async (redshiftWarehouseNode: RedshiftWarehouseNode) => {
         redshiftWarehouseNode.connectionParams = undefined
-        await updateConnectionParamsState(redshiftWarehouseNode.arn, deleteConnection)
+        await globals.globalState.saveRedshiftConnection(redshiftWarehouseNode.arn, deleteConnection)
         await vscode.commands.executeCommand('aws.refreshAwsExplorerNode', redshiftWarehouseNode)
     }
 }

packages/core/src/awsService/redshift/explorer/redshiftState.ts

@@ -5,20 +5,20 @@
 import { AWSResourceNode } from '../../../shared/treeview/nodes/awsResourceNode'
 import { AWSTreeNodeBase } from '../../../shared/treeview/nodes/awsTreeNodeBase'
 import * as vscode from 'vscode'
+import globals from '../../../shared/extensionGlobals'
 import { RedshiftNode } from './redshiftNode'
 import { makeChildrenNodes } from '../../../shared/treeview/utils'
 import { RedshiftDatabaseNode } from './redshiftDatabaseNode'
 import { localize } from '../../../shared/utilities/vsCodeUtils'
 import { LoadMoreNode } from '../../../shared/treeview/nodes/loadMoreNode'
 import { ChildNodeLoader, ChildNodePage } from '../../../awsexplorer/childNodeLoader'
 import { DefaultRedshiftClient } from '../../../shared/clients/redshiftClient'
-import { ConnectionParams, ConnectionType, RedshiftWarehouseType } from '../models/models'
+import { deleteConnection, ConnectionParams, ConnectionType, RedshiftWarehouseType } from '../models/models'
 import { RedshiftNodeConnectionWizard } from '../wizards/connectionWizard'
 import { ListDatabasesResponse } from 'aws-sdk/clients/redshiftdata'
 import { getIcon } from '../../../shared/icons'
 import { AWSCommandTreeNode } from '../../../shared/treeview/nodes/awsCommandTreeNode'
 import { telemetry } from '../../../shared/telemetry/telemetry'
-import { deleteConnection, getConnectionParamsState, updateConnectionParamsState } from './redshiftState'
 import { createLogsConnectionMessage, showConnectionMessage } from '../messageUtils'
 
 export class CreateNotebookNode extends AWSCommandTreeNode {
@@ -50,7 +50,7 @@ export class RedshiftWarehouseNode extends AWSTreeNodeBase implements AWSResourc
         this.arn = redshiftWarehouse.arn
         this.name = redshiftWarehouse.name
         this.redshiftClient = parent.redshiftClient
-        const existingConnectionParams = getConnectionParamsState(this.arn)
+        const existingConnectionParams = globals.globalState.getRedshiftConnection(this.arn)
         if (existingConnectionParams && existingConnectionParams !== deleteConnection) {
             this.connectionParams = existingConnectionParams as ConnectionParams
             this.iconPath = getIcon('aws-redshift-cluster-connected')
@@ -108,14 +108,14 @@ export class RedshiftWarehouseNode extends AWSTreeNodeBase implements AWSResourc
         return await makeChildrenNodes({
             getChildNodes: async () => {
                 this.childLoader.clearChildren()
-                const existingConnectionParams = getConnectionParamsState(this.arn)
+                const existingConnectionParams = globals.globalState.getRedshiftConnection(this.arn)
                 if (existingConnectionParams && existingConnectionParams === deleteConnection) {
                     // connection is deleted but explorer is not refreshed: return clickToEstablishConnectionNode
-                    await updateConnectionParamsState(this.arn, undefined)
+                    await globals.globalState.saveRedshiftConnection(this.arn, undefined)
                     return this.getClickToEstablishConnectionNode()
-                } else if (existingConnectionParams && existingConnectionParams !== deleteConnection) {
+                } else if (existingConnectionParams) {
                     // valid connectionParams: update the redshiftWarehouseNode
-                    this.connectionParams = existingConnectionParams as ConnectionParams
+                    this.connectionParams = existingConnectionParams
                 } else {
                     // No connectionParams: trigger connection wizard to get user input
                     this.connectionParams = await new RedshiftNodeConnectionWizard(this).run()
@@ -137,11 +137,11 @@ export class RedshiftWarehouseNode extends AWSTreeNodeBase implements AWSResourc
                     const childNodes = await this.childLoader.getChildren()
                     const startButtonNode = new CreateNotebookNode(this)
                     childNodes.unshift(startButtonNode)
-                    await updateConnectionParamsState(this.arn, this.connectionParams)
+                    await globals.globalState.saveRedshiftConnection(this.arn, this.connectionParams)
                     return childNodes
                 } catch (error) {
                     showConnectionMessage(this.redshiftWarehouse.name, error as Error)
-                    await updateConnectionParamsState(this.arn, undefined)
+                    await globals.globalState.saveRedshiftConnection(this.arn, undefined)
                     return this.getRetryNode()
                 }
             },

packages/core/src/awsService/redshift/explorer/redshiftWarehouseNode.ts

@@ -5,6 +5,10 @@
 
 import { Region } from '../../../shared/regions/endpoints'
 
+// Sigil treated such that the connection wizard is not triggered during explorer node refresh after
+// connection deletion.
+export const deleteConnection = 'DELETE_CONNECTION' as const
+
 export class ConnectionParams {
     constructor(
         public readonly connectionType: ConnectionType,

packages/core/src/awsService/redshift/models/models.ts

@@ -35,7 +35,7 @@ import { LoginManager } from './auth/deprecated/loginManager'
 import { CredentialsStore } from './auth/credentials/store'
 import { initializeAwsCredentialsStatusBarItem } from './auth/ui/statusBarItem'
 import { RegionProvider, getEndpointsFromFetcher } from './shared/regions/regionProvider'
-import { getMachineId } from './shared/vscode/env'
+import { getMachineId, isAutomation } from './shared/vscode/env'
 import { registerCommandErrorHandler } from './shared/vscode/commands2'
 import { registerWebviewErrorHandler } from './webviews/server'
 import { showQuickStartWebview } from './shared/extensionStartup'
@@ -77,7 +77,7 @@ export async function activateCommon(
     const homeDirLogs = await fs.init(context, homeDir => {
         void showViewLogsMessage(`Invalid home directory (check $HOME): "${homeDir}"`)
     })
-    errors.init(fs.getUsername())
+    errors.init(fs.getUsername(), isAutomation())
     await initializeComputeRegion()
 
     globals.contextPrefix = '' //todo: disconnect supplied argument

packages/core/src/extensionCommon.ts

@@ -13,14 +13,15 @@ import { isNonNullable } from './utilities/tsUtils'
 import type * as nodefs from 'fs'
 import type * as os from 'os'
 import { CodeWhispererStreamingServiceException } from '@amzn/codewhisperer-streaming'
-import { isAutomation } from './vscode/env'
 import { driveLetterRegex } from './utilities/pathUtils'
 
 let _username = 'unknown-user'
+let _isAutomation = false
 
-/** One-time initialization for this module. */
-export function init(username: string) {
+/** Performs one-time initialization, to avoid circular dependencies. */
+export function init(username: string, isAutomation: boolean) {
     _username = username
+    _isAutomation = isAutomation
 }
 
 export const errorCode = {
@@ -667,7 +668,7 @@ function vscodeModeToString(mode: vscode.FileStat['permissions']) {
     }
 
     // XXX: future-proof in case vscode.FileStat.permissions gains more granularity.
-    if (isAutomation()) {
+    if (_isAutomation) {
         throw new Error('vscode.FileStat.permissions gained new fields, update this logic')
     }
 }

packages/core/src/shared/errors.ts

@@ -147,7 +147,7 @@ export function initialize(context: ExtensionContext, isWeb: boolean = false): T
         context,
         clock: copyClock(),
         didReload: checkDidReload(context),
-        globalState: new GlobalState(context),
+        globalState: new GlobalState(context.globalState),
         manifestPaths: {} as ToolkitGlobals['manifestPaths'],
         visualizationResourcePaths: {} as ToolkitGlobals['visualizationResourcePaths'],
         isWeb,
@@ -171,7 +171,7 @@ export { globals as default }
  */
 interface ToolkitGlobals {
     readonly context: ExtensionContext
-    /** Global, shared, mutable, persisted state (survives IDE restart), namespaced to the extension (i.e. not shared with other vscode extensions). */
+    /** Global, shared (with all vscode instances, including remote!), mutable, persisted state (survives IDE restart), namespaced to the extension (not shared with other vscode extensions). */
     readonly globalState: GlobalState
     /** Decides the prefix for package.json extension parameters, e.g. commands, 'setContext' values, etc. */
     contextPrefix: string