Merge pull request #186 from devchat-ai/fix-edt-errors

Fix edt errors
devchat-ai · Jul 3, 2024 · 1a1537d · 1a1537d
2 parents 0104d90 + 3f41c1b
commit 1a1537d
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 13 deletions.
diff --git a/src/main/kotlin/ai/devchat/core/handlers/NewSrcFile.kt b/src/main/kotlin/ai/devchat/core/handlers/NewSrcFile.kt
@@ -6,6 +6,7 @@ import ai.devchat.plugin.currentProject
 import com.alibaba.fastjson.JSONObject
 import com.intellij.lang.Language
 import com.intellij.openapi.application.ApplicationManager
+import com.intellij.openapi.application.runInEdt
 import com.intellij.openapi.fileEditor.FileEditorManager
 import com.intellij.psi.PsiFile
 import com.intellij.psi.PsiFileFactory
@@ -20,9 +21,9 @@ class NewSrcFile(requestAction: String, metadata: JSONObject?, payload: JSONObje
     override fun action() {
         val content = payload!!.getString("content")
         val language = payload!!.getString("language")
-        ApplicationManager.getApplication().invokeLater {
-            val project = currentProject ?: return@invokeLater
-            val dir = FileEditorManager.getInstance(project).selectedEditor?.file?.parent ?: return@invokeLater
+        runInEdt {
+            val project = currentProject ?: return@runInEdt
+            val dir = FileEditorManager.getInstance(project).selectedEditor?.file?.parent ?: return@runInEdt
             ApplicationManager.getApplication().runWriteAction {
                 val psiDirectory = PsiManager.getInstance(project).findDirectory(dir) ?: return@runWriteAction
                 val (fileLanguage, ext) = getLanguageByName(language) ?: return@runWriteAction

diff --git a/src/main/kotlin/ai/devchat/plugin/IDEServer.kt b/src/main/kotlin/ai/devchat/plugin/IDEServer.kt
@@ -48,6 +48,7 @@ import kotlinx.serialization.Serializable
 import java.awt.Point
 import java.io.File
 import java.net.ServerSocket
+import java.util.concurrent.FutureTask
 import kotlin.reflect.full.memberFunctions
 
 
@@ -365,21 +366,36 @@ fun getAvailablePort(startPort: Int): Int {
     }
 }
 
-fun Project.getPsiFile(filePath: String): PsiFile = ReadAction.compute<PsiFile, Throwable> {
-    val virtualFile = LocalFileSystem.getInstance().findFileByIoFile(File(filePath))
-    PsiManager.getInstance(this).findFile(virtualFile!!)
+fun <T> runInEdtAndGet(block: () -> T): T {
+    val app = ApplicationManager.getApplication()
+    return if (app.isDispatchThread) { block() } else {
+        val future = FutureTask(block)
+        app.invokeAndWait { future.run() }
+        future.get()
+    }
+}
+
+fun Project.getPsiFile(filePath: String): PsiFile = runInEdtAndGet {
+    ReadAction.compute<PsiFile, Throwable> {
+        val virtualFile = LocalFileSystem.getInstance().findFileByIoFile(File(filePath))
+        PsiManager.getInstance(this).findFile(virtualFile!!)
+    }
 }
 
-fun Project.getDocument(filePath: String): Document = ReadAction.compute<Document, Throwable> {
-    LocalFileSystem.getInstance().findFileByIoFile(File(filePath))?.let {
-        FileDocumentManager.getInstance().getDocument(it)
+fun Project.getDocument(filePath: String): Document = runInEdtAndGet {
+    ReadAction.compute<Document, Throwable> {
+        LocalFileSystem.getInstance().findFileByIoFile(File(filePath))?.let {
+            FileDocumentManager.getInstance().getDocument(it)
+        }
     }
 }
 
-fun Project.getCurrentFile(): VirtualFile = ReadAction.compute<VirtualFile, Throwable> {
-    val editor: Editor? = FileEditorManager.getInstance(this).selectedTextEditor
-    editor?.document?.let { document ->
-        FileDocumentManager.getInstance().getFile(document)
+fun Project.getCurrentFile(): VirtualFile = runInEdtAndGet {
+    ReadAction.compute<VirtualFile, Throwable> {
+        val editor: Editor? = FileEditorManager.getInstance(this).selectedTextEditor
+        editor?.document?.let { document ->
+            FileDocumentManager.getInstance().getFile(document)
+        }
     }
 }