fix(amazonq): move workspace lsp messages off blocking thread calls

rli · rli · commit 4a127de12dbc · 2025-06-11T17:49:14.000-07:00
we are inappropriately freezing EDT while passing workspace context to the LSP
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
@@ -282,11 +282,32 @@ class AmazonQLspService(private val project: Project, private val cs: CoroutineS
         return runnable(lsp)
     }
 
+    suspend fun<T> executeIfRunning(runnable: suspend AmazonQLspService.(AmazonQLanguageServer) -> T): T? {
+        val lsp = try {
+            withTimeout(1.seconds) {
+                val holder = mutex.withLock { instance }.await()
+                holder.initializeResult.join()
+
+                holder.languageServer
+            }
+        } catch (_: Exception) {
+            LOG.debug { "LSP not running" }
+            null
+        }
+
+        return lsp?.let { runnable(it) }
+    }
+
     fun<T> executeSync(runnable: suspend AmazonQLspService.(AmazonQLanguageServer) -> T): T =
         runBlocking(cs.coroutineContext) {
             execute(runnable)
         }
 
+    fun<T> syncExecuteIfRunning(runnable: suspend AmazonQLspService.(AmazonQLanguageServer) -> T): T? =
+        runBlocking(cs.coroutineContext) {
+            executeIfRunning(runnable)
+        }
+
     companion object {
         private val LOG = getLogger<AmazonQLspService>()
         private const val MAX_RESTARTS = 5
@@ -295,10 +316,10 @@ class AmazonQLspService(private val project: Project, private val cs: CoroutineS
 
         @Deprecated("Easy to accidentally freeze EDT")
         fun <T> executeIfRunning(project: Project, runnable: AmazonQLspService.(AmazonQLanguageServer) -> T): T? =
-            project.serviceIfCreated<AmazonQLspService>()?.executeSync(runnable)
+            project.serviceIfCreated<AmazonQLspService>()?.syncExecuteIfRunning(runnable)
 
-        suspend fun <T> asyncExecuteIfRunning(project: Project, runnable: suspend AmazonQLspService.(AmazonQLanguageServer) -> T): T? =
-            project.serviceIfCreated<AmazonQLspService>()?.execute(runnable)
+        suspend fun <T> executeAsyncIfRunning(project: Project, runnable: suspend AmazonQLspService.(AmazonQLanguageServer) -> T): T? =
+            project.serviceIfCreated<AmazonQLspService>()?.executeIfRunning(runnable)
 
         fun didChangeConfiguration(project: Project) {
             executeIfRunning(project) {
@@ -496,13 +517,13 @@ private class AmazonQServerInstance(private val project: Project, private val cs
                 DefaultAuthCredentialsService(project, encryptionManager).also {
                     Disposer.register(this, it)
                 }
-                TextDocumentServiceHandler(project).also {
+                TextDocumentServiceHandler(project, cs).also {
                     Disposer.register(this, it)
                 }
-                WorkspaceServiceHandler(project, lspInitResult).also {
+                WorkspaceServiceHandler(project, cs, lspInitResult).also {
                     Disposer.register(this, it)
                 }
-                DefaultModuleDependenciesService(project).also {
+                DefaultModuleDependenciesService(project, cs).also {
                     Disposer.register(this, it)
                 }
             }
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
@@ -8,13 +8,15 @@ import com.intellij.openapi.module.ModuleManager
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.roots.ModuleRootEvent
 import com.intellij.openapi.roots.ModuleRootListener
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.launch
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLspService
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.dependencies.ModuleDependencyProvider.Companion.EP_NAME
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.dependencies.DidChangeDependencyPathsParams
-import software.aws.toolkits.jetbrains.utils.pluginAwareExecuteOnPooledThread
 
 class DefaultModuleDependenciesService(
     private val project: Project,
+    private val cs: CoroutineScope,
 ) : ModuleDependenciesService,
     ModuleRootListener,
     Disposable {
@@ -35,7 +37,7 @@ class DefaultModuleDependenciesService(
 
     override fun didChangeDependencyPaths(params: DidChangeDependencyPathsParams) {
         AmazonQLspService.executeIfRunning(project) { languageServer ->
-            pluginAwareExecuteOnPooledThread {
+            cs.launch {
                 languageServer.didChangeDependencyPaths(params)
             }
         }
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
@@ -19,6 +19,8 @@ import com.intellij.openapi.vfs.VirtualFileManager
 import com.intellij.openapi.vfs.newvfs.BulkFileListener
 import com.intellij.openapi.vfs.newvfs.events.VFileContentChangeEvent
 import com.intellij.openapi.vfs.newvfs.events.VFileEvent
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.launch
 import org.eclipse.lsp4j.DidChangeTextDocumentParams
 import org.eclipse.lsp4j.DidCloseTextDocumentParams
 import org.eclipse.lsp4j.DidOpenTextDocumentParams
@@ -29,10 +31,10 @@ import org.eclipse.lsp4j.TextDocumentItem
 import org.eclipse.lsp4j.VersionedTextDocumentIdentifier
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLspService
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.util.LspEditorUtil.toUriString
-import software.aws.toolkits.jetbrains.utils.pluginAwareExecuteOnPooledThread
 
 class TextDocumentServiceHandler(
     private val project: Project,
+    private val cs: CoroutineScope,
 ) : FileDocumentManagerListener,
     FileEditorManagerListener,
     BulkFileListener,
@@ -78,7 +80,7 @@ class TextDocumentServiceHandler(
         }
         AmazonQLspService.executeIfRunning(project) { languageServer ->
             toUriString(file)?.let { uri ->
-                pluginAwareExecuteOnPooledThread {
+                cs.launch {
                     languageServer.textDocumentService.didOpen(
                         DidOpenTextDocumentParams().apply {
                             textDocument = TextDocumentItem().apply {
@@ -98,7 +100,7 @@ class TextDocumentServiceHandler(
         AmazonQLspService.executeIfRunning(project) { languageServer ->
             val file = FileDocumentManager.getInstance().getFile(document) ?: return@executeIfRunning
             toUriString(file)?.let { uri ->
-                pluginAwareExecuteOnPooledThread {
+                cs.launch {
                     languageServer.textDocumentService.didSave(
                         DidSaveTextDocumentParams().apply {
                             textDocument = TextDocumentIdentifier().apply {
@@ -114,7 +116,7 @@ class TextDocumentServiceHandler(
 
     override fun after(events: MutableList<out VFileEvent>) {
         AmazonQLspService.executeIfRunning(project) { languageServer ->
-            pluginAwareExecuteOnPooledThread {
+            cs.launch {
                 events.filterIsInstance<VFileContentChangeEvent>().forEach { event ->
                     val document = FileDocumentManager.getInstance().getCachedDocument(event.file) ?: return@forEach
                     toUriString(event.file)?.let { uri ->
@@ -168,8 +170,8 @@ class TextDocumentServiceHandler(
 
     private fun realTimeEdit(event: DocumentEvent) {
         AmazonQLspService.executeIfRunning(project) { languageServer ->
-            pluginAwareExecuteOnPooledThread {
-                val vFile = FileDocumentManager.getInstance().getFile(event.document) ?: return@pluginAwareExecuteOnPooledThread
+            cs.launch {
+                val vFile = FileDocumentManager.getInstance().getFile(event.document) ?: return@launch
                 toUriString(vFile)?.let { uri ->
                     languageServer.textDocumentService.didChange(
                         DidChangeTextDocumentParams().apply {
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
@@ -16,6 +16,8 @@ import com.intellij.openapi.vfs.newvfs.events.VFileDeleteEvent
 import com.intellij.openapi.vfs.newvfs.events.VFileEvent
 import com.intellij.openapi.vfs.newvfs.events.VFileMoveEvent
 import com.intellij.openapi.vfs.newvfs.events.VFilePropertyChangeEvent
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.launch
 import org.eclipse.lsp4j.CreateFilesParams
 import org.eclipse.lsp4j.DeleteFilesParams
 import org.eclipse.lsp4j.DidChangeWatchedFilesParams
@@ -37,13 +39,13 @@ import org.eclipse.lsp4j.WorkspaceFoldersChangeEvent
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLspService
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.util.LspEditorUtil.toUriString
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.util.WorkspaceFolderUtil.createWorkspaceFolders
-import software.aws.toolkits.jetbrains.utils.pluginAwareExecuteOnPooledThread
 import java.nio.file.FileSystems
 import java.nio.file.PathMatcher
 import java.nio.file.Paths
 
 class WorkspaceServiceHandler(
     private val project: Project,
+    private val cs: CoroutineScope,
     initializeResult: InitializeResult,
 ) : BulkFileListener,
     ModuleRootListener,
@@ -281,7 +283,7 @@ class WorkspaceServiceHandler(
 
     override fun after(events: List<VFileEvent>) {
         // since we are using synchronous FileListener
-        pluginAwareExecuteOnPooledThread {
+        cs.launch {
             didCreateFiles(events.filter { it is VFileCreateEvent || it is VFileMoveEvent || it is VFileCopyEvent })
             didDeleteFiles(events.filter { it is VFileMoveEvent || it is VFileDeleteEvent })
             didRenameFiles(events.filterIsInstance<VFilePropertyChangeEvent>())
