feat: improve code completion user experience (#763)

src/main/java/ee/carlrobert/codegpt/settings/service/llama/form/InfillPromptTemplatePanel.java

@@ -17,6 +17,6 @@ public class InfillPromptTemplatePanel extends BasePromptTemplatePanel<InfillPro
 
   @Override
   protected String buildPromptDescription(InfillPromptTemplate template) {
-    return template.buildPrompt(new InfillRequest.Builder("PREFIX", "SUFFIX").build());
+    return template.buildPrompt(new InfillRequest.Builder("PREFIX", "SUFFIX", 0).build());
   }
 }

src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionCompletionEventListener.kt

@@ -1,113 +0,0 @@
-package ee.carlrobert.codegpt.codecompletions
-
-import ai.grazie.nlp.utils.takeLastWhitespaces
-import ai.grazie.nlp.utils.takeWhitespaces
-import com.intellij.notification.NotificationType
-import com.intellij.openapi.diagnostic.thisLogger
-import com.intellij.openapi.editor.Editor
-import ee.carlrobert.codegpt.CodeGPTKeys.IS_FETCHING_COMPLETION
-import ee.carlrobert.codegpt.ui.OverlayUtil.showNotification
-import ee.carlrobert.llm.client.openai.completion.ErrorDetails
-import ee.carlrobert.llm.completion.CompletionEventListener
-import okhttp3.sse.EventSource
-import java.util.concurrent.atomic.AtomicBoolean
-
-private const val MAX_LINES_TO_REQUEST = 4
-private const val MAX_LINES_TO_DISPLAY = 2
-
-abstract class CodeCompletionCompletionEventListener(
-    private val editor: Editor,
-    private val infillRequest: InfillRequest? = null,
-) : CompletionEventListener<String> {
-
-    companion object {
-        private val logger = thisLogger()
-    }
-
-    private val stringBuilder = StringBuilder()
-    private val isCancelled = AtomicBoolean(false)
-    private val isSending = AtomicBoolean(true)
-
-    open fun onComplete(fullMessage: String) {}
-    open fun onMessage(message: String) {}
-
-    override fun onOpen() {
-        setLoading(true)
-    }
-
-    override fun onMessage(message: String, eventSource: EventSource) {
-        if (isCancelled.get()) return
-
-        val processedMessage = if (infillRequest != null && stringBuilder.isEmpty()) {
-            message.tryTrimStart(infillRequest.prefix.lines())
-        } else {
-            message
-        }
-
-        val newLineCount = (stringBuilder.toString() + processedMessage).count { it == '\n' }
-        if (newLineCount >= MAX_LINES_TO_REQUEST) {
-            cancelStreaming(processedMessage, eventSource)
-            return
-        }
-
-        stringBuilder.append(processedMessage)
-
-        if (newLineCount <= MAX_LINES_TO_DISPLAY && isSending.get()) {
-            if (newLineCount == MAX_LINES_TO_DISPLAY && processedMessage.contains('\n')) {
-                isSending.set(false)
-                onMessage(processedMessage.substring(0, processedMessage.lastIndexOf('\n')))
-            } else {
-                onMessage(processedMessage)
-            }
-        }
-    }
-
-    override fun onComplete(messageBuilder: StringBuilder) {
-        setLoading(false)
-        onComplete(stringBuilder.trimEnd().toString())
-    }
-
-    override fun onCancelled(messageBuilder: StringBuilder) {
-        setLoading(false)
-        onComplete(stringBuilder.trimEnd().toString())
-    }
-
-    override fun onError(error: ErrorDetails, ex: Throwable) {
-        if (ex.message == null || (ex.message != null && ex.message != "Canceled")) {
-            showNotification(error.message, NotificationType.ERROR)
-            logger.error(error.message, ex)
-        }
-    }
-
-    private fun setLoading(loading: Boolean) {
-        IS_FETCHING_COMPLETION.set(editor, loading)
-        editor.project?.messageBus
-            ?.syncPublisher(CodeCompletionProgressNotifier.CODE_COMPLETION_PROGRESS_TOPIC)
-            ?.loading(loading)
-    }
-
-    private fun cancelStreaming(processedMessage: String, eventSource: EventSource) {
-        stringBuilder.append(processedMessage.substring(0, processedMessage.lastIndexOf('\n')))
-        isCancelled.set(true)
-        isSending.set(false)
-        eventSource.cancel()
-    }
-
-    private fun String.tryTrimStart(lines: List<String>): String {
-        val whiteSpaces = this.takeWhitespaces()
-        if (lines.size >= 2
-            && whiteSpaces.isNotEmpty()
-            && lines[lines.size - 1].trim().isEmpty()
-        ) {
-            return this.trimStart()
-        }
-
-        if (lines.isNotEmpty()) {
-            val lastLine = lines[lines.size - 1]
-            if (lastLine.takeLastWhitespaces().isNotEmpty()) {
-                return this.trimStart()
-            }
-        }
-        return this
-    }
-}

src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionEventListener.kt

