From 2158ffea040ba5afe98b9daa0555f2b11f644804 Mon Sep 17 00:00:00 2001 From: Luo Tim Date: Fri, 8 Dec 2023 22:49:54 +0800 Subject: [PATCH] Use coroutines to make prompt responsive --- build.gradle.kts | 1 + .../kotlin/ai/devchat/cli/DevChatWrapper.kt | 91 +++++++++++++------ src/main/kotlin/ai/devchat/common/Log.kt | 5 +- .../handler/SendMessageRequestHandler.kt | 4 +- .../ai/devchat/idea/DevChatToolWindow.kt | 14 ++- .../kotlin/ai/devchat/idea/JSJavaBridge.kt | 8 +- 6 files changed, 81 insertions(+), 42 deletions(-) diff --git a/build.gradle.kts b/build.gradle.kts index baf92ad..a29d42a 100644 --- a/build.gradle.kts +++ b/build.gradle.kts @@ -16,6 +16,7 @@ dependencies { implementation("com.alibaba:fastjson:2.0.42") implementation("com.fasterxml.jackson.dataformat:jackson-dataformat-yaml:2.16.0") implementation(kotlin("stdlib-jdk8")) + implementation("org.jetbrains.kotlinx:kotlinx-coroutines-core:1.7.3") } // Configure Gradle Changelog Plugin - read more: https://github.com/JetBrains/gradle-changelog-plugin diff --git a/src/main/kotlin/ai/devchat/cli/DevChatWrapper.kt b/src/main/kotlin/ai/devchat/cli/DevChatWrapper.kt index 1046059..e0db8e3 100644 --- a/src/main/kotlin/ai/devchat/cli/DevChatWrapper.kt +++ b/src/main/kotlin/ai/devchat/cli/DevChatWrapper.kt @@ -6,11 +6,38 @@ import ai.devchat.common.Settings import com.alibaba.fastjson.JSON import com.alibaba.fastjson.JSONArray import com.intellij.util.containers.addIfNotNull -import java.io.BufferedReader +import kotlinx.coroutines.* import java.io.IOException private const val DEFAULT_LOG_MAX_COUNT = 10000 + +private suspend fun Process.await( + onOutput: (String) -> Unit, + onError: (String) -> Unit +): Int = coroutineScope { + launch(Dispatchers.IO) { + inputStream.bufferedReader().forEachLine { onOutput(it) } + errorStream.bufferedReader().forEachLine { onError(it) } + } + val processExitCode = this@await.waitFor() + processExitCode +} + +suspend fun executeCommand( + command: List, + env: Map, + onOutputLine: (String) -> Unit, + onErrorLine: (String) -> Unit +): Int { + val processBuilder = ProcessBuilder(command) + env.forEach { (key, value) -> processBuilder.environment()[key] = value} + val process = withContext(Dispatchers.IO) { + processBuilder.start() + } + return process.await(onOutputLine, onErrorLine) +} + class DevChatWrapper( private val command: String = DevChatPathUtil.devchatBinPath, private var apiBase: String? = null, @@ -26,10 +53,8 @@ class DevChatWrapper( } } - private fun execCommand(commands: List, callback: ((String) -> Unit)?): String? { - val pb = ProcessBuilder(commands) - val env = pb.environment() - + private fun getEnv(): Map { + val env: MutableMap = mutableMapOf() apiBase?.let { env["OPENAI_API_BASE"] = it Log.info("api_base: $it") @@ -38,36 +63,50 @@ class DevChatWrapper( env["OPENAI_API_KEY"] = it Log.info("api_key: ${it.substring(0, 5)}...${it.substring(it.length - 4)}") } + return env + } + + private fun execCommand(commands: List): String { + Log.info("Executing command: ${commands.joinToString(" ")}}") return try { - Log.info("Executing command: ${commands.joinToString(" ")}}") - val process = pb.start() - val text = process.inputStream.bufferedReader().use { reader -> - callback?.let { - reader.forEachLine(it) - "" - } ?: reader.readText() + val outputLines: MutableList = mutableListOf() + val errorLines: MutableList = mutableListOf() + val exitCode = runBlocking { + executeCommand(commands, getEnv(), outputLines::add, errorLines::add) } - val errors = process.errorStream.bufferedReader().use(BufferedReader::readText) - process.waitFor() - val exitCode = process.exitValue() + val errors = errorLines.joinToString("\n") if (exitCode != 0) { - Log.error("Failed to execute command: $commands Exit Code: $exitCode Error: $errors") - throw RuntimeException( - "Failed to execute command: $commands Exit Code: $exitCode Error: $errors" - ) + throw RuntimeException("Command failure with exit Code: $exitCode, Errors: $errors") } else { - text + outputLines.joinToString("\n") } } catch (e: IOException) { - Log.error("Failed to execute command: $commands") - throw RuntimeException("Failed to execute command: $commands", e) - } catch (e: InterruptedException) { - Log.error("Failed to execute command: $commands") + Log.error("Failed to execute command: $commands, Exception: $e") throw RuntimeException("Failed to execute command: $commands", e) } } + private fun execCommandAsync( + commands: List, + onOutput: (String) -> Unit, + onError: (String) -> Unit = Log::error + ): Job { + Log.info("Executing command: ${commands.joinToString(" ")}}") + val exceptionHandler = CoroutineExceptionHandler { _, exception -> + Log.error("Failed to execute command: $commands, Exception: $exception") + throw RuntimeException("Failed to execute command: $commands", exception) + } + val cmdScope = CoroutineScope(SupervisorJob() + Dispatchers.Default) + + return cmdScope.launch(exceptionHandler) { + val exitCode = executeCommand(commands, getEnv(), onOutput, onError) + if (exitCode != 0) { + throw RuntimeException("Command failure with exit Code: $exitCode") + } + } + } + val prompt: (MutableList>, String, ((String) -> Unit)?) -> Unit get() = { flags: MutableList>, message: String, callback: ((String) -> Unit)? -> flags.addAll(listOf("model" to currentModel, "" to message)) @@ -98,7 +137,7 @@ class DevChatWrapper( cmd.addIfNotNull(value) } return try { - execCommand(cmd, callback) + callback?.let { execCommandAsync(cmd, callback); "" } ?: execCommand(cmd) } catch (e: Exception) { Log.error("Failed to run command $cmd: ${e.message}") throw RuntimeException("Failed to run command $cmd", e) @@ -114,7 +153,7 @@ class DevChatWrapper( cmd.addIfNotNull(value) } try { - execCommand(cmd, callback) + callback?.let { execCommandAsync(cmd, callback); "" } ?: execCommand(cmd) } catch (e: Exception) { Log.error("Failed to run command $cmd: ${e.message}") throw RuntimeException("Failed to run command $cmd", e) diff --git a/src/main/kotlin/ai/devchat/common/Log.kt b/src/main/kotlin/ai/devchat/common/Log.kt index 9e3e453..46e667c 100644 --- a/src/main/kotlin/ai/devchat/common/Log.kt +++ b/src/main/kotlin/ai/devchat/common/Log.kt @@ -1,13 +1,10 @@ package ai.devchat.common -import ai.devchat.cli.DevChatInstallationManager import com.intellij.openapi.diagnostic.LogLevel import com.intellij.openapi.diagnostic.Logger object Log { - private val LOG = Logger.getInstance( - DevChatInstallationManager::class.java - ) + private val LOG = Logger.getInstance("DevChat") private const val PREFIX = "[DevChat] " private fun setLevel(level: LogLevel) { LOG.setLevel(level) diff --git a/src/main/kotlin/ai/devchat/devchat/handler/SendMessageRequestHandler.kt b/src/main/kotlin/ai/devchat/devchat/handler/SendMessageRequestHandler.kt index a40ce23..8db0f9f 100644 --- a/src/main/kotlin/ai/devchat/devchat/handler/SendMessageRequestHandler.kt +++ b/src/main/kotlin/ai/devchat/devchat/handler/SendMessageRequestHandler.kt @@ -54,6 +54,7 @@ class SendMessageRequestHandler(metadata: JSONObject?, payload: JSONObject?) : B response.update(line) promptCallback(response) } + /* TODO: update messages cache with new one val currentTopic = ActiveConversation.topic ?: response.promptHash!! val newMessage = wrapper.logTopic(currentTopic, 1).getJSONObject(0) @@ -62,6 +63,7 @@ class SendMessageRequestHandler(metadata: JSONObject?, payload: JSONObject?) : B } else { ActiveConversation.reset(currentTopic, listOf(newMessage)) } + */ } override fun except(exception: Exception) { @@ -106,7 +108,7 @@ class SendMessageRequestHandler(metadata: JSONObject?, payload: JSONObject?) : B // Loop through the command names and check if message starts with it for (command in commandNames) { if (message.startsWith("/$command ")) { - if (message.length > command!!.length + 2) { + if (message.length > command.length + 2) { message = message.substring(command.length + 2) // +2 to take into account the '/' and the space ' ' } runResult = wrapper.runCommand(listOf(command), null) diff --git a/src/main/kotlin/ai/devchat/idea/DevChatToolWindow.kt b/src/main/kotlin/ai/devchat/idea/DevChatToolWindow.kt index ec48f86..34a07aa 100644 --- a/src/main/kotlin/ai/devchat/idea/DevChatToolWindow.kt +++ b/src/main/kotlin/ai/devchat/idea/DevChatToolWindow.kt @@ -1,8 +1,6 @@ package ai.devchat.idea -import ai.devchat.common.Log.error -import ai.devchat.common.Log.info -import ai.devchat.common.Log.setLevelInfo +import ai.devchat.common.Log import ai.devchat.devchat.DevChatActionHandler.Companion.instance import com.intellij.openapi.editor.colors.EditorColors import com.intellij.openapi.editor.colors.EditorColorsManager @@ -40,12 +38,12 @@ internal class DevChatToolWindowContent(project: Project) { private val project: Project init { - setLevelInfo() + Log.setLevelInfo() this.project = project content = JPanel(BorderLayout()) // Check if JCEF is supported if (!JBCefApp.isSupported()) { - error("JCEF is not supported.") + Log.error("JCEF is not supported.") content.add(JLabel("JCEF is not supported", SwingConstants.CENTER)) // TODO: 'return' is not allowed here // return @@ -56,17 +54,17 @@ internal class DevChatToolWindowContent(project: Project) { // Read static files var htmlContent = readStaticFile("/static/main.html") if (htmlContent!!.isEmpty()) { - error("main.html is missing.") + Log.error("main.html is missing.") htmlContent = "

