feat(amazonq): allow opt-out for workspace context server

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<WorkspaceFolder, Map<CodewhispererLanguage, FileMetadata[]>>
+    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<void> {
+    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

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<BaseDependencyInfo>[] = []
     private initializedWorkspaceFolder = new Map<WorkspaceFolder, boolean>()
-    // 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) {

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<T extends BaseDependencyInfo> {
     public language: CodewhispererLanguage
@@ -28,7 +33,7 @@ export abstract class LanguageDependencyHandler<T extends BaseDependencyInfo> {
     // key: workspaceFolder, value: {key: dependency name, value: Dependency}
     protected dependencyMap = new Map<WorkspaceFolder, Map<string, Dependency>>()
     protected dependencyUploadedSizeMap = new Map<WorkspaceFolder, number>()
-    protected dependencyUploadedSizeSum: Uint32Array<SharedArrayBuffer>
+    protected dependencyHandlerSharedState: DependencyHandlerSharedState
     protected dependencyWatchers: Map<string, DependencyWatcher> = new Map<string, DependencyWatcher>()
     protected artifactManager: ArtifactManager
     protected dependenciesFolderName: string
@@ -48,7 +53,7 @@ export abstract class LanguageDependencyHandler<T extends BaseDependencyInfo> {
         workspaceFolders: WorkspaceFolder[],
         artifactManager: ArtifactManager,
         dependenciesFolderName: string,
-        dependencyUploadedSizeSum: Uint32Array<SharedArrayBuffer>
+        dependencyHandlerSharedState: DependencyHandlerSharedState
     ) {
         this.language = language
         this.workspace = workspace
@@ -62,7 +67,7 @@ export abstract class LanguageDependencyHandler<T extends BaseDependencyInfo> {
         this.workspaceFolders.forEach(workSpaceFolder =>
             this.dependencyMap.set(workSpaceFolder, new Map<string, Dependency>())
         )
-        this.dependencyUploadedSizeSum = dependencyUploadedSizeSum
+        this.dependencyHandlerSharedState = dependencyHandlerSharedState
     }
 
     /*
@@ -126,6 +131,9 @@ export abstract class LanguageDependencyHandler<T extends BaseDependencyInfo> {
     async zipDependencyMap(folders: WorkspaceFolder[]): Promise<void> {
         // 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<T extends BaseDependencyInfo> {
         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<T extends BaseDependencyInfo> {
                     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<T extends BaseDependencyInfo> {
      * 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<T extends BaseDependencyInfo> {
 
     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<T extends BaseDependencyInfo> {
     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
+            })
+        })
+    }
 }

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<SharedArrayBuffer>
+        dependencyHandlerSharedState: DependencyHandlerSharedState
     ): LanguageDependencyHandler<BaseDependencyInfo> | 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

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)
+        }
     }
 }