@@ -0,0 +1,121 @@
+package ee.carlrobert.codegpt.codecompletions
+
+import com.intellij.notification.NotificationType
+import com.intellij.openapi.application.runReadAction
+import com.intellij.openapi.diagnostic.thisLogger
+import com.intellij.openapi.editor.Editor
+import com.intellij.openapi.util.TextRange
+import ee.carlrobert.codegpt.CodeGPTKeys.IS_FETCHING_COMPLETION
+import ee.carlrobert.codegpt.CodeGPTKeys.REMAINING_EDITOR_COMPLETION
+import ee.carlrobert.codegpt.ui.OverlayUtil.showNotification
+import ee.carlrobert.codegpt.util.StringUtil
+import ee.carlrobert.llm.client.openai.completion.ErrorDetails
+import ee.carlrobert.llm.completion.CompletionEventListener
+import okhttp3.sse.EventSource
+
+abstract class CodeCompletionEventListener(
+    private val editor: Editor
+) : CompletionEventListener<String> {
+
+    companion object {
+        private val logger = thisLogger()
+    }
+
+    private var isFirstLine = true
+    private val currentLineBuffer = StringBuilder()
+    private val incomingTextBuffer = StringBuilder()
+
+    open fun onLineReceived(completionLine: String) {}
+
+    override fun onOpen() {
+        setLoading(true)
+        REMAINING_EDITOR_COMPLETION.set(editor, "")
+    }
+
+    override fun onMessage(message: String, eventSource: EventSource) {
+        incomingTextBuffer.append(message)
+
+        while (incomingTextBuffer.contains("\n")) {
+            val lineEndIndex = incomingTextBuffer.indexOf("\n")
+            val line = incomingTextBuffer.substring(0, lineEndIndex) + '\n'
+            processCompletionLine(line)
+            incomingTextBuffer.delete(0, lineEndIndex + 1)
+        }
+    }
+
+    private fun processCompletionLine(line: String) {
+        currentLineBuffer.append(line)
+
+        if (currentLineBuffer.trim().isNotEmpty()) {
+            val completionText = if (isFirstLine) {
+                line.adjustWhitespaces().also {
+                    isFirstLine = false
+                    onLineReceived(it)
+                }
+            } else {
+                currentLineBuffer.toString()
+            }
+
+            appendRemainingCompletion(completionText)
+            currentLineBuffer.clear()
+        }
+    }
+
+    override fun onComplete(messageBuilder: StringBuilder) {
+        handleCompleted(messageBuilder)
+    }
+
+    override fun onCancelled(messageBuilder: StringBuilder) {
+        handleCompleted(messageBuilder)
+    }
+
+    override fun onError(error: ErrorDetails, ex: Throwable) {
+        if (ex.message == null || (ex.message != null && ex.message != "Canceled")) {
+            showNotification(error.message, NotificationType.ERROR)
+            logger.error(error.message, ex)
+        }
+        setLoading(false)
+    }
+
+    private fun String.adjustWhitespaces(): String {
+        val adjustedLine = runReadAction {
+            val lineNumber = editor.document.getLineNumber(editor.caretModel.offset)
+            val editorLine = editor.document.getText(
+                TextRange(
+                    editor.document.getLineStartOffset(lineNumber),
+                    editor.document.getLineEndOffset(lineNumber)
+                )
+            )
+
+            StringUtil.adjustWhitespace(this, editorLine)
+        }
+
+        return if (adjustedLine.length != this.length) adjustedLine else this
+    }
+
+    private fun handleCompleted(messageBuilder: StringBuilder) {
+        setLoading(false)
+
+        if (incomingTextBuffer.isNotEmpty()) {
+            appendRemainingCompletion(incomingTextBuffer.toString())
+        }
+
+        if (isFirstLine) {
+            val completionLine = messageBuilder.toString().adjustWhitespaces()
+            REMAINING_EDITOR_COMPLETION.set(editor, completionLine)
+            onLineReceived(completionLine)
+        }
+    }
+
+    private fun appendRemainingCompletion(text: String) {
+        val previousRemainingText = REMAINING_EDITOR_COMPLETION.get(editor) ?: ""
+        REMAINING_EDITOR_COMPLETION.set(editor, previousRemainingText + text)
+    }
+
+    private fun setLoading(loading: Boolean) {
+        IS_FETCHING_COMPLETION.set(editor, loading)
+        editor.project?.messageBus
+            ?.syncPublisher(CodeCompletionProgressNotifier.CODE_COMPLETION_PROGRESS_TOPIC)
+            ?.loading(loading)
+    }
+}

src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionInsertAction.kt

