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

server/aws-lsp-codewhisperer/src/language-server/workspaceContext/artifactManager.ts

@@ -124,7 +124,12 @@ export class ArtifactManager {
         return zipFileMetadata
     }
 
-    async removeWorkspaceFolders(workspaceFolders: WorkspaceFolder[]): Promise<void> {
+    dispose(): void {
+        this.filesByWorkspaceFolderAndLanguage.clear()
+        this.workspaceFolders = []
+    }
+
+    removeWorkspaceFolders(workspaceFolders: WorkspaceFolder[]): void {
         workspaceFolders.forEach(workspaceToRemove => {
             // Find the matching workspace folder by URI
             let folderToDelete: WorkspaceFolder | undefined

server/aws-lsp-codewhisperer/src/language-server/workspaceContext/dependency/dependencyDiscoverer.ts

@@ -137,6 +137,7 @@ export class DependencyDiscoverer {
     async reSyncDependenciesToS3(folders: WorkspaceFolder[]) {
         Atomics.store(this.dependencyUploadedSizeSum, 0, 0)
         for (const dependencyHandler of this.dependencyHandlerRegistry) {
+            dependencyHandler.markAllDependenciesAsUnZipped()
             await dependencyHandler.zipDependencyMap(folders)
         }
     }

server/aws-lsp-codewhisperer/src/language-server/workspaceContext/dependency/dependencyHandler/LanguageDependencyHandler.ts

@@ -353,4 +353,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/fileUploadJobManager.ts

@@ -157,6 +157,8 @@ export class FileUploadJobManager {
     }
 
     public dispose(): void {
-        clearInterval(this.jobConsumerInterval)
+        if (this.jobConsumerInterval) {
+            clearInterval(this.jobConsumerInterval)
+        }
     }
 }

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

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
+                fileUploadJobManager?.dispose()
+                dependencyEventBundler?.dispose()
                 await workspaceFolderManager.clearAllWorkspaceResources()
+                // Delete remote workspace when user chooses to opt-out
+                await workspaceFolderManager.deleteRemoteWorkspace()
             }
         } catch (error) {
             logging.error(`Error in getConfiguration: ${error}`)
@@ -304,6 +320,8 @@ export const WorkspaceContextServer = (): Server => features => {
                     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
+                        fileUploadJobManager?.dispose()
+                        dependencyEventBundler?.dispose()
                         await workspaceFolderManager.clearAllWorkspaceResources()
                     }
                     isWorkflowInitialized = false