fix(amazonq): fix issue where AmazonQLspService subservices are never disposed (#5753)

rli · web-flow · commit 7cb2b07a8156 · 2025-05-27T11:30:20.000-07:00
diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLspService.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLspService.kt
@@ -476,10 +476,18 @@ private class AmazonQServerInstance(private val project: Project, private val cs
             }
 
             this@AmazonQServerInstance.apply {
-                DefaultAuthCredentialsService(project, encryptionManager, this)
-                TextDocumentServiceHandler(project, this)
-                WorkspaceServiceHandler(project, lspInitResult, this)
-                DefaultModuleDependenciesService(project, this)
+                DefaultAuthCredentialsService(project, encryptionManager).also {
+                    Disposer.register(this, it)
+                }
+                TextDocumentServiceHandler(project).also {
+                    Disposer.register(this, it)
+                }
+                WorkspaceServiceHandler(project, lspInitResult).also {
+                    Disposer.register(this, it)
+                }
+                DefaultModuleDependenciesService(project).also {
+                    Disposer.register(this, it)
+                }
             }
         }
     }
diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/auth/DefaultAuthCredentialsService.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/auth/DefaultAuthCredentialsService.kt
@@ -38,7 +38,6 @@ import java.util.concurrent.TimeUnit
 class DefaultAuthCredentialsService(
     private val project: Project,
     private val encryptionManager: JwtEncryptionManager,
-    serverInstance: Disposable,
 ) : AuthCredentialsService,
     BearerTokenProviderListener,
     ToolkitConnectionManagerListener,