@@ -0,0 +1,113 @@
+package ee.carlrobert.codegpt.codecompletions
+
+import ai.grazie.nlp.utils.takeWhitespaces
+import com.intellij.codeInsight.hint.HintManagerImpl
+import com.intellij.codeInsight.inline.completion.InlineCompletionInsertEnvironment
+import com.intellij.codeInsight.inline.completion.session.InlineCompletionContext
+import com.intellij.codeInsight.inline.completion.session.InlineCompletionSession
+import com.intellij.codeInsight.lookup.LookupManager
+import com.intellij.openapi.actionSystem.DataContext
+import com.intellij.openapi.editor.Caret
+import com.intellij.openapi.editor.Editor
+import com.intellij.openapi.editor.actionSystem.EditorAction
+import com.intellij.openapi.editor.actionSystem.EditorWriteActionHandler
+import com.intellij.psi.PsiDocumentManager
+import com.intellij.util.concurrency.ThreadingAssertions
+import ee.carlrobert.codegpt.CodeGPTKeys.REMAINING_EDITOR_COMPLETION
+
+class CodeCompletionInsertAction :
+    EditorAction(InsertInlineCompletionHandler()), HintManagerImpl.ActionToIgnore {
+
+    class InsertInlineCompletionHandler : EditorWriteActionHandler() {
+        override fun executeWriteAction(editor: Editor, caret: Caret?, dataContext: DataContext) {
+            ThreadingAssertions.assertEventDispatchThread()
+            ThreadingAssertions.assertWriteAccess()
+
+            val session = InlineCompletionSession.getOrNull(editor) ?: return
+            val context = session.context
+            val elements = context.state.elements
+                .filter { it.element is CodeCompletionTextElement }
+                .map { it.element as CodeCompletionTextElement }
+
+            for (element in elements) {
+                val insertEnvironment = InlineCompletionInsertEnvironment(
+                    editor,
+                    session.request.file,
+                    element.textRange
+                )
+                context.copyUserDataTo(insertEnvironment)
+
+                editor.document.insertString(element.textRange.startOffset, element.text)
+
+                if (element.originalText == element.text) {
+                    processStandardCompletionElement(element, editor)
+                } else {
+                    processPartialCompletionElement(element, editor)
+                }
+
+                PsiDocumentManager.getInstance(session.request.file.project)
+                    .commitDocument(editor.document)
+
+                session.provider.insertHandler.afterInsertion(insertEnvironment, elements)
+
+                LookupManager.getActiveLookup(editor)?.hideLookup(false)
+            }
+        }
+
+        override fun isEnabledForCaret(
+            editor: Editor,
+            caret: Caret,
+            dataContext: DataContext
+        ): Boolean {
+            val completionContext = InlineCompletionContext.getOrNull(editor)
+            val element = completionContext?.state?.elements?.firstOrNull()?.element
+            if (element is CodeCompletionTextElement) {
+                return completionContext.startOffset() == (caret.offset + element.offsetDelta)
+            }
+            return completionContext?.startOffset() == caret.offset
+        }
+
+        private fun processStandardCompletionElement(
+            element: CodeCompletionTextElement,
+            editor: Editor
+        ) {
+            val endOffset = element.textRange.endOffset
+            editor.caretModel.moveToOffset(endOffset)
+
+            val remainingCompletionLine = REMAINING_EDITOR_COMPLETION.get(editor)
+                .removePrefix(element.text)
+
+            processRemainingCompletion(remainingCompletionLine, editor, endOffset)
+        }
+
+        private fun processPartialCompletionElement(element: CodeCompletionTextElement, editor: Editor) {
+            val lineNumber = editor.document.getLineNumber(editor.caretModel.offset)
+            val lineEndOffset = editor.document.getLineEndOffset(lineNumber)
+            editor.caretModel.moveToOffset(lineEndOffset)
+
+            val remainingText = REMAINING_EDITOR_COMPLETION.get(editor)
+            val remainingCompletionLine = if (element.originalText.length > remainingText.length) {
+                remainingText.removePrefix(element.text)
+            } else {
+                remainingText.removePrefix(element.originalText)
+            }
+
+            processRemainingCompletion(remainingCompletionLine, editor, lineEndOffset + 1)
+        }
+
+        private fun processRemainingCompletion(
+            remainingCompletion: String,
+            editor: Editor,
+            offset: Int
+        ) {
+            val whitespaces = remainingCompletion.takeWhitespaces()
+            if (whitespaces.isNotEmpty()) {
+                editor.document.insertString(offset, whitespaces)
+                editor.caretModel.moveToOffset(offset + whitespaces.length)
+            }
+
+            val nextCompletionLine = remainingCompletion.removePrefix(whitespaces)
+            REMAINING_EDITOR_COMPLETION.set(editor, nextCompletionLine)
+        }
+    }
+}

src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionInsertHandler.kt

@@ -5,76 +5,20 @@ import com.intellij.codeInsight.inline.completion.InlineCompletionEvent
 import com.intellij.codeInsight.inline.completion.InlineCompletionInsertEnvironment
 import com.intellij.codeInsight.inline.completion.InlineCompletionInsertHandler
 import com.intellij.codeInsight.inline.completion.elements.InlineCompletionElement
-import com.intellij.openapi.application.readAction
-import com.intellij.openapi.components.service
-import com.intellij.openapi.editor.Editor
-import com.intellij.openapi.util.TextRange
 import ee.carlrobert.codegpt.CodeGPTKeys.REMAINING_EDITOR_COMPLETION
-import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.Dispatchers
-import kotlinx.coroutines.SupervisorJob
-import kotlinx.coroutines.launch
-
-private const val NEXT_COMPLETION_LINE_COUNT_THRESHOLD = 4
 
 class CodeCompletionInsertHandler : InlineCompletionInsertHandler {
 
-    private val scope = CoroutineScope(Dispatchers.Default + SupervisorJob())
-
     override fun afterInsertion(
         environment: InlineCompletionInsertEnvironment,
         elements: List<InlineCompletionElement>
     ) {
         val editor = environment.editor
-        val appliedText = elements.joinToString("") { it.text }
-        val existingCompletion = REMAINING_EDITOR_COMPLETION.get(editor) ?: ""
-        val remainingCompletion =
-            existingCompletion.substring(appliedText.length, existingCompletion.length)
-
-        REMAINING_EDITOR_COMPLETION.set(editor, remainingCompletion)
-
+        val remainingCompletion = REMAINING_EDITOR_COMPLETION.get(editor)
         if (remainingCompletion.isNotEmpty()) {
             InlineCompletion.getHandlerOrNull(editor)?.invoke(
                 InlineCompletionEvent.DirectCall(editor, editor.caretModel.currentCaret)
             )
-
-            if (remainingCompletion.count { it == '\n' } <= NEXT_COMPLETION_LINE_COUNT_THRESHOLD) {
-                scope.launch {
-                    fetchNextCompletion(editor, remainingCompletion)
-                }
-            }
         }
     }
-
-    private suspend fun fetchNextCompletion(editor: Editor, remainingCompletion: String) {
-        val project = editor.project ?: return
-        project.service<CodeCompletionService>().getCodeCompletionAsync(
-            buildNextRequest(editor, remainingCompletion),
-            object : CodeCompletionCompletionEventListener(editor) {
-                override fun onComplete(fullMessage: String) {
-                    val nextCompletion =
-                        (REMAINING_EDITOR_COMPLETION.get(editor) ?: "") + fullMessage
-                    REMAINING_EDITOR_COMPLETION.set(editor, nextCompletion)
-                }
-            }
-        )
-    }
-
-    private suspend fun buildNextRequest(
-        editor: Editor,
-        remainingCompletion: String
-    ): InfillRequest {
-        val caretOffset = readAction { editor.caretModel.offset }
-        val prefix =
-            (editor.document.getText(TextRange(0, caretOffset)) + remainingCompletion)
-                .truncateText(MAX_PROMPT_TOKENS, false)
-        val suffix =
-            editor.document.getText(
-                TextRange(
-                    caretOffset,
-                    editor.document.textLength
-                )
-            ).truncateText(MAX_PROMPT_TOKENS)
-        return InfillRequest.Builder(prefix, suffix).build()
-    }
 }

