diff --git a/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/artifactManager.ts b/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/artifactManager.ts index 2c26b2f552..49a32e8e97 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/artifactManager.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/artifactManager.ts @@ -69,6 +69,7 @@ export class ArtifactManager { private workspaceFolders: WorkspaceFolder[] // TODO, how to handle when two workspace folders have the same name but different URI private filesByWorkspaceFolderAndLanguage: Map> + private isDisposed: boolean = false constructor(workspace: Workspace, logging: Logging, workspaceFolders: WorkspaceFolder[]) { this.workspace = workspace @@ -124,7 +125,17 @@ export class ArtifactManager { return zipFileMetadata } - async removeWorkspaceFolders(workspaceFolders: WorkspaceFolder[]): Promise { + public resetFromDisposal(): void { + this.isDisposed = false + } + + dispose(): void { + this.filesByWorkspaceFolderAndLanguage.clear() + this.workspaceFolders = [] + this.isDisposed = true + } + + removeWorkspaceFolders(workspaceFolders: WorkspaceFolder[]): void { workspaceFolders.forEach(workspaceToRemove => { // Find the matching workspace folder by URI let folderToDelete: WorkspaceFolder | undefined @@ -543,6 +554,9 @@ export class ArtifactManager { } for (const workspaceFolder of workspaceFolders) { + if (this.isDisposed) { + break + } const workspacePath = URI.parse(workspaceFolder.uri).path try { @@ -579,6 +593,9 @@ export class ArtifactManager { const zipFileMetadata: FileMetadata[] = [] await this.updateWorkspaceFiles(workspaceFolder, filesByLanguage) for (const [language, files] of filesByLanguage.entries()) { + if (this.isDisposed) { + break + } // Generate java .classpath and .project files const processedFiles = language === 'java' ? await this.processJavaProjectConfig(workspaceFolder, files) : files diff --git a/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/dependency/dependencyDiscoverer.ts b/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/dependency/dependencyDiscoverer.ts index f0d41e9801..027bba4796 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/dependency/dependencyDiscoverer.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/dependency/dependencyDiscoverer.ts @@ -3,7 +3,11 @@ import * as fs from 'fs' import { Logging, Workspace, WorkspaceFolder } from '@aws/language-server-runtimes/server-interface' import { URI } from 'vscode-uri' import { DependencyHandlerFactory } from './dependencyHandler/LanguageDependencyHandlerFactory' -import { BaseDependencyInfo, LanguageDependencyHandler } from './dependencyHandler/LanguageDependencyHandler' +import { + BaseDependencyInfo, + DependencyHandlerSharedState, + LanguageDependencyHandler, +} from './dependencyHandler/LanguageDependencyHandler' import { ArtifactManager } from '../artifactManager' import { supportedWorkspaceContextLanguages } from '../../../shared/languageDetection' @@ -12,8 +16,7 @@ export class DependencyDiscoverer { private workspaceFolders: WorkspaceFolder[] public dependencyHandlerRegistry: LanguageDependencyHandler[] = [] private initializedWorkspaceFolder = new Map() - // Create a SharedArrayBuffer with 4 bytes (for a 32-bit unsigned integer) for thread-safe counter - protected dependencyUploadedSizeSum = new Uint32Array(new SharedArrayBuffer(4)) + private sharedState: DependencyHandlerSharedState = { isDisposed: false, dependencyUploadedSizeSum: 0 } constructor( workspace: Workspace, @@ -23,7 +26,6 @@ export class DependencyDiscoverer { ) { this.workspaceFolders = workspaceFolders this.logging = logging - this.dependencyUploadedSizeSum[0] = 0 let jstsHandlerCreated = false supportedWorkspaceContextLanguages.forEach(language => { @@ -33,7 +35,7 @@ export class DependencyDiscoverer { logging, workspaceFolders, artifactManager, - this.dependencyUploadedSizeSum + this.sharedState ) if (handler) { // Share handler for javascript and typescript @@ -135,8 +137,9 @@ export class DependencyDiscoverer { } async reSyncDependenciesToS3(folders: WorkspaceFolder[]) { - Atomics.store(this.dependencyUploadedSizeSum, 0, 0) + this.sharedState.dependencyUploadedSizeSum = 0 for (const dependencyHandler of this.dependencyHandlerRegistry) { + dependencyHandler.markAllDependenciesAsUnZipped() await dependencyHandler.zipDependencyMap(folders) } } @@ -150,12 +153,17 @@ export class DependencyDiscoverer { } } + public resetFromDisposal(): void { + this.sharedState.isDisposed = false + this.sharedState.dependencyUploadedSizeSum = 0 + } + public dispose(): void { this.initializedWorkspaceFolder.clear() this.dependencyHandlerRegistry.forEach(dependencyHandler => { dependencyHandler.dispose() }) - Atomics.store(this.dependencyUploadedSizeSum, 0, 0) + this.sharedState.isDisposed = true } public disposeWorkspaceFolder(workspaceFolder: WorkspaceFolder) { diff --git a/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/dependency/dependencyHandler/LanguageDependencyHandler.ts b/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/dependency/dependencyHandler/LanguageDependencyHandler.ts index 226e16ef57..e381d7625f 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/dependency/dependencyHandler/LanguageDependencyHandler.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/dependency/dependencyHandler/LanguageDependencyHandler.ts @@ -19,6 +19,11 @@ export interface BaseDependencyInfo { workspaceFolder: WorkspaceFolder } +export interface DependencyHandlerSharedState { + isDisposed: boolean + dependencyUploadedSizeSum: number +} + // Abstract base class for all language dependency handlers export abstract class LanguageDependencyHandler { public language: CodewhispererLanguage @@ -28,7 +33,7 @@ export abstract class LanguageDependencyHandler { // key: workspaceFolder, value: {key: dependency name, value: Dependency} protected dependencyMap = new Map>() protected dependencyUploadedSizeMap = new Map() - protected dependencyUploadedSizeSum: Uint32Array + protected dependencyHandlerSharedState: DependencyHandlerSharedState protected dependencyWatchers: Map = new Map() protected artifactManager: ArtifactManager protected dependenciesFolderName: string @@ -48,7 +53,7 @@ export abstract class LanguageDependencyHandler { workspaceFolders: WorkspaceFolder[], artifactManager: ArtifactManager, dependenciesFolderName: string, - dependencyUploadedSizeSum: Uint32Array + dependencyHandlerSharedState: DependencyHandlerSharedState ) { this.language = language this.workspace = workspace @@ -62,7 +67,7 @@ export abstract class LanguageDependencyHandler { this.workspaceFolders.forEach(workSpaceFolder => this.dependencyMap.set(workSpaceFolder, new Map()) ) - this.dependencyUploadedSizeSum = dependencyUploadedSizeSum + this.dependencyHandlerSharedState = dependencyHandlerSharedState } /* @@ -126,6 +131,9 @@ export abstract class LanguageDependencyHandler { async zipDependencyMap(folders: WorkspaceFolder[]): Promise { // Process each workspace folder sequentially for (const [workspaceFolder, correspondingDependencyMap] of this.dependencyMap) { + if (this.dependencyHandlerSharedState.isDisposed) { + return + } // Check if the workspace folder is in the provided folders if (!folders.includes(workspaceFolder)) { continue @@ -144,6 +152,9 @@ export abstract class LanguageDependencyHandler { let currentChunkSize = 0 let currentChunk: Dependency[] = [] for (const dependency of dependencyList) { + if (this.dependencyHandlerSharedState.isDisposed) { + return + } // If adding this dependency would exceed the chunk size limit, // process the current chunk first if (currentChunkSize + dependency.size > MAX_CHUNK_SIZE_BYTES && currentChunk.length > 0) { @@ -177,7 +188,7 @@ export abstract class LanguageDependencyHandler { workspaceFolder, (this.dependencyUploadedSizeMap.get(workspaceFolder) || 0) + dependency.size ) - Atomics.add(this.dependencyUploadedSizeSum, 0, dependency.size) + this.dependencyHandlerSharedState.dependencyUploadedSizeSum += dependency.size // Mark this dependency that has been zipped dependency.zipped = true this.dependencyMap.get(workspaceFolder)?.set(dependency.name, dependency) @@ -298,8 +309,7 @@ export abstract class LanguageDependencyHandler { * However, everytime flare server restarts, this dependency map will be initialized. */ private validateWorkspaceDependencySize(workspaceFolder: WorkspaceFolder): boolean { - let uploadedSize = Atomics.load(this.dependencyUploadedSizeSum, 0) - if (uploadedSize && this.MAX_WORKSPACE_DEPENDENCY_SIZE < uploadedSize) { + if (this.MAX_WORKSPACE_DEPENDENCY_SIZE < this.dependencyHandlerSharedState.dependencyUploadedSizeSum) { return false } return true @@ -314,7 +324,8 @@ export abstract class LanguageDependencyHandler { disposeWorkspaceFolder(workspaceFolder: WorkspaceFolder): void { this.dependencyMap.delete(workspaceFolder) - Atomics.sub(this.dependencyUploadedSizeSum, 0, this.dependencyUploadedSizeMap.get(workspaceFolder) || 0) + this.dependencyHandlerSharedState.dependencyUploadedSizeSum -= + this.dependencyUploadedSizeMap.get(workspaceFolder) || 0 this.dependencyUploadedSizeMap.delete(workspaceFolder) this.disposeWatchers(workspaceFolder) this.disposeDependencyInfo(workspaceFolder) @@ -353,4 +364,12 @@ export abstract class LanguageDependencyHandler { protected isDependencyZipped(dependencyName: string, workspaceFolder: WorkspaceFolder): boolean | undefined { return this.dependencyMap.get(workspaceFolder)?.get(dependencyName)?.zipped } + + markAllDependenciesAsUnZipped(): void { + this.dependencyMap.forEach(correspondingDependencyMap => { + correspondingDependencyMap.forEach(dependency => { + dependency.zipped = false + }) + }) + } } diff --git a/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/dependency/dependencyHandler/LanguageDependencyHandlerFactory.ts b/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/dependency/dependencyHandler/LanguageDependencyHandlerFactory.ts index d5cadbe520..a6ce6892f1 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/dependency/dependencyHandler/LanguageDependencyHandlerFactory.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/dependency/dependencyHandler/LanguageDependencyHandlerFactory.ts @@ -1,7 +1,11 @@ import { JavaDependencyHandler } from './JavaDependencyHandler' import { PythonDependencyHandler } from './PythonDependencyHandler' import { JSTSDependencyHandler } from './JSTSDependencyHandler' -import { BaseDependencyInfo, LanguageDependencyHandler } from './LanguageDependencyHandler' +import { + BaseDependencyInfo, + DependencyHandlerSharedState, + LanguageDependencyHandler, +} from './LanguageDependencyHandler' import { Logging, Workspace, WorkspaceFolder } from '@aws/language-server-runtimes/server-interface' import { ArtifactManager } from '../../artifactManager' import { CodewhispererLanguage } from '../../../../shared/languageDetection' @@ -13,7 +17,7 @@ export class DependencyHandlerFactory { logging: Logging, workspaceFolders: WorkspaceFolder[], artifactManager: ArtifactManager, - dependencyUploadedSizeSum: Uint32Array + dependencyHandlerSharedState: DependencyHandlerSharedState ): LanguageDependencyHandler | null { switch (language.toLowerCase()) { case 'python': @@ -24,7 +28,7 @@ export class DependencyHandlerFactory { workspaceFolders, artifactManager, 'site-packages', - dependencyUploadedSizeSum + dependencyHandlerSharedState ) case 'javascript': case 'typescript': @@ -35,7 +39,7 @@ export class DependencyHandlerFactory { workspaceFolders, artifactManager, 'node_modules', - dependencyUploadedSizeSum + dependencyHandlerSharedState ) case 'java': return new JavaDependencyHandler( @@ -45,7 +49,7 @@ export class DependencyHandlerFactory { workspaceFolders, artifactManager, 'dependencies', - dependencyUploadedSizeSum + dependencyHandlerSharedState ) default: return null diff --git a/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/fileUploadJobManager.ts b/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/fileUploadJobManager.ts index beac88cb51..3f35ded00a 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/fileUploadJobManager.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/fileUploadJobManager.ts @@ -157,6 +157,8 @@ export class FileUploadJobManager { } public dispose(): void { - clearInterval(this.jobConsumerInterval) + if (this.jobConsumerInterval) { + clearInterval(this.jobConsumerInterval) + } } } diff --git a/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/workspaceContextServer.test.ts b/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/workspaceContextServer.test.ts index f1c342c8a2..f3bbb28a2a 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/workspaceContextServer.test.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/workspaceContextServer.test.ts @@ -2,14 +2,13 @@ import { InitializeParams, Server } from '@aws/language-server-runtimes/server-i import { TestFeatures } from '@aws/language-server-runtimes/testing' import sinon from 'ts-sinon' import { WorkspaceContextServer } from './workspaceContextServer' -import { AmazonQTokenServiceManager } from '../../shared/amazonQServiceManager/AmazonQTokenServiceManager' describe('WorkspaceContext Server', () => { let features: TestFeatures let server: Server let disposeServer: () => void - before(() => { + beforeEach(() => { features = new TestFeatures() server = WorkspaceContextServer() disposeServer = server(features) @@ -45,4 +44,149 @@ describe('WorkspaceContext Server', () => { sinon.assert.calledWith(features.logging.warn, sinon.match(/No workspaceIdentifier set/)) }) }) + + describe('UpdateConfiguration', () => { + it('should opt in for VSCode extension with server-sideContext enabled', async () => { + features.lsp.getClientInitializeParams.returns({ + initializationOptions: { + aws: { + clientInfo: { + name: 'AmazonQ-For-VSCode', + version: '0.0.1', + extension: { + name: 'AmazonQ-For-VSCode', + version: '0.0.1', + }, + }, + }, + }, + } as InitializeParams) + + features.lsp.workspace.getConfiguration.withArgs('amazonQ').resolves({ + 'server-sideContext': true, + }) + + await features.initialize(server) + await features.doChangeConfiguration() + + sinon.assert.calledWith(features.logging.log, sinon.match(/Workspace context server opt-in flag is: true/)) + }) + + it('should opt out for VSCode extension with server-sideContext disabled', async () => { + features.lsp.getClientInitializeParams.returns({ + initializationOptions: { + aws: { + clientInfo: { + name: 'AmazonQ-For-VSCode', + version: '0.0.1', + extension: { + name: 'AmazonQ-For-VSCode', + version: '0.0.1', + }, + }, + }, + }, + } as InitializeParams) + + features.lsp.workspace.getConfiguration.withArgs('amazonQ').resolves({ + 'server-sideContext': false, + }) + + await features.initialize(server) + await features.doChangeConfiguration() + + sinon.assert.calledWith(features.logging.log, sinon.match(/Workspace context server opt-in flag is: false/)) + }) + + it('should opt in for VSCode extension with server-sideContext missing for internal & BuilderID users', async () => { + features.lsp.getClientInitializeParams.returns({ + initializationOptions: { + aws: { + clientInfo: { + name: 'AmazonQ-For-VSCode', + version: '0.0.1', + extension: { + name: 'AmazonQ-For-VSCode', + version: '0.0.1', + }, + }, + }, + }, + } as InitializeParams) + + features.lsp.workspace.getConfiguration.withArgs('amazonQ').resolves({}) + await features.initialize(server) + + // Internal users + features.credentialsProvider.getConnectionMetadata.returns({ + sso: { + startUrl: 'https://amzn.awsapps.com/start', + }, + }) + await features.doChangeConfiguration() + sinon.assert.calledWith(features.logging.log, sinon.match(/Workspace context server opt-in flag is: true/)) + + // BuilderID users + sinon.restore() + features.credentialsProvider.getConnectionMetadata.returns({ + sso: { + startUrl: 'https://view.awsapps.com/start', + }, + }) + await features.doChangeConfiguration() + sinon.assert.calledWith(features.logging.log, sinon.match(/Workspace context server opt-in flag is: true/)) + }) + + it('should opt in for JetBrains extension with server-sideContext enabled', async () => { + features.lsp.getClientInitializeParams.returns({ + initializationOptions: { + aws: { + clientInfo: { + name: 'Amazon Q For JetBrains', + version: '0.0.1', + extension: { + name: 'Amazon Q For JetBrains', + version: '0.0.1', + }, + }, + }, + }, + } as InitializeParams) + + features.lsp.workspace.getConfiguration.withArgs('aws.codeWhisperer').resolves({ + workspaceContext: true, + }) + + await features.initialize(server) + await features.doChangeConfiguration() + + sinon.assert.calledWith(features.logging.log, sinon.match(/Workspace context server opt-in flag is: true/)) + }) + + it('should opt out for JetBrains extension with server-sideContext disabled', async () => { + features.lsp.getClientInitializeParams.returns({ + initializationOptions: { + aws: { + clientInfo: { + name: 'Amazon Q For JetBrains', + version: '0.0.1', + extension: { + name: 'Amazon Q For JetBrains', + version: '0.0.1', + }, + }, + }, + }, + } as InitializeParams) + + features.lsp.workspace.getConfiguration.withArgs('aws.codeWhisperer').resolves({ + workspaceContext: false, + }) + + await features.initialize(server) + await features.doChangeConfiguration() + + sinon.assert.calledWith(features.logging.log, sinon.match(/Workspace context server opt-in flag is: false/)) + }) + }) }) diff --git a/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/workspaceContextServer.ts b/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/workspaceContextServer.ts index 76c7275c41..c474869eab 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/workspaceContextServer.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/workspaceContextServer.ts @@ -23,7 +23,7 @@ import { AmazonQTokenServiceManager } from '../../shared/amazonQServiceManager/A import { FileUploadJobManager, FileUploadJobType } from './fileUploadJobManager' import { DependencyEventBundler } from './dependency/dependencyEventBundler' import ignore = require('ignore') -import { INTERNAL_USER_START_URL } from '../../shared/constants' +import { BUILDER_ID_START_URL, INTERNAL_USER_START_URL } from '../../shared/constants' const Q_CONTEXT_CONFIGURATION_SECTION = 'aws.q.workspaceContext' @@ -141,19 +141,35 @@ export const WorkspaceContextServer = (): Server => features => { const updateConfiguration = async () => { try { - let workspaceContextConfig = (await lsp.workspace.getConfiguration('amazonQ.workspaceContext')) || false - const configJetBrains = await lsp.workspace.getConfiguration('aws.codeWhisperer') - if (configJetBrains) { - workspaceContextConfig = workspaceContextConfig || configJetBrains['workspaceContext'] + const clientInitializParams = safeGet(lsp.getClientInitializeParams()) + const extensionName = clientInitializParams.initializationOptions?.aws?.clientInfo?.extension.name + if (extensionName === 'AmazonQ-For-VSCode') { + const amazonQSettings = (await lsp.workspace.getConfiguration('amazonQ'))?.['server-sideContext'] + isOptedIn = amazonQSettings || false + + // We want this temporary override for Amazon internal users and BuilderId users who are still using + // the old VSCode extension versions. Will remove this later. + if (amazonQSettings === undefined) { + const startUrl = credentialsProvider.getConnectionMetadata()?.sso?.startUrl + const isInternalOrBuilderIdUser = + startUrl && + (startUrl.includes(INTERNAL_USER_START_URL) || startUrl.includes(BUILDER_ID_START_URL)) + if (isInternalOrBuilderIdUser) { + isOptedIn = true + } + } + } else { + isOptedIn = (await lsp.workspace.getConfiguration('aws.codeWhisperer'))?.['workspaceContext'] || false } - - // TODO, removing client side opt in temporarily - isOptedIn = true - // isOptedIn = workspaceContextConfig === true + logging.log(`Workspace context server opt-in flag is: ${isOptedIn}`) if (!isOptedIn) { isWorkflowInitialized = false - await workspaceFolderManager.clearAllWorkspaceResources() + fileUploadJobManager?.dispose() + dependencyEventBundler?.dispose() + workspaceFolderManager.clearAllWorkspaceResources() + // Delete remote workspace when user chooses to opt-out + await workspaceFolderManager.deleteRemoteWorkspace() } } catch (error) { logging.error(`Error in getConfiguration: ${error}`) @@ -292,19 +308,18 @@ export const WorkspaceContextServer = (): Server => features => { fileUploadJobManager.startFileUploadJobConsumer() dependencyEventBundler.startDependencyEventBundler() - workspaceFolderManager.initializeWorkspaceStatusMonitor().catch(error => { - logging.error(`Error while initializing workspace status monitoring: ${error}`) - }) + await Promise.all([ + workspaceFolderManager.initializeWorkspaceStatusMonitor(), + workspaceFolderManager.processNewWorkspaceFolders(workspaceFolders), + ]) logging.log(`Workspace context workflow initialized`) - artifactManager.updateWorkspaceFolders(workspaceFolders) - workspaceFolderManager.processNewWorkspaceFolders(workspaceFolders).catch(error => { - logging.error(`Error while processing new workspace folders: ${error}`) - }) } else if (!isLoggedIn) { if (isWorkflowInitialized) { // If user is not logged in but the workflow is marked as initialized, it means user was logged in and is now logged out // In this case, clear the resources and stop the monitoring - await workspaceFolderManager.clearAllWorkspaceResources() + fileUploadJobManager?.dispose() + dependencyEventBundler?.dispose() + workspaceFolderManager.clearAllWorkspaceResources() } isWorkflowInitialized = false } @@ -523,11 +538,7 @@ export const WorkspaceContextServer = (): Server => features => { dependencyEventBundler.dispose() } if (workspaceFolderManager) { - workspaceFolderManager.clearAllWorkspaceResources().catch(error => { - logging.warn( - `Error while clearing workspace resources: ${error instanceof Error ? error.message : 'Unknown error'}` - ) - }) + workspaceFolderManager.clearAllWorkspaceResources() } } } diff --git a/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/workspaceFolderManager.ts b/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/workspaceFolderManager.ts index d0147219eb..bbab80e2d1 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/workspaceFolderManager.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/workspaceFolderManager.ts @@ -37,13 +37,17 @@ export class WorkspaceFolderManager { private static instance: WorkspaceFolderManager | undefined private readonly workspaceIdentifier: string private workspaceState: WorkspaceState - private remoteWorkspaceIdPromise: Promise - private remoteWorkspaceIdResolver!: (id: string) => void + // Promise that gates operations until workspace ID is ready or cancelled + private remoteWorkspaceIdPromise: Promise + // Resolves the remoteWorkspaceIdPromise to signal whether operations should proceed + private remoteWorkspaceIdResolver!: (id: boolean) => void + // Tracks whether the existing remoteWorkspaceIdPromise has been resolved + private remoteWorkspaceIdPromiseResolved: boolean = false private workspaceFolders: WorkspaceFolder[] private credentialsProvider: CredentialsProvider private readonly INITIAL_CHECK_INTERVAL = 40 * 1000 // 40 seconds private readonly INITIAL_CONNECTION_TIMEOUT = 2 * 60 * 1000 // 2 minutes - private readonly CONTINUOUS_MONITOR_INTERVAL = 30 * 60 * 1000 // 30 minutes + private readonly CONTINUOUS_MONITOR_INTERVAL = 5 * 60 * 1000 // 30 minutes private readonly MESSAGE_PUBLISH_INTERVAL: number = 100 // 100 milliseconds private continuousMonitorInterval: NodeJS.Timeout | undefined private optOutMonitorInterval: NodeJS.Timeout | undefined @@ -105,7 +109,7 @@ export class WorkspaceFolderManager { }) }) - this.remoteWorkspaceIdPromise = new Promise(resolve => { + this.remoteWorkspaceIdPromise = new Promise(resolve => { this.remoteWorkspaceIdResolver = resolve }) this.workspaceState = { @@ -137,7 +141,10 @@ export class WorkspaceFolderManager { async processNewWorkspaceFolders(folders: WorkspaceFolder[]) { // Wait for remote workspace id - await this.waitForRemoteWorkspaceId() + const shouldProceed = await this.remoteWorkspaceIdPromise + if (!shouldProceed) { + return + } // Sync workspace source codes await this.syncSourceCodesToS3(folders).catch(e => { @@ -201,12 +208,15 @@ export class WorkspaceFolderManager { return s3Url } - async clearAllWorkspaceResources() { + clearAllWorkspaceResources() { this.stopContinuousMonitoring() + this.stopOptOutMonitoring() + this.remoteWorkspaceIdResolver(false) + this.remoteWorkspaceIdPromiseResolved = true this.stopMessageQueueConsumer() - this.resetRemoteWorkspaceId() this.workspaceState.webSocketClient?.destroyClient() this.dependencyDiscoverer.dispose() + this.artifactManager.dispose() } /** @@ -215,7 +225,10 @@ export class WorkspaceFolderManager { * @param workspaceFolder */ async processWorkspaceFoldersDeletion(workspaceFolders: WorkspaceFolder[]) { - const workspaceId = await this.waitForRemoteWorkspaceId() + const shouldProceed = await this.remoteWorkspaceIdPromise + if (!shouldProceed) { + return + } for (const folder of workspaceFolders) { const languagesMap = this.artifactManager.getLanguagesForWorkspaceFolder(folder) const programmingLanguages = languagesMap ? Array.from(languagesMap.keys()) : [] @@ -234,7 +247,7 @@ export class WorkspaceFolderManager { ], }, workspaceChangeMetadata: { - workspaceId: workspaceId, + workspaceId: this.workspaceState.workspaceId, programmingLanguage: language, }, }, @@ -243,7 +256,7 @@ export class WorkspaceFolderManager { } this.dependencyDiscoverer.disposeWorkspaceFolder(folder) } - await this.artifactManager.removeWorkspaceFolders(workspaceFolders) + this.artifactManager.removeWorkspaceFolders(workspaceFolders) } private async uploadDependencyZipAndQueueEvent(zip: FileMetadata, addWSFolderPathInS3: boolean): Promise { @@ -302,36 +315,43 @@ export class WorkspaceFolderManager { async initializeWorkspaceStatusMonitor() { this.logging.log(`Initializing workspace status check for workspace [${this.workspaceIdentifier}]`) + // Reset workspace ID to force operations to wait for new remote workspace information + this.resetRemoteWorkspaceId() + + this.artifactManager.resetFromDisposal() + this.dependencyDiscoverer.resetFromDisposal() + // Set up message queue consumer - this.messageQueueConsumerInterval = setInterval(() => { - if (this.workspaceState.webSocketClient && this.workspaceState.webSocketClient.isConnected()) { - const message = this.workspaceState.messageQueue[0] - if (message) { - try { - this.workspaceState.webSocketClient.send(message) - this.workspaceState.messageQueue.shift() - } catch (error) { - this.logging.error(`Error sending message: ${error}`) + if (this.messageQueueConsumerInterval === undefined) { + this.messageQueueConsumerInterval = setInterval(() => { + if (this.workspaceState.webSocketClient && this.workspaceState.webSocketClient.isConnected()) { + const message = this.workspaceState.messageQueue[0] + if (message) { + try { + this.workspaceState.webSocketClient.send(message) + this.workspaceState.messageQueue.shift() + } catch (error) { + this.logging.error(`Error sending message: ${error}`) + } } } - } - }, this.MESSAGE_PUBLISH_INTERVAL) + }, this.MESSAGE_PUBLISH_INTERVAL) + } // Perform a one-time checkRemoteWorkspaceStatusAndReact first // Pass skipUploads as true since it would be handled by processNewWorkspaceFolders await this.checkRemoteWorkspaceStatusAndReact(true) // Set up continuous monitoring which periodically invokes checkRemoteWorkspaceStatusAndReact - if (!this.isOptedOut) { + if (!this.isOptedOut && this.continuousMonitorInterval === undefined) { this.logging.log(`Starting continuous monitor for workspace [${this.workspaceIdentifier}]`) - const intervalId = setInterval(async () => { + this.continuousMonitorInterval = setInterval(async () => { try { await this.checkRemoteWorkspaceStatusAndReact() } catch (error) { this.logging.error(`Error monitoring workspace status: ${error}`) } }, this.CONTINUOUS_MONITOR_INTERVAL) - this.continuousMonitorInterval = intervalId } } @@ -353,8 +373,8 @@ export class WorkspaceFolderManager { if (optOut) { this.logging.log(`User opted out during initial connection`) this.isOptedOut = true - await this.clearAllWorkspaceResources() - await this.startOptOutMonitor() + this.clearAllWorkspaceResources() + this.startOptOutMonitor() return resolve(false) } @@ -407,8 +427,8 @@ export class WorkspaceFolderManager { if (optOut) { this.logging.log('User opted out, clearing all resources and starting opt-out monitor') this.isOptedOut = true - await this.clearAllWorkspaceResources() - await this.startOptOutMonitor() + this.clearAllWorkspaceResources() + this.startOptOutMonitor() return } @@ -426,8 +446,7 @@ export class WorkspaceFolderManager { this.workspaceState.remoteWorkspaceState = metadata.workspaceStatus if (this.workspaceState.workspaceId === undefined) { - this.workspaceState.workspaceId = metadata.workspaceId - this.remoteWorkspaceIdResolver(this.workspaceState.workspaceId) + this.setRemoteWorkspaceId(metadata.workspaceId) } switch (metadata.workspaceStatus) { @@ -454,37 +473,24 @@ export class WorkspaceFolderManager { } } - async waitForRemoteWorkspaceId(): Promise { - // If workspaceId is already set, return it immediately - if (this.workspaceState.workspaceId) { - return this.workspaceState.workspaceId - } - // Otherwise, wait for the promise to resolve - let waitedWorkspaceId = undefined - while (!waitedWorkspaceId) { - waitedWorkspaceId = await this.remoteWorkspaceIdPromise - } - return waitedWorkspaceId + private setRemoteWorkspaceId(workspaceId: string) { + this.workspaceState.workspaceId = workspaceId + this.remoteWorkspaceIdResolver(true) + this.remoteWorkspaceIdPromiseResolved = true } private resetRemoteWorkspaceId() { this.workspaceState.workspaceId = undefined - // Store the old resolver - const oldResolver = this.remoteWorkspaceIdResolver - - // Create new promise first - this.remoteWorkspaceIdPromise = new Promise(resolve => { - this.remoteWorkspaceIdResolver = resolve - }) - - // Reset the old promise - if (oldResolver) { - oldResolver('') + if (this.remoteWorkspaceIdPromiseResolved) { + this.remoteWorkspaceIdPromise = new Promise(resolve => { + this.remoteWorkspaceIdResolver = resolve + }) + this.remoteWorkspaceIdPromiseResolved = false } } - private async startOptOutMonitor() { + private startOptOutMonitor() { if (this.optOutMonitorInterval === undefined) { const intervalId = setInterval(async () => { try { @@ -492,10 +498,17 @@ export class WorkspaceFolderManager { if (!optOut) { this.isOptedOut = false - this.logging.log('User opted back in, stopping opt-out monitor and re-initializing workspace') + this.logging.log( + "User's administrator opted in, stopping opt-out monitor and initializing remote workspace" + ) clearInterval(intervalId) this.optOutMonitorInterval = undefined - await this.initializeWorkspaceStatusMonitor() + this.initializeWorkspaceStatusMonitor().catch(error => { + this.logging.error(`Error while initializing workspace status monitoring: ${error}`) + }) + this.processNewWorkspaceFolders(this.workspaceFolders).catch(error => { + this.logging.error(`Error while processing workspace folders: ${error}`) + }) } } catch (error) { this.logging.error(`Error in opt-out monitor: ${error}`) @@ -555,16 +568,23 @@ export class WorkspaceFolderManager { } private stopContinuousMonitoring() { - this.logging.log(`Stopping monitoring for workspace [${this.workspaceIdentifier}]`) if (this.continuousMonitorInterval) { + this.logging.log(`Stopping monitoring for workspace [${this.workspaceIdentifier}]`) clearInterval(this.continuousMonitorInterval) this.continuousMonitorInterval = undefined } } + private stopOptOutMonitoring() { + if (this.optOutMonitorInterval) { + clearInterval(this.optOutMonitorInterval) + this.optOutMonitorInterval = undefined + } + } + private stopMessageQueueConsumer() { - this.logging.log(`Stopping message queue consumer`) if (this.messageQueueConsumerInterval) { + this.logging.log(`Stopping message queue consumer`) clearInterval(this.messageQueueConsumerInterval) this.messageQueueConsumerInterval = undefined } @@ -580,8 +600,7 @@ export class WorkspaceFolderManager { this.workspaceState.remoteWorkspaceState = workspaceDetails.workspace.workspaceStatus if (this.workspaceState.workspaceId === undefined) { - this.workspaceState.workspaceId = workspaceDetails.workspace.workspaceId - this.remoteWorkspaceIdResolver(this.workspaceState.workspaceId) + this.setRemoteWorkspaceId(workspaceDetails.workspace.workspaceId) } return createWorkspaceResult @@ -640,14 +659,19 @@ export class WorkspaceFolderManager { } } - // TODO, this function is unused at the moment - private async deleteWorkspace(workspaceId: string) { + public async deleteRemoteWorkspace() { + const workspaceId = this.workspaceState.workspaceId + this.resetRemoteWorkspaceId() try { + if (!workspaceId) { + this.logging.warn(`No remote workspaceId found, skipping workspace deletion`) + return + } if (isLoggedInUsingBearerToken(this.credentialsProvider)) { await this.serviceManager.getCodewhispererService().deleteWorkspace({ workspaceId: workspaceId, }) - this.logging.log(`Workspace (${workspaceId}) deleted successfully`) + this.logging.log(`Remote workspace (${workspaceId}) deleted successfully`) } else { this.logging.log(`Skipping workspace (${workspaceId}) deletion because user is not logged in`) } @@ -681,7 +705,7 @@ export class WorkspaceFolderManager { e?.__type?.includes('AccessDeniedException') && e?.reason === 'UNAUTHORIZED_WORKSPACE_CONTEXT_FEATURE_ACCESS' ) { - this.logging.log(`Server side opt-out detected for workspace context`) + this.logging.log(`User's administrator opted out server-side workspace context`) optOut = true } }