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

plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLspService.kt

@@ -91,8 +91,6 @@
 import java.io.OutputStreamWriter
 import java.io.PipedInputStream
 import java.io.PipedOutputStream
-import java.io.PrintWriter
-import java.io.StringWriter
 import java.net.Proxy
 import java.net.URI
 import java.nio.charset.StandardCharsets
@@ -286,11 +284,32 @@
         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
@@ -299,10 +318,10 @@
 
         @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) {
@@ -479,10 +498,17 @@
         }
             .wrapMessages { consumer ->
                 MessageConsumer { message ->
+                    // logging
+                    val traceLogger = LOG.atLevel(if (isDeveloperMode()) Level.INFO else Level.DEBUG)
+                    val direction = if (consumer is RemoteEndpoint) "Sent" else "Received"
+                    traceLogger.log { "$direction: $message" }
+
                     if (message is ResponseMessage && message.result is AwsExtendedInitializeResult) {
                         val result = message.result as AwsExtendedInitializeResult
                         AwsServerCapabilitiesProvider.getInstance(project).setAwsServerCapabilities(result.getAwsServerCapabilities())
                     }
+
+                    // required
                     consumer?.consume(message)
                 }
             }
@@ -495,18 +521,7 @@
                 // otherwise Gson treats all numbers as double which causes deser issues
                 it.setObjectToNumberStrategy(ToNumberPolicy.LONG_OR_DOUBLE)
                 it.registerTypeAdapterFactory(AmazonQLspTypeAdapterFactory())
-            }.traceMessages(
-                PrintWriter(
-                    object : StringWriter() {
-                        private val traceLogger = LOG.atLevel(if (isDeveloperMode()) Level.INFO else Level.DEBUG)
-
-                        override fun flush() {
-                            traceLogger.log { buffer.toString() }
-                            buffer.setLength(0)
-                        }
-                    }
-                )
-            )
+            }
             .setInput(inputWrapper.inputStream)
             .setOutput(launcherHandler.process.outputStream)
             .create()
@@ -544,13 +559,13 @@
                 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)
                 }
             }

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

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 {

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>())