src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionTextElement.kt

@@ -0,0 +1,37 @@
+package ee.carlrobert.codegpt.codecompletions
+
+import com.intellij.codeInsight.inline.completion.elements.InlineCompletionElement
+import com.intellij.codeInsight.inline.completion.elements.InlineCompletionGrayTextElement
+import com.intellij.openapi.editor.Editor
+import com.intellij.openapi.util.TextRange
+
+class CodeCompletionTextElement(
+    override val text: String,
+    private val insertOffset: Int,
+    val textRange: TextRange,
+    val offsetDelta: Int = 0,
+    val originalText: String = text,
+) : InlineCompletionElement {
+
+    override fun toPresentable(): InlineCompletionElement.Presentable =
+        Presentable(this, insertOffset, textRange)
+
+    open class Presentable(
+        element: InlineCompletionElement,
+        private val insertOffset: Int,
+        private val textRange: TextRange,
+    ) : InlineCompletionGrayTextElement.Presentable(element) {
+
+        override fun render(editor: Editor, offset: Int) {
+            super.render(editor, insertOffset)
+        }
+
+        override fun startOffset(): Int {
+            return textRange.startOffset
+        }
+
+        override fun endOffset(): Int {
+            return textRange.endOffset
+        }
+    }
+}

src/main/kotlin/ee/carlrobert/codegpt/codecompletions/DebouncedCodeCompletionProvider.kt

@@ -2,11 +2,11 @@ package ee.carlrobert.codegpt.codecompletions
 
 import com.intellij.codeInsight.inline.completion.*
 import com.intellij.codeInsight.inline.completion.elements.InlineCompletionElement
-import com.intellij.codeInsight.inline.completion.elements.InlineCompletionGrayTextElement
 import com.intellij.openapi.application.runInEdt
 import com.intellij.openapi.components.service
 import com.intellij.openapi.diagnostic.thisLogger
 import com.intellij.openapi.editor.Editor
+import com.intellij.openapi.util.TextRange
 import ee.carlrobert.codegpt.CodeGPTKeys.IS_FETCHING_COMPLETION
 import ee.carlrobert.codegpt.CodeGPTKeys.REMAINING_EDITOR_COMPLETION
 import ee.carlrobert.codegpt.settings.GeneralSettings
@@ -16,6 +16,7 @@ import ee.carlrobert.codegpt.settings.service.custom.CustomServiceSettings
 import ee.carlrobert.codegpt.settings.service.llama.LlamaSettings
 import ee.carlrobert.codegpt.settings.service.ollama.OllamaSettings
 import ee.carlrobert.codegpt.settings.service.openai.OpenAISettings
+import ee.carlrobert.codegpt.util.StringUtil.findCompletionParts
 import kotlinx.coroutines.channels.ProducerScope
 import kotlinx.coroutines.channels.awaitClose
 import kotlinx.coroutines.flow.channelFlow