Error: main.html is missing.

" } var jsContent = readStaticFile("/static/main.js") if (jsContent!!.isEmpty()) { - error("main.js is missing.") + Log.error("main.js is missing.") jsContent = "console.log('Error: main.js not found')" } val HtmlWithCssContent = insertCSSToHTML(htmlContent) val HtmlWithJsContent = insertJStoHTML(HtmlWithCssContent, jsContent) - info("main.html and main.js are loaded.") + Log.info("main.html and main.js are loaded.") // enable dev tools val myDevTools = jbCefBrowser.cefBrowser.devTools diff --git a/src/main/kotlin/ai/devchat/idea/JSJavaBridge.kt b/src/main/kotlin/ai/devchat/idea/JSJavaBridge.kt index 01c0557..988b657 100644 --- a/src/main/kotlin/ai/devchat/idea/JSJavaBridge.kt +++ b/src/main/kotlin/ai/devchat/idea/JSJavaBridge.kt @@ -3,6 +3,7 @@ package ai.devchat.idea import ai.devchat.common.Log.info import ai.devchat.devchat.ActionHandlerFactory import com.alibaba.fastjson.JSON +import com.intellij.openapi.application.ApplicationManager import com.intellij.ui.jcef.JBCefBrowser import com.intellij.ui.jcef.JBCefBrowserBase import com.intellij.ui.jcef.JBCefJSQuery @@ -12,10 +13,9 @@ import org.cef.handler.CefLoadHandlerAdapter import org.cef.network.CefRequest class JSJavaBridge(private val jbCefBrowser: JBCefBrowser) { - private val jsQuery: JBCefJSQuery + private val jsQuery: JBCefJSQuery = JBCefJSQuery.create((jbCefBrowser as JBCefBrowserBase)) init { - jsQuery = JBCefJSQuery.create((jbCefBrowser as JBCefBrowserBase)) jsQuery.addHandler { arg: String -> callJava(arg) } } @@ -33,7 +33,9 @@ class JSJavaBridge(private val jbCefBrowser: JBCefBrowser) { val payload = jsonObject.getJSONObject("payload") info("Got action: $action") val actionHandler = ActionHandlerFactory().createActionHandler(action, metadata, payload) - actionHandler.executeAction() + ApplicationManager.getApplication().invokeLater { + actionHandler.executeAction() + } return JBCefJSQuery.Response("ignore me") }