fix(amazonq): properly deposit workspace context server resources on exit (#1647)

* fix(amazonq): cancel workspace context server initialization workflow on exit

* fix(amazonq): cancel message queue consumer interval on exit

* fix(amazonq): add check before start workspace context server init workflow

---------

Co-authored-by: Jiatong Li <[email protected]>

Author: LiGaCu

server/aws-lsp-codewhisperer/src/language-server/workspaceContext/workspaceContextServer.test.ts

@@ -7,14 +7,18 @@ import { AmazonQTokenServiceManager } from '../../shared/amazonQServiceManager/A
 describe('WorkspaceContext Server', () => {
     let features: TestFeatures
     let server: Server
+    let disposeServer: () => void
 
     before(() => {
         features = new TestFeatures()
         server = WorkspaceContextServer()
+        disposeServer = server(features)
     })
 
     afterEach(() => {
         sinon.restore()
+        disposeServer()
+        features.dispose()
     })
 
     describe('Initialization', () => {
@@ -35,10 +39,6 @@ describe('WorkspaceContext Server', () => {
                 },
             } as InitializeParams)
 
-            // Create a stub for the static getInstance method
-            const getInstanceStub = sinon.stub(AmazonQTokenServiceManager, 'getInstance')
-            getInstanceStub.throws(new Error('Deliberate error to exit the test at onInitialized'))
-
             await features.initialize(server)
 
             // Verify that a warning was logged (indicating the workspaceIdentifier was generated)

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

@@ -26,6 +26,7 @@ export const WorkspaceContextServer = (): Server => features => {
     let artifactManager: ArtifactManager
     let dependencyDiscoverer: DependencyDiscoverer
     let workspaceFolderManager: WorkspaceFolderManager
+    let workflowInitializationInterval: NodeJS.Timeout
     let isWorkflowInitialized: boolean = false
     let isOptedIn: boolean = false
     let abTestingEvaluated = false
@@ -236,7 +237,10 @@ export const WorkspaceContextServer = (): Server => features => {
              * of workspace folders is updated using *artifactManager.updateWorkspaceFolders(workspaceFolders)* before
              * initializing again.
              */
-            setInterval(async () => {
+            if (workflowInitializationInterval) {
+                return
+            }
+            workflowInitializationInterval = setInterval(async () => {
                 if (!isOptedIn) {
                     return
                 }
@@ -501,10 +505,13 @@ export const WorkspaceContextServer = (): Server => features => {
     logging.log('Workspace context server has been initialized')
 
     return () => {
-        workspaceFolderManager.clearAllWorkspaceResources().catch(error => {
-            logging.warn(
-                `Error while clearing workspace resources: ${error instanceof Error ? error.message : 'Unknown error'}`
-            )
-        })
+        clearInterval(workflowInitializationInterval)
+        if (workspaceFolderManager) {
+            workspaceFolderManager.clearAllWorkspaceResources().catch(error => {
+                logging.warn(
+                    `Error while clearing workspace resources: ${error instanceof Error ? error.message : 'Unknown error'}`
+                )
+            })
+        }
     }
 }

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

@@ -114,20 +114,6 @@ export class WorkspaceFolderManager {
             remoteWorkspaceState: 'CREATION_PENDING',
             messageQueue: [],
         }
-
-        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)
     }
 
     /**
@@ -229,6 +215,7 @@ export class WorkspaceFolderManager {
 
     async clearAllWorkspaceResources() {
         this.stopContinuousMonitoring()
+        this.stopMessageQueueConsumer()
         this.resetRemoteWorkspaceId()
         this.workspaceState.webSocketClient?.destroyClient()
         this.dependencyDiscoverer.dispose()
@@ -333,6 +320,22 @@ export class WorkspaceFolderManager {
 
     async initializeWorkspaceStatusMonitor() {
         this.logging.log(`Initializing workspace status check for workspace [${this.workspaceIdentifier}]`)
+
+        // 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}`)
+                    }
+                }
+            }
+        }, 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)
@@ -578,6 +581,14 @@ export class WorkspaceFolderManager {
         }
     }
 
+    private stopMessageQueueConsumer() {
+        this.logging.log(`Stopping message queue consumer`)
+        if (this.messageQueueConsumerInterval) {
+            clearInterval(this.messageQueueConsumerInterval)
+            this.messageQueueConsumerInterval = undefined
+        }
+    }
+
     private async createNewWorkspace() {
         const createWorkspaceResult = await this.createWorkspace(this.workspaceIdentifier)
         const workspaceDetails = createWorkspaceResult.response