@@ -44,6 +45,14 @@ class DebouncedCodeCompletionProvider : DebouncedInlineCompletionProvider() {
     override val providerPresentation: InlineCompletionProviderPresentation
         get() = CodeCompletionProviderPresentation()
 
+    private fun String.extractUntilNewline(): String {
+        val index = this.indexOf('\n')
+        if (index == -1) {
+            return this
+        }
+        return this.substring(0, index + 1)
+    }
+
     override suspend fun getSuggestionDebounced(request: InlineCompletionRequest): InlineCompletionSuggestion {
         val editor = request.editor
         val remainingCompletion = REMAINING_EDITOR_COMPLETION.get(editor)
@@ -51,7 +60,7 @@ class DebouncedCodeCompletionProvider : DebouncedInlineCompletionProvider() {
             && remainingCompletion != null
             && remainingCompletion.isNotEmpty()
         ) {
-            return sendNextSuggestion(remainingCompletion)
+            return sendNextSuggestion(remainingCompletion.extractUntilNewline(), request)
         }
 
         val project = editor.project
@@ -69,10 +78,12 @@ class DebouncedCodeCompletionProvider : DebouncedInlineCompletionProvider() {
                 ?.loading(true)
 
             val infillRequest = InfillRequestUtil.buildInfillRequest(request)
-            val call = project.service<CodeCompletionService>().getCodeCompletionAsync(
-                infillRequest,
-                getEventListener(request.editor, infillRequest)
-            )
+            val call = project
+                .service<CodeCompletionService>()
+                .getCodeCompletionAsync(
+                    infillRequest,
+                    getEventListener(request.editor, infillRequest)
+                )
             currentCallRef.set(call)
             awaitClose { currentCallRef.getAndSet(null)?.cancel() }
         })
@@ -109,27 +120,62 @@ class DebouncedCodeCompletionProvider : DebouncedInlineCompletionProvider() {
     private fun ProducerScope<InlineCompletionElement>.getEventListener(
         editor: Editor,
         infillRequest: InfillRequest
-    ) = object : CodeCompletionCompletionEventListener(editor, infillRequest) {
-        override fun onMessage(message: String) {
-            runInEdt {
-                trySend(InlineCompletionGrayTextElement(message))
-            }
-        }
+    ) = object : CodeCompletionEventListener(editor) {
 
-        override fun onComplete(fullMessage: String) {
-            REMAINING_EDITOR_COMPLETION.set(editor, fullMessage)
+        override fun onLineReceived(completionLine: String) {
+            runInEdt {
+                val editorLineSuffix = editor.getLineSuffixAfterCaret()
+                if (editorLineSuffix.isEmpty()) {
+                    trySend(
+                        CodeCompletionTextElement(
+                            completionLine,
+                            infillRequest.caretOffset,
+                            TextRange.from(infillRequest.caretOffset, completionLine.length),
+                        )
+                    )
+                } else {
+                    var prevStartOffset = infillRequest.caretOffset
+                    val completionParts =
+                        findCompletionParts(editorLineSuffix, completionLine.trimEnd())
+
+                    completionParts.forEach { (completionPart, offsetDelta) ->
+                        val element = CodeCompletionTextElement(
+                            completionPart,
+                            infillRequest.caretOffset + offsetDelta,
+                            TextRange.from(prevStartOffset + offsetDelta, completionPart.length),
+                            offsetDelta,
+                            completionLine
+                        )
+                        prevStartOffset += completionPart.length
+
+                        trySend(element)
+                    }
+                }
+
+            }
         }
     }
 
-    private fun sendNextSuggestion(nextCompletion: String): InlineCompletionSuggestion {
+    private fun Editor.getLineSuffixAfterCaret(): String {
+        val lineEndOffset = document.getLineEndOffset(document.getLineNumber(caretModel.offset))
+        return document.getText(TextRange(caretModel.offset, lineEndOffset))
+    }
+
+    private fun sendNextSuggestion(
+        nextCompletion: String,
+        request: InlineCompletionRequest
+    ): InlineCompletionSuggestion {
+
         return InlineCompletionSuggestion.Default(channelFlow {
             launch {
                 trySend(
-                    InlineCompletionGrayTextElement(
-                        nextCompletion.lines().take(2).joinToString("\n")
+                    CodeCompletionTextElement(
+                        nextCompletion,
+                        request.startOffset,
+                        TextRange.from(request.startOffset, nextCompletion.length),
                     )
                 )
             }
         })
     }
-}
+}

src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillRequest.kt

@@ -13,13 +13,15 @@ const val MAX_PROMPT_TOKENS = 128
 class InfillRequest private constructor(
     val prefix: String,
     val suffix: String,
+    val caretOffset: Int,
     val fileDetails: FileDetails?,
     val vcsDetails: VcsDetails?,
     val context: InfillContext?
 ) {
 
     companion object {
-        fun builder(prefix: String, suffix: String) = Builder(prefix, suffix)
+        fun builder(prefix: String, suffix: String, caretOffset: Int) =
+            Builder(prefix, suffix, caretOffset)
     }
 
     data class VcsDetails(val stagedDiff: String? = null, val unstagedDiff: String? = null)
@@ -28,13 +30,15 @@ class InfillRequest private constructor(
     class Builder {
         private val prefix: String
         private val suffix: String
+        private val caretOffset: Int
         private var fileDetails: FileDetails? = null
         private var vcsDetails: VcsDetails? = null
         private var context: InfillContext? = null
 
-        constructor(prefix: String, suffix: String) {
+        constructor(prefix: String, suffix: String, caretOffset: Int) {
             this.prefix = prefix
             this.suffix = suffix
+            this.caretOffset = caretOffset
         }
 
         constructor(document: Document, caretOffset: Int) {
@@ -48,6 +52,7 @@ class InfillRequest private constructor(
                         document.textLength
                     )
                 ).truncateText(MAX_PROMPT_TOKENS)
+            this.caretOffset = caretOffset
         }
 
         fun fileDetails(fileDetails: FileDetails) = apply { this.fileDetails = fileDetails }
@@ -55,7 +60,7 @@ class InfillRequest private constructor(
         fun context(context: InfillContext) = apply { this.context = context }
 
         fun build() =
-            InfillRequest(prefix, suffix, fileDetails, vcsDetails, context)
+            InfillRequest(prefix, suffix, caretOffset, fileDetails, vcsDetails, context)
     }
 }
 

src/main/kotlin/ee/carlrobert/codegpt/settings/service/CodeCompletionConfigurationForm.kt

@@ -65,7 +65,7 @@ class CodeCompletionConfigurationForm(
         promptTemplateHelpText.setToolTipText(null)
 
         val description = StringEscapeUtils.escapeHtml4(
-            template.buildPrompt(InfillRequest.Builder("PREFIX", "SUFFIX").build())
+            template.buildPrompt(InfillRequest.Builder("PREFIX", "SUFFIX", 0).build())
         )
         HelpTooltip()
             .setTitle(template.toString())

src/main/kotlin/ee/carlrobert/codegpt/settings/service/custom/form/CustomServiceCodeCompletionForm.kt

@@ -163,7 +163,7 @@ class CustomServiceCodeCompletionForm(
     private fun testConnection() {
         CompletionRequestService.getInstance().getCustomOpenAICompletionAsync(
             CodeCompletionRequestFactory.buildCustomRequest(
-                InfillRequest.Builder("Hello", "!").build(),
+                InfillRequest.Builder("Hello", "!", 0).build(),
                 urlField.text,
                 tabbedPane.headers,
                 tabbedPane.body,
@@ -206,7 +206,7 @@ class CustomServiceCodeCompletionForm(
 
         val description = StringEscapeUtils.escapeHtml4(
             template.buildPrompt(
-                InfillRequest.Builder("PREFIX", "SUFFIX").build(),
+                InfillRequest.Builder("PREFIX", "SUFFIX", 0).build(),
             )
         )
         HelpTooltip()

src/main/kotlin/ee/carlrobert/codegpt/util/StringUtil.kt

@@ -0,0 +1,66 @@
+package ee.carlrobert.codegpt.util
+
+import ai.grazie.nlp.utils.takeWhitespaces
+import com.intellij.util.diff.Diff
+
+object StringUtil {
+
+    fun adjustWhitespace(
+        completionLine: String,
+        editorLine: String
+    ): String {
+        val editorWhitespaces = editorLine.takeWhitespaces()
+
+        if (completionLine.isNotEmpty() && editorWhitespaces.isNotEmpty()) {
+            if (completionLine.startsWith(editorWhitespaces)) {
+                return completionLine.substring(editorWhitespaces.length)
+            }
+            if (editorLine.isBlank()) {
+                val completionWhitespaces = completionLine.takeWhitespaces()
+                return completionLine.substring(completionWhitespaces.length)
+            }
+        }
+
+        return completionLine
+    }
+
+    fun findCompletionParts(
+        editorLineSuffix: String,
+        completionLine: String
+    ): List<Pair<String, Int>> {
+        val nonOverlappingPart = findNonOverlappingPart(editorLineSuffix, completionLine)
+        if (nonOverlappingPart.length == completionLine.length) {
+            return listOf(Pair(completionLine, 0))
+        }
+
+        val result = ArrayList<Pair<String, Int>>()
+        val editorChars: IntArray = editorLineSuffix.chars().toArray()
+        val completionChars: IntArray = completionLine.chars().toArray()
+        val changes: List<Diff.Change> =
+            Diff.buildChanges(editorChars, completionChars)?.toList() ?: emptyList()
+        for (change in changes) {
+            val part = completionLine.substring(change.line1, change.line1 + change.inserted)
+            result.add(Pair(part, change.line0))
+        }
+
+        return result
+    }
+
+    private fun findNonOverlappingPart(
+        editorLineSuffix: String,
+        completionLine: String
+    ): String {
+        var i = editorLineSuffix.length - 1
+        var j = completionLine.length - 1
+        while (i >= 0 && j >= 0 && editorLineSuffix[i] == completionLine[j]) {
+            i--
+            j--
+        }
+
+        if (j >= 0) {
+            return completionLine.substring(0, j + 1)
+        }
+
+        return ""
+    }
+}

src/main/resources/META-INF/plugin.xml

@@ -84,6 +84,11 @@
     <resource-bundle>messages.codegpt</resource-bundle>
 
     <actions>
+        <action id="InsertInlineCompletionAction" class="ee.carlrobert.codegpt.codecompletions.CodeCompletionInsertAction" overrides="true">
+          <add-to-group group-id="InlineCompletion" anchor="first"/>
+          <keyboard-shortcut first-keystroke="TAB" keymap="$default"/>
+        </action>
+
         <group id="CodeGPTEditorPopup">
             <group id="action.editor.group.EditorActionGroup"
                 text="CodeGPT"

src/test/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionServiceTest.kt

@@ -2,7 +2,9 @@ package ee.carlrobert.codegpt.codecompletions
 
 import com.intellij.codeInsight.inline.completion.session.InlineCompletionSession.Companion.getOrNull
 import com.intellij.openapi.editor.VisualPosition
+import com.intellij.openapi.util.TextRange
 import com.intellij.testFramework.PlatformTestUtil
+import ee.carlrobert.codegpt.CodeGPTKeys.REMAINING_EDITOR_COMPLETION
 import ee.carlrobert.codegpt.util.file.FileUtil
 import ee.carlrobert.llm.client.http.RequestEntity
 import ee.carlrobert.llm.client.http.exchange.StreamHttpExchange
@@ -86,7 +88,249 @@ class CodeCompletionServiceTest : IntegrationTest() {
         }
     }
 
-    private fun assertInlineSuggestion(errorMessage: String, onAssert: (String) -> Boolean) {
+    fun `test apply inline suggestions without initial following text`() {
+        useCodeGPTService()
+        myFixture.configureByText(
+            "CompletionTest.java",
+            "class Node {\n  "
+        )
+        myFixture.editor.caretModel.moveToVisualPosition(VisualPosition(1, 2))
+        expectCodeGPT(StreamHttpExchange { request: RequestEntity ->
+            assertThat(request.uri.path).isEqualTo("/v1/code/completions")
+            assertThat(request.method).isEqualTo("POST")
+            assertThat(request.body)
+                .extracting("model", "prefix", "suffix", "fileExtension")
+                .containsExactly(
+                    "TEST_CODE_MODEL",
+                    "class Node {\n   ",
+                    "",
+                    "java"
+                )
+            listOf(
+                jsonMapResponse("choices", jsonArray(jsonMap("text", "\n   int data;"))),
+                jsonMapResponse("choices", jsonArray(jsonMap("text", "\n   Node lef"))),
+                jsonMapResponse("choices", jsonArray(jsonMap("text", "t;\n   Node ri"))),
+                jsonMapResponse("choices", jsonArray(jsonMap("text", "ght;\n\n   public"))),
+                jsonMapResponse("choices", jsonArray(jsonMap("text", " Node(int data"))),
+                jsonMapResponse("choices", jsonArray(jsonMap("text", ") {\n"))),
+                jsonMapResponse("choices", jsonArray(jsonMap("text", "      this.data ="))),
+                jsonMapResponse("choices", jsonArray(jsonMap("text", " data;\n   }"))),
+                jsonMapResponse("choices", jsonArray(jsonMap("text", "\n}"))),
+            )
+        })
+
+        myFixture.type(' ')
+        assertRemainingCompletion {
+            it == "int data;\n" +
+                    "   Node left;\n" +
+                    "   Node right;\n" +
+                    "\n" +
+                    "   public Node(int data) {\n" +
+                    "      this.data = data;\n" +
+                    "   }\n" +
+                    "}"
+        }
+        assertInlineSuggestion {
+            it == "int data;\n"
+        }
+        myFixture.type('\t')
+        assertRemainingCompletion {
+            it == "Node left;\n" +
+                    "   Node right;\n" +
+                    "\n" +
+                    "   public Node(int data) {\n" +
+                    "      this.data = data;\n" +
+                    "   }\n" +
+                    "}"
+        }
+        assertInlineSuggestion {
+            it == "Node left;\n"
+        }
+        assertThat(myFixture.editor.caretModel.visualPosition).isEqualTo(VisualPosition(2, 3))
+        myFixture.type('\t')
+        assertRemainingCompletion {
+            it == "Node right;\n" +
+                    "\n" +
+                    "   public Node(int data) {\n" +
+                    "      this.data = data;\n" +
+                    "   }\n" +
+                    "}"
+        }
+        assertInlineSuggestion("Failed to assert remaining completion.") {
+            it == "Node right;\n"
+        }
+        assertThat(myFixture.editor.caretModel.visualPosition).isEqualTo(VisualPosition(3, 3))
+        myFixture.type('\t')
+        assertRemainingCompletion {
+            it == "public Node(int data) {\n" +
+                    "      this.data = data;\n" +
+                    "   }\n" +
+                    "}"
+        }
+        assertInlineSuggestion {
+            it == "public Node(int data) {\n"
+        }
+        assertThat(myFixture.editor.caretModel.visualPosition).isEqualTo(VisualPosition(5, 3))
+        myFixture.type('\t')
+        assertRemainingCompletion {
+            it == "this.data = data;\n" +
+                    "   }\n" +
+                    "}"
+        }
+        assertInlineSuggestion {
+            it == "this.data = data;\n"
+        }
+        assertThat(myFixture.editor.caretModel.visualPosition).isEqualTo(VisualPosition(6, 6))
+        myFixture.type('\t')
+        assertRemainingCompletion {
+            it == "}\n" +
+                    "}"
+        }
+        assertInlineSuggestion {
+            it == "}\n"
+        }
+        assertThat(myFixture.editor.caretModel.visualPosition).isEqualTo(VisualPosition(7, 3))
+        myFixture.type('\t')
+        assertRemainingCompletion {
+            it == "}"
+        }
+        assertInlineSuggestion {
+            it == "}"
+        }
+        assertThat(myFixture.editor.caretModel.visualPosition).isEqualTo(VisualPosition(8, 0))
+        myFixture.type('\t')
+        assertRemainingCompletion {
+            it == ""
+        }
+    }
+
+    fun `test apply inline suggestions with initial following text`() {
+        useCodeGPTService()
+        myFixture.configureByText(
+            "CompletionTest.java",
+            "if () {\n   \n} else {\n}"
+        )
+        myFixture.editor.caretModel.moveToVisualPosition(VisualPosition(0, 4))
+        expectCodeGPT(StreamHttpExchange { request: RequestEntity ->
+            assertThat(request.uri.path).isEqualTo("/v1/code/completions")
+            assertThat(request.method).isEqualTo("POST")
+            assertThat(request.body)
+                .extracting("model", "prefix", "suffix", "fileExtension")
+                .containsExactly(
+                    "TEST_CODE_MODEL",
+                    "if (r",
+                    ") {\n   \n} else {\n}",
+                    "java"
+                )
+            listOf(
+                jsonMapResponse("choices", jsonArray(jsonMap("text", "oot == n"))),
+                jsonMapResponse("choices", jsonArray(jsonMap("text", "ull) {\n"))),
+                jsonMapResponse("choices", jsonArray(jsonMap("text", "   root = new Node"))),
+                jsonMapResponse("choices", jsonArray(jsonMap("text", "(data);\n"))),
+                jsonMapResponse("choices", jsonArray(jsonMap("text", "   return;"))),
+                jsonMapResponse("choices", jsonArray(jsonMap("text", "\n} else {"))),
+            )
+        })
+        myFixture.type('r')
+        assertRemainingCompletion {
+            it == "oot == null) {\n" +
+                    "   root = new Node(data);\n" +
+                    "   return;\n" +
+                    "} else {"
+        }
+        assertInlineSuggestion {
+            it == "oot == null"
+        }
+        myFixture.type('\t')
+        assertRemainingCompletion {
+            it == "root = new Node(data);\n" +
+                    "   return;\n" +
+                    "} else {"
+        }
+        assertInlineSuggestion {
+            it == "root = new Node(data);\n"
+        }
+        assertThat(myFixture.editor.caretModel.visualPosition).isEqualTo(VisualPosition(1, 3))
+        myFixture.type('\t')
+        assertRemainingCompletion {
+            it == "return;\n" +
+                    "} else {"
+        }
+        assertInlineSuggestion("Failed to assert remaining completion.") {
+            it == "return;\n"
+        }
+        assertThat(myFixture.editor.caretModel.visualPosition).isEqualTo(VisualPosition(2, 3))
+        myFixture.type('\t')
+        assertRemainingCompletion {
+            it == "} else {"
+        }
+        assertInlineSuggestion {
+            it == "} else {"
+        }
+        assertThat(myFixture.editor.caretModel.visualPosition).isEqualTo(VisualPosition(3, 0))
+        myFixture.type('\t')
+        assertRemainingCompletion {
+            it == ""
+        }
+    }
+
+    fun `test adjust completion line whitespaces`() {
+        useCodeGPTService()
+        myFixture.configureByText(
+            "CompletionTest.java",
+            "class Node {\n" +
+                    "  \n" +
+                    "}"
+        )
+        myFixture.editor.caretModel.moveToVisualPosition(VisualPosition(1, 3))
+        expectCodeGPT(StreamHttpExchange { request: RequestEntity ->
+            assertThat(request.uri.path).isEqualTo("/v1/code/completions")
+            assertThat(request.method).isEqualTo("POST")
+            assertThat(request.body)
+                .extracting("model", "prefix", "suffix", "fileExtension")
+                .containsExactly(
+                    "TEST_CODE_MODEL",
+                    "class Node {\n   ",
+                    "\n}",
+                    "java"
+                )
+            listOf(
+                jsonMapResponse("choices", jsonArray(jsonMap("text", "\n   int data;"))),
+                jsonMapResponse("choices", jsonArray(jsonMap("text", "\n   Node"))),
+                jsonMapResponse("choices", jsonArray(jsonMap("text", " left;\n   N"))),
+                jsonMapResponse("choices", jsonArray(jsonMap("text", "ode right;\n"))),
+            )
+        })
+        myFixture.type(' ')
+        assertRemainingCompletion {
+            it == "int data;\n" +
+                    "   Node left;\n" +
+                    "   Node right;\n"
+        }
+        assertInlineSuggestion {
+            it == "int data;\n"
+        }
+    }
+
+    private fun assertRemainingCompletion(
+        errorMessage: String = "Failed to assert remaining suggestion",
+        onAssert: (String) -> Boolean
+    ) {
+        PlatformTestUtil.waitWithEventsDispatching(
+            errorMessage,
+            {
+                val remainingCompletion = REMAINING_EDITOR_COMPLETION.get(myFixture.editor)
+                    ?: return@waitWithEventsDispatching false
+                onAssert(remainingCompletion)
+            },
+            5
+        )
+    }
+
+    private fun assertInlineSuggestion(
+        errorMessage: String = "Failed to assert inline suggestion",
+        onAssert: (String) -> Boolean
+    ) {
         PlatformTestUtil.waitWithEventsDispatching(
             errorMessage,
             {

src/test/kotlin/ee/carlrobert/codegpt/util/StringUtilTest.kt

@@ -0,0 +1,54 @@
+package ee.carlrobert.codegpt.util
+
+import ee.carlrobert.codegpt.util.StringUtil.findCompletionParts
+import org.assertj.core.api.Assertions.assertThat
+import org.junit.Test
+
+class StringUtilTest {
+
+    @Test
+    fun `should parse completion without brackets and braces`() {
+        val completionLine = "root != null"
+        val editorLineSuffix = ") {\n"
+
+        val result = findCompletionParts(editorLineSuffix, completionLine)
+
+        assertThat(result[0].second).isEqualTo(0)
+        assertThat(result[0].first).isEqualTo("root != null")
+    }
+
+    @Test
+    fun `should parse completion with closing bracket and brace into separate parts`() {
+        val completionLine = "root != null) {\n"
+        val editorLineSuffix = ")\n"
+
+        val result = findCompletionParts(editorLineSuffix, completionLine)
+
+        assertThat(result[0].second).isEqualTo(0)
+        assertThat(result[0].first).isEqualTo("root != null")
+        assertThat(result[1].second).isEqualTo(1)
+        assertThat(result[1].first).isEqualTo(" {")
+    }
+
+    @Test
+    fun `should parse completion when editor suffix contains closing bracket and brace`() {
+        val completionLine = "root != null) {\n"
+        val editorLineSuffix = ") {\n"
+
+        val result = findCompletionParts(editorLineSuffix, completionLine)
+
+        assertThat(result[0].second).isEqualTo(0)
+        assertThat(result[0].first).isEqualTo("root != null")
+    }
+
+    @Test
+    fun `should parse completion between opening and closing brackets`() {
+        val completionLine = "(root != null) {\n"
+        val editorLineSuffix = "() {\n"
+
+        val result = findCompletionParts(editorLineSuffix, completionLine)
+
+        assertThat(result[0].second).isEqualTo(1)
+        assertThat(result[0].first).isEqualTo("root != null")
+    }
+}