add request cancellation and fix phantom completion text

src/main/kotlin/com/smallcloud/refactai/io/AsyncConnection.kt

@@ -87,9 +87,10 @@ class AsyncConnection : Disposable {
             requestProperties: Map<String, String>? = null,
             stat: UsageStatistic = UsageStatistic(),
             dataReceiveEnded: (String) -> Unit = {},
-            dataReceived: (String) -> Unit = {},
+            dataReceived: (String, String) -> Unit =  { _: String, _: String -> },
             errorDataReceived: (JsonObject) -> Unit = {},
             failedDataReceiveEnded: (Throwable?) -> Unit = {},
+            requestId: String = ""
     ): CompletableFuture<Future<*>> {
         val requestProducer: AsyncRequestProducer = BasicRequestProducer(
                 BasicRequestBuilder
@@ -106,20 +107,22 @@ class AsyncConnection : Disposable {
                 requestProducer, uri,
                 stat = stat,
                 dataReceiveEnded = dataReceiveEnded, dataReceived = dataReceived,
-                errorDataReceived = errorDataReceived, failedDataReceiveEnded = failedDataReceiveEnded
+                errorDataReceived = errorDataReceived, failedDataReceiveEnded = failedDataReceiveEnded,
+                requestId=requestId
         )
     }
 
     fun post(
-            uri: URI,
-            body: String? = null,
-            headers: Map<String, String>? = null,
-            requestProperties: Map<String, String>? = null,
-            stat: UsageStatistic,
-            dataReceiveEnded: (String) -> Unit = {},
-            dataReceived: (String) -> Unit = {},
-            errorDataReceived: (JsonObject) -> Unit = {},
-            failedDataReceiveEnded: (Throwable?) -> Unit = {},
+        uri: URI,
+        body: String? = null,
+        headers: Map<String, String>? = null,
+        requestProperties: Map<String, String>? = null,
+        stat: UsageStatistic,
+        dataReceiveEnded: (String) -> Unit = {},
+        dataReceived: (String, String) -> Unit = { _: String, _: String -> },
+        errorDataReceived: (JsonObject) -> Unit = {},
+        failedDataReceiveEnded: (Throwable?) -> Unit = {},
+        requestId: String = ""
     ): CompletableFuture<Future<*>> {
         val requestProducer: AsyncRequestProducer = BasicRequestProducer(
                 BasicRequestBuilder
@@ -135,7 +138,8 @@ class AsyncConnection : Disposable {
                 requestProducer, uri,
                 stat = stat,
                 dataReceiveEnded = dataReceiveEnded, dataReceived = dataReceived,
-                errorDataReceived = errorDataReceived, failedDataReceiveEnded = failedDataReceiveEnded
+                errorDataReceived = errorDataReceived, failedDataReceiveEnded = failedDataReceiveEnded,
+                requestId=requestId
         )
     }
 
@@ -144,9 +148,10 @@ class AsyncConnection : Disposable {
             uri: URI,
             stat: UsageStatistic,
             dataReceiveEnded: (String) -> Unit,
-            dataReceived: (String) -> Unit,
+            dataReceived: (String, String) -> Unit,
             errorDataReceived: (JsonObject) -> Unit,
             failedDataReceiveEnded: (Throwable?) -> Unit = {},
+            requestId: String = ""
     ): CompletableFuture<Future<*>> {
         return CompletableFuture.supplyAsync {
             return@supplyAsync client.execute(
@@ -179,7 +184,7 @@ class AsyncConnection : Disposable {
                                 } catch (_: JsonSyntaxException) {
                                 }
                                 val (dataPieces, maybeLeftOverBuffer) = lookForCompletedDataInStreamingBuf(bufferStr)
-                                dataPieces.forEach { dataReceived(it) }
+                                dataPieces.forEach { dataReceived(it, requestId) }
                                 if (maybeLeftOverBuffer == null) {
                                     return
                                 } else {

src/main/kotlin/com/smallcloud/refactai/io/RequestHelpers.kt

@@ -61,22 +61,24 @@ fun streamedInferenceFetch(
             uri, body, headers,
             stat = request.stat,
             dataReceiveEnded = dataReceiveEnded,
-            dataReceived = {
-                val rawJson = gson.fromJson(it, JsonObject::class.java)
+            dataReceived = {body: String, reqId: String ->
+                val rawJson = gson.fromJson(body, JsonObject::class.java)
                 if (rawJson.has("metering_balance")) {
                     AccountManager.instance.meteringBalance = rawJson.get("metering_balance").asInt
                 }
 
-                val json = gson.fromJson(it, SMCStreamingPeace::class.java)
+                val json = gson.fromJson(body, SMCStreamingPeace::class.java)
                 InferenceGlobalContext.lastAutoModel = json.model
+                json.requestId = reqId
                 UsageStats.addStatistic(true, request.stat, request.uri.toString(), "")
                 dataReceived(json)
             },
             errorDataReceived = {
                 lookForCommonErrors(it, request)?.let { message ->
                     throw SMCExceptions(message)
                 }
-            }
+            },
+            requestId=request.id
     )
 
     return job

src/main/kotlin/com/smallcloud/refactai/lsp/LSPProcessHolder.kt

@@ -224,7 +224,6 @@ class LSPProcessHolder: Disposable {
         var res = LSPCapabilities()
         InferenceGlobalContext.connection.get(url.resolve("/v1/caps"),
                 dataReceiveEnded = {},
-                dataReceived = {},
                 errorDataReceived = {}).also {
             var requestFuture: ComplexFuture<*>? = null
             try {

src/main/kotlin/com/smallcloud/refactai/modes/completion/CompletionMode.kt

@@ -34,14 +34,15 @@ import com.smallcloud.refactai.privacy.PrivacyService.Companion.instance as Priv
 import com.smallcloud.refactai.statistic.UsageStats.Companion.instance as UsageStats
 
 private val specialSymbolsRegex = "^[:\\s\\t\\n\\r(){},.\"'\\];]*\$".toRegex()
-
 class CompletionMode(
     override var needToRender: Boolean = true
 ) : Mode, CaretListener {
     private val scope: String = "completion"
     private val app = ApplicationManager.getApplication()
     private val scheduler = AppExecutorUtil.createBoundedScheduledExecutorService("SMCCompletionScheduler", 1)
     private var processTask: Future<*>? = null
+    private var lastRequestId: String = ""
+    private var requestTask: Future<*>? = null
     private var completionLayout: AsyncCompletionLayout? = null
     private val logger = Logger.getInstance("StreamedCompletionMode")
     private var completionInProgress: Boolean = false
@@ -72,7 +73,6 @@ class CompletionMode(
 
         val isMultiline = currentLine.all { it == ' ' || it == '\t' }
 
-
         if (!event.force) {
             val docEvent = event.event ?: return
             if (docEvent.offset + docEvent.newLength > editor.document.text.length) return
@@ -202,12 +202,23 @@ class CompletionMode(
             request.body.parameters.maxNewTokens = 50
             request.body.noCache = true
         }
+        if (requestTask != null && requestTask!!.isDone && requestTask!!.isCancelled ) {
+            requestTask!!.cancel(true)
+        }
+        lastRequestId = request.id
         streamedInferenceFetch(request, dataReceiveEnded = {
             InferenceGlobalContext.status = ConnectionStatus.CONNECTED
             InferenceGlobalContext.lastErrorMsg = null
         }) { prediction ->
             val choice = prediction.choices.first()
-            if ((!completionInProgress) || (choice.delta.isEmpty() && !choice.finishReason.isNullOrEmpty())) {
+            if (lastRequestId != prediction.requestId) {
+                completionLayout?.dispose()
+                completionLayout = null
+                return@streamedInferenceFetch
+            }
+
+            if ((!completionInProgress)
+                || (choice.delta.isEmpty() && !choice.finishReason.isNullOrEmpty())) {
                 return@streamedInferenceFetch
             }
             val completion: Completion = if (completionLayout?.lastCompletionData == null ||
@@ -226,26 +237,23 @@ class CompletionMode(
                 completion.snippetTelemetryId = prediction.snippetTelemetryId
             }
             completion.updateCompletion(choice.delta)
-
             synchronized(this) {
                 renderCompletion(
                     editorState.editor, editorState, completion, !prediction.cached
                 )
             }
         }?.also {
-            var requestFuture: Future<*>? = null
             try {
-                requestFuture = it.get()
-                requestFuture.get()
+                requestTask = it.get()
+                requestTask!!.get()
                 logger.debug("Completion request finished")
                 completionInProgress = false
             } catch (e: InterruptedException) {
-                handleInterruptedException(requestFuture, editorState.editor)
+                handleInterruptedException(requestTask, editorState.editor)
             } catch (e: InterruptedIOException) {
-                handleInterruptedException(requestFuture, editorState.editor)
+                handleInterruptedException(requestTask, editorState.editor)
             } catch (e: ExecutionException) {
                 cancelOrClose()
-                requestFuture?.cancel(true)
                 catchNetExceptions(e.cause)
             } catch (e: Exception) {
                 InferenceGlobalContext.status = ConnectionStatus.ERROR

src/main/kotlin/com/smallcloud/refactai/panes/gptchat/ChatGPTProvider.kt

@@ -122,7 +122,7 @@ class ChatGPTProvider : ActionListener {
             InferenceGlobalContext.connection.post(req.uri,
                     reqStr,
                     mapOf("Authorization" to "Bearer ${req.token}"),
-                    dataReceived = { response ->
+                    dataReceived = { response, _ ->
                         fun parse(response: String?): String? {
                             val gson = Gson()
                             val obj = gson.fromJson(response, JsonObject::class.java)

src/main/kotlin/com/smallcloud/refactai/struct/SMCPrediction.kt

@@ -1,13 +1,16 @@
 package com.smallcloud.refactai.struct
 
+import com.google.gson.annotations.Expose
 import com.google.gson.annotations.SerializedName
 
 data class SMCStreamingPeace(
     val choices: List<StreamingChoice>,
     val created: Double,
     val model: String,
     @SerializedName("snippet_telemetry_id") val snippetTelemetryId: Int? = null,
-    val cached: Boolean = false
+    val cached: Boolean = false,
+    @Expose
+    var requestId: String = ""
 )
 
 

src/main/kotlin/com/smallcloud/refactai/struct/SMCRequest.kt

@@ -2,8 +2,9 @@ package com.smallcloud.refactai.struct
 
 import com.google.gson.annotations.SerializedName
 import com.smallcloud.refactai.statistic.UsageStatistic
+import org.apache.commons.lang.RandomStringUtils
 import java.net.URI
-
+private fun uuid() = RandomStringUtils.randomAlphanumeric(8)
 
 data class SMCCursor(
     val file: String = "",
@@ -30,8 +31,9 @@ data class SMCRequestBody(
 )
 
 data class SMCRequest(
-        var uri: URI,
-        var body: SMCRequestBody,
-        var token: String,
-        var stat: UsageStatistic = UsageStatistic(),
+    var uri: URI,
+    var body: SMCRequestBody,
+    var token: String,
+    var id: String = uuid(),
+    var stat: UsageStatistic = UsageStatistic(),
 )