@@ -50,7 +49,7 @@ class DefaultAuthCredentialsService(
     private val tokenSyncIntervalMinutes = 5L
 
     init {
-        project.messageBus.connect(serverInstance).apply {
+        project.messageBus.connect(this).apply {
             subscribe(BearerTokenProviderListener.TOPIC, this@DefaultAuthCredentialsService)
             subscribe(ToolkitConnectionManagerListener.TOPIC, this@DefaultAuthCredentialsService)
             subscribe(QRegionProfileSelectedListener.TOPIC, this@DefaultAuthCredentialsService)
diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/dependencies/DefaultModuleDependenciesService.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/dependencies/DefaultModuleDependenciesService.kt
@@ -15,12 +15,11 @@ import software.aws.toolkits.jetbrains.utils.pluginAwareExecuteOnPooledThread
 
 class DefaultModuleDependenciesService(
     private val project: Project,
-    serverInstance: Disposable,
 ) : ModuleDependenciesService,
-    ModuleRootListener {
-
+    ModuleRootListener,
+    Disposable {
     init {
-        project.messageBus.connect(serverInstance).subscribe(
+        project.messageBus.connect(this).subscribe(
             ModuleRootListener.TOPIC,
             this
         )
@@ -52,4 +51,7 @@ class DefaultModuleDependenciesService(
             }
         }
     }
+
+    override fun dispose() {
+    }
 }
diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/textdocument/TextDocumentServiceHandler.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/textdocument/TextDocumentServiceHandler.kt
@@ -32,27 +32,26 @@ import software.aws.toolkits.jetbrains.utils.pluginAwareExecuteOnPooledThread
 
 class TextDocumentServiceHandler(
     private val project: Project,
-    private val serverInstance: Disposable,
 ) : FileDocumentManagerListener,
     FileEditorManagerListener,
     BulkFileListener,
-    DocumentListener {
-
+    DocumentListener,
+    Disposable {
     init {
         // didOpen & didClose events
-        project.messageBus.connect(serverInstance).subscribe(
+        project.messageBus.connect(this).subscribe(
             FileEditorManagerListener.FILE_EDITOR_MANAGER,
             this
         )
 
         // didChange events
-        project.messageBus.connect(serverInstance).subscribe(
+        project.messageBus.connect(this).subscribe(
             VirtualFileManager.VFS_CHANGES,
             this
         )
 
         // didSave events
-        project.messageBus.connect(serverInstance).subscribe(
+        project.messageBus.connect(this).subscribe(
             FileDocumentManagerListener.TOPIC,
             this
         )
@@ -72,7 +71,7 @@ class TextDocumentServiceHandler(
                         realTimeEdit(event)
                     }
                 },
-                serverInstance
+                this
             )
         }
         AmazonQLspService.executeIfRunning(project) { languageServer ->
@@ -183,4 +182,7 @@ class TextDocumentServiceHandler(
         }
         // Process document changes here
     }
+
+    override fun dispose() {
+    }
 }
diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/workspace/WorkspaceServiceHandler.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/workspace/WorkspaceServiceHandler.kt
@@ -45,22 +45,22 @@ import java.nio.file.Paths
 class WorkspaceServiceHandler(
     private val project: Project,
     initializeResult: InitializeResult,
-    serverInstance: Disposable,
 ) : BulkFileListener,
-    ModuleRootListener {
+    ModuleRootListener,
+    Disposable {
 
     private var lastSnapshot: List<WorkspaceFolder> = emptyList()
     private val operationMatchers: MutableMap<FileOperationType, List<Pair<PathMatcher, String>>> = mutableMapOf()
 
     init {
         operationMatchers.putAll(initializePatterns(initializeResult))
 
-        project.messageBus.connect(serverInstance).subscribe(
+        project.messageBus.connect(this).subscribe(
             VirtualFileManager.VFS_CHANGES,
             this
         )
 
-        project.messageBus.connect(serverInstance).subscribe(
+        project.messageBus.connect(this).subscribe(
             ModuleRootListener.TOPIC,
             this
         )
@@ -313,4 +313,7 @@ class WorkspaceServiceHandler(
             lastSnapshot = currentSnapshot
         }
     }
+
+    override fun dispose() {
+    }
 }
diff --git a/plugins/amazonq/shared/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/lsp/auth/DefaultAuthCredentialsServiceTest.kt b/plugins/amazonq/shared/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/lsp/auth/DefaultAuthCredentialsServiceTest.kt
@@ -143,7 +143,7 @@ class DefaultAuthCredentialsServiceTest {
 
     @Test
     fun `activeConnectionChanged updates token when connection ID matches Q connection`() {
-        sut = DefaultAuthCredentialsService(project, mockEncryptionManager, mockk())
+        sut = DefaultAuthCredentialsService(project, mockEncryptionManager)
         val newConnection = createMockConnection("new-token", "connection-id")
         every { mockConnection.id } returns "connection-id"
 
@@ -154,7 +154,7 @@ class DefaultAuthCredentialsServiceTest {
 
     @Test
     fun `activeConnectionChanged does not update token when connection ID differs`() {
-        sut = DefaultAuthCredentialsService(project, mockEncryptionManager, mockk())
+        sut = DefaultAuthCredentialsService(project, mockEncryptionManager)
         val newConnection = createMockConnection("new-token", "different-id")
         every { mockConnection.id } returns "q-connection-id"
 
@@ -165,7 +165,7 @@ class DefaultAuthCredentialsServiceTest {
 
     @Test
     fun `onChange updates token with new connection`() {
-        sut = DefaultAuthCredentialsService(project, mockEncryptionManager, mockk())
+        sut = DefaultAuthCredentialsService(project, mockEncryptionManager)
         setupMockConnectionManager("updated-token")
 
         sut.onChange("providerId", listOf("new-scope"))
@@ -178,7 +178,7 @@ class DefaultAuthCredentialsServiceTest {
         every { isQConnected(project) } returns false
         every { isQExpired(project) } returns false
 
-        sut = DefaultAuthCredentialsService(project, mockEncryptionManager, mockk())
+        sut = DefaultAuthCredentialsService(project, mockEncryptionManager)
 
         verify(exactly = 0) { mockLanguageServer.updateTokenCredentials(any()) }
     }
@@ -188,15 +188,15 @@ class DefaultAuthCredentialsServiceTest {
         every { isQConnected(project) } returns true
         every { isQExpired(project) } returns true
 
-        sut = DefaultAuthCredentialsService(project, mockEncryptionManager, mockk())
+        sut = DefaultAuthCredentialsService(project, mockEncryptionManager)
 
         verify(exactly = 0) { mockLanguageServer.updateTokenCredentials(any()) }
     }
 
     @Test
     fun `test updateTokenCredentials unencrypted success`() {
         val isEncrypted = false
-        sut = DefaultAuthCredentialsService(project, mockEncryptionManager, mockk())
+        sut = DefaultAuthCredentialsService(project, mockEncryptionManager)
 
         sut.updateTokenCredentials(mockConnection, isEncrypted)
 
@@ -215,7 +215,7 @@ class DefaultAuthCredentialsServiceTest {
 
     @Test
     fun `test updateTokenCredentials encrypted success`() {
-        sut = DefaultAuthCredentialsService(project, mockEncryptionManager, mockk())
+        sut = DefaultAuthCredentialsService(project, mockEncryptionManager)
 
         val encryptedToken = "encryptedToken"
         val isEncrypted = true
@@ -239,7 +239,7 @@ class DefaultAuthCredentialsServiceTest {
 
     @Test
     fun `test deleteTokenCredentials success`() {
-        sut = DefaultAuthCredentialsService(project, mockEncryptionManager, mockk())
+        sut = DefaultAuthCredentialsService(project, mockEncryptionManager)
 
         every { mockLanguageServer.deleteTokenCredentials() } returns CompletableFuture.completedFuture(Unit)
 
@@ -252,7 +252,7 @@ class DefaultAuthCredentialsServiceTest {
     fun `init results in token update`() {
         every { isQConnected(any()) } returns true
         every { isQExpired(any()) } returns false
-        sut = DefaultAuthCredentialsService(project, mockEncryptionManager, mockk())
+        sut = DefaultAuthCredentialsService(project, mockEncryptionManager)
 
         verify(exactly = 1) { mockLanguageServer.updateTokenCredentials(any()) }
     }
diff --git a/plugins/amazonq/shared/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/lsp/dependencies/DefaultModuleDependenciesServiceTest.kt b/plugins/amazonq/shared/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/lsp/dependencies/DefaultModuleDependenciesServiceTest.kt
@@ -95,7 +95,7 @@ class DefaultModuleDependenciesServiceTest {
         every { mockModuleManager.modules } returns arrayOf(module)
         prepDependencyProvider(listOf(Pair(module, params)))
 
-        sut = DefaultModuleDependenciesService(project, mockk())
+        sut = DefaultModuleDependenciesService(project)
 
         verify { mockLanguageServer.didChangeDependencyPaths(params) }
     }
@@ -127,7 +127,7 @@ class DefaultModuleDependenciesServiceTest {
             )
         )
 
-        sut = DefaultModuleDependenciesService(project, mockk())
+        sut = DefaultModuleDependenciesService(project)
 
         verify { mockLanguageServer.didChangeDependencyPaths(params1) }
         verify { mockLanguageServer.didChangeDependencyPaths(params2) }
@@ -148,7 +148,7 @@ class DefaultModuleDependenciesServiceTest {
         val event = mockk<ModuleRootEvent>()
         every { event.isCausedByFileTypesChange } returns true
 
-        sut = DefaultModuleDependenciesService(project, mockk())
+        sut = DefaultModuleDependenciesService(project)
 
         sut.rootsChanged(event)
 
@@ -173,7 +173,7 @@ class DefaultModuleDependenciesServiceTest {
 
         prepDependencyProvider(listOf(Pair(module, params)))
 
-        sut = DefaultModuleDependenciesService(project, mockk())
+        sut = DefaultModuleDependenciesService(project)
 
         sut.rootsChanged(event)
 
diff --git a/plugins/amazonq/shared/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/lsp/textdocument/TextDocumentServiceHandlerTest.kt b/plugins/amazonq/shared/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/lsp/textdocument/TextDocumentServiceHandlerTest.kt
@@ -94,7 +94,7 @@ class TextDocumentServiceHandlerTest {
         every { mockTextDocumentService.didOpen(any()) } returns Unit
         every { mockTextDocumentService.didClose(any()) } returns Unit
 
-        sut = TextDocumentServiceHandler(projectRule.project, mockk())
+        sut = TextDocumentServiceHandler(projectRule.project)
     }
 
     @Test
@@ -137,7 +137,7 @@ class TextDocumentServiceHandlerTest {
             projectRule.fixture.createFile("name", content).also { projectRule.fixture.openFileInEditor(it) }
         }
 
-        sut = TextDocumentServiceHandler(projectRule.project, mockk())
+        sut = TextDocumentServiceHandler(projectRule.project)
 
         val paramsSlot = mutableListOf<DidOpenTextDocumentParams>()
         verify { mockTextDocumentService.didOpen(capture(paramsSlot)) }
diff --git a/plugins/amazonq/shared/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/lsp/workspace/WorkspaceServiceHandlerTest.kt b/plugins/amazonq/shared/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/lsp/workspace/WorkspaceServiceHandlerTest.kt
@@ -146,7 +146,7 @@ class WorkspaceServiceHandlerTest {
         every { mockInitializeResult.capabilities } returns mockCapabilities
 
         // Create WorkspaceServiceHandler with mocked InitializeResult
-        sut = WorkspaceServiceHandler(project, mockInitializeResult, mockk())
+        sut = WorkspaceServiceHandler(project, mockInitializeResult)
     }
 
     @Test

Original file line number	Diff line number	Diff line change
`@@ -476,10 +476,18 @@ private class AmazonQServerInstance(private val project: Project, private val cs`
`476`	`476`	`}`
`477`	`477`
`478`	`478`	`this@AmazonQServerInstance.apply {`
`479`		`- DefaultAuthCredentialsService(project, encryptionManager, this)`
`480`		`- TextDocumentServiceHandler(project, this)`
`481`		`- WorkspaceServiceHandler(project, lspInitResult, this)`
`482`		`- DefaultModuleDependenciesService(project, this)`
	`479`	`+ DefaultAuthCredentialsService(project, encryptionManager).also {`
	`480`	`+ Disposer.register(this, it)`
	`481`	`+ }`
	`482`	`+ TextDocumentServiceHandler(project).also {`
	`483`	`+ Disposer.register(this, it)`
	`484`	`+ }`
	`485`	`+ WorkspaceServiceHandler(project, lspInitResult).also {`
	`486`	`+ Disposer.register(this, it)`
	`487`	`+ }`
	`488`	`+ DefaultModuleDependenciesService(project).also {`
	`489`	`+ Disposer.register(this, it)`
	`490`	`+ }`
`483`	`491`	`}`
`484`	`492`	`}`
`485`	`493`	`}`
Original file line number	Diff line number	Diff line change
`@@ -15,12 +15,11 @@ import software.aws.toolkits.jetbrains.utils.pluginAwareExecuteOnPooledThread`
`15`	`15`
`16`	`16`	`class DefaultModuleDependenciesService(`
`17`	`17`	`private val project: Project,`
`18`		`- serverInstance: Disposable,`
`19`	`18`	`) : ModuleDependenciesService,`
`20`		`- ModuleRootListener {`
`21`		`-`
	`19`	`+ ModuleRootListener,`
	`20`	`+ Disposable {`
`22`	`21`	`init {`
`23`		`- project.messageBus.connect(serverInstance).subscribe(`
	`22`	`+ project.messageBus.connect(this).subscribe(`
`24`	`23`	`ModuleRootListener.TOPIC,`
`25`	`24`	`this`
`26`	`25`	`)`
`@@ -52,4 +51,7 @@ class DefaultModuleDependenciesService(`
`52`	`51`	`}`
`53`	`52`	`}`
`54`	`53`	`}`
	`54`	`+`
	`55`	`+ override fun dispose() {`
	`56`	`+ }`
`55`	`57`	`}`
Original file line number	Diff line number	Diff line change
`@@ -32,27 +32,26 @@ import software.aws.toolkits.jetbrains.utils.pluginAwareExecuteOnPooledThread`
`32`	`32`
`33`	`33`	`class TextDocumentServiceHandler(`
`34`	`34`	`private val project: Project,`
`35`		`- private val serverInstance: Disposable,`
`36`	`35`	`) : FileDocumentManagerListener,`
`37`	`36`	`FileEditorManagerListener,`
`38`	`37`	`BulkFileListener,`
`39`		`- DocumentListener {`
`40`		`-`
	`38`	`+ DocumentListener,`
	`39`	`+ Disposable {`
`41`	`40`	`init {`
`42`	`41`	`// didOpen & didClose events`
`43`		`- project.messageBus.connect(serverInstance).subscribe(`
	`42`	`+ project.messageBus.connect(this).subscribe(`
`44`	`43`	`FileEditorManagerListener.FILE_EDITOR_MANAGER,`
`45`	`44`	`this`
`46`	`45`	`)`
`47`	`46`
`48`	`47`	`// didChange events`
`49`		`- project.messageBus.connect(serverInstance).subscribe(`
	`48`	`+ project.messageBus.connect(this).subscribe(`
`50`	`49`	`VirtualFileManager.VFS_CHANGES,`
`51`	`50`	`this`
`52`	`51`	`)`
`53`	`52`
`54`	`53`	`// didSave events`
`55`		`- project.messageBus.connect(serverInstance).subscribe(`
	`54`	`+ project.messageBus.connect(this).subscribe(`
`56`	`55`	`FileDocumentManagerListener.TOPIC,`
`57`	`56`	`this`
`58`	`57`	`)`
`@@ -72,7 +71,7 @@ class TextDocumentServiceHandler(`
`72`	`71`	`realTimeEdit(event)`
`73`	`72`	`}`
`74`	`73`	`},`
`75`		`- serverInstance`
	`74`	`+ this`
`76`	`75`	`)`
`77`	`76`	`}`
`78`	`77`	`AmazonQLspService.executeIfRunning(project) { languageServer ->`
`@@ -183,4 +182,7 @@ class TextDocumentServiceHandler(`
`183`	`182`	`}`
`184`	`183`	`// Process document changes here`
`185`	`184`	`}`
	`185`	`+`
	`186`	`+ override fun dispose() {`
	`187`	`+ }`
`186`	`188`	`}`
Original file line number	Diff line number	Diff line change
`@@ -94,7 +94,7 @@ class TextDocumentServiceHandlerTest {`
`94`	`94`	`every { mockTextDocumentService.didOpen(any()) } returns Unit`
`95`	`95`	`every { mockTextDocumentService.didClose(any()) } returns Unit`
`96`	`96`
`97`		`- sut = TextDocumentServiceHandler(projectRule.project, mockk())`
	`97`	`+ sut = TextDocumentServiceHandler(projectRule.project)`
`98`	`98`	`}`
`99`	`99`
`100`	`100`	`@Test`
`@@ -137,7 +137,7 @@ class TextDocumentServiceHandlerTest {`
`137`	`137`	`projectRule.fixture.createFile("name", content).also { projectRule.fixture.openFileInEditor(it) }`
`138`	`138`	`}`
`139`	`139`
`140`		`- sut = TextDocumentServiceHandler(projectRule.project, mockk())`
	`140`	`+ sut = TextDocumentServiceHandler(projectRule.project)`
`141`	`141`
`142`	`142`	`val paramsSlot = mutableListOf<DidOpenTextDocumentParams>()`
`143`	`143`	`verify { mockTextDocumentService.didOpen(capture(paramsSlot)) }`
Original file line number	Diff line number	Diff line change
`@@ -146,7 +146,7 @@ class WorkspaceServiceHandlerTest {`
`146`	`146`	`every { mockInitializeResult.capabilities } returns mockCapabilities`
`147`	`147`
`148`	`148`	`// Create WorkspaceServiceHandler with mocked InitializeResult`
`149`		`- sut = WorkspaceServiceHandler(project, mockInitializeResult, mockk())`
	`149`	`+ sut = WorkspaceServiceHandler(project, mockInitializeResult)`
`150`	`150`	`}`
`151`	`151`
`152`	`152`	`@Test`