feat(intellij): add completion replace range support. (#613)

clients/intellij/build.gradle.kts

@@ -6,7 +6,7 @@ plugins {
 }
 
 group = "com.tabbyml"
-version = "0.6.0"
+version = "1.0.0-dev"
 
 repositories {
   mavenCentral()

clients/intellij/src/main/kotlin/com/tabbyml/intellijtabby/agent/Agent.kt

@@ -127,37 +127,14 @@ class Agent : ProcessAdapter() {
 
   data class Config(
     val server: Server? = null,
-    val completion: Completion? = null,
     val logs: Logs? = null,
     val anonymousUsageTracking: AnonymousUsageTracking? = null,
   ) {
     data class Server(
       val endpoint: String? = null,
       val requestHeaders: Map<String, String>? = null,
-      val requestTimeout: Int? = null,
     )
 
-    data class Completion(
-      val prompt: Prompt? = null,
-      val debounce: Debounce? = null,
-      val timeout: Timeout? = null,
-    ) {
-      data class Prompt(
-        val maxPrefixLines: Int? = null,
-        val maxSuffixLines: Int? = null,
-      )
-
-      data class Debounce(
-        val mode: String? = null,
-        val interval: Int? = null,
-      )
-
-      data class Timeout(
-        val auto: Int? = null,
-        val manually: Int? = null,
-      )
-    }
-
     data class Logs(
       val level: String? = null,
     )
@@ -252,7 +229,13 @@ class Agent : ProcessAdapter() {
     data class Choice(
       val index: Int,
       val text: String,
-    )
+      val replaceRange: Range,
+    ) {
+      data class Range(
+        val start: Int,
+        val end: Int,
+      )
+    }
   }
 
   suspend fun provideCompletions(request: CompletionRequest): CompletionResponse? {

clients/intellij/src/main/kotlin/com/tabbyml/intellijtabby/editor/EditorListener.kt

@@ -1,5 +1,6 @@
 package com.tabbyml.intellijtabby.editor
 
+import com.intellij.openapi.application.invokeLater
 import com.intellij.openapi.components.service
 import com.intellij.openapi.diagnostic.Logger
 import com.intellij.openapi.editor.Editor
@@ -38,8 +39,9 @@ class EditorListener : EditorFactoryListener {
       override fun documentChanged(event: DocumentEvent) {
         if (editorManager.selectedTextEditor == editor) {
           if (settings.completionTriggerMode == ApplicationSettingsState.TriggerMode.AUTOMATIC) {
-            val offset = event.offset + event.newFragment.length
-            completionProvider.provideCompletion(editor, offset)
+            invokeLater {
+              completionProvider.provideCompletion(editor, editor.caretModel.primaryCaret.offset)
+            }
           }
         }
       }

clients/intellij/src/main/kotlin/com/tabbyml/intellijtabby/editor/InlineCompletionService.kt

@@ -10,8 +10,12 @@ import com.intellij.openapi.editor.EditorCustomElementRenderer
 import com.intellij.openapi.editor.Inlay
 import com.intellij.openapi.editor.colors.EditorFontType
 import com.intellij.openapi.editor.impl.FontInfo
+import com.intellij.openapi.editor.markup.HighlighterLayer
+import com.intellij.openapi.editor.markup.HighlighterTargetArea
+import com.intellij.openapi.editor.markup.RangeHighlighter
 import com.intellij.openapi.editor.markup.TextAttributes
 import com.intellij.openapi.util.Disposer
+import com.intellij.openapi.util.TextRange
 import com.intellij.ui.JBColor
 import com.intellij.util.ui.UIUtil
 import com.tabbyml.intellijtabby.agent.Agent
@@ -30,8 +34,8 @@ class InlineCompletionService {
     val editor: Editor,
     val offset: Int,
     val completion: Agent.CompletionResponse,
-    val text: String,
     val inlays: List<Inlay<*>>,
+    val markups: List<RangeHighlighter>,
   )
 
   var shownInlineCompletion: InlineCompletion? = null
@@ -43,14 +47,109 @@ class InlineCompletionService {
       return
     }
     invokeLater {
-      // FIXME: support multiple choices
-      val text = completion.choices.first().text
-      logger.info("Showing inline completion at $offset: $text")
-      val lines = text.split("\n")
-      val inlays = lines
-        .mapIndexed { index, line -> createInlayLine(editor, offset, line, index) }
-        .filterNotNull()
-      shownInlineCompletion = InlineCompletion(editor, offset, completion, text, inlays)
+      if (editor.caretModel.offset != offset) {
+        return@invokeLater
+      }
+
+      // only support multiple choices for now
+      val choice = completion.choices.first()
+      logger.info("Showing inline completion at $offset: $choice")
+
+      val prefixReplaceLength = offset - choice.replaceRange.start
+      val suffixReplaceLength = choice.replaceRange.end - offset
+      val text = choice.text.substring(prefixReplaceLength)
+      if (text.isEmpty()) {
+        return@invokeLater
+      }
+      val currentLineNumber = editor.document.getLineNumber(offset)
+      val currentLineEndOffset = editor.document.getLineEndOffset(currentLineNumber)
+      if (currentLineEndOffset - offset < suffixReplaceLength) {
+        return@invokeLater
+      }
+      val currentLineSuffix = editor.document.getText(TextRange(offset, currentLineEndOffset))
+
+      val textLines = text.split("\n").toMutableList()
+
+      val inlays = mutableListOf<Inlay<*>>()
+      val markups = mutableListOf<RangeHighlighter>()
+      if (suffixReplaceLength == 0) {
+        // No replace range to handle
+        createInlayText(editor, textLines[0], offset, 0)?.let { inlays.add(it) }
+        if (textLines.size > 1) {
+          if (currentLineSuffix.isNotEmpty()) {
+            markupReplaceText(editor, offset, currentLineEndOffset).let { markups.add(it) }
+            textLines[textLines.lastIndex] += currentLineSuffix
+          }
+          textLines.forEachIndexed { index, line ->
+            if (index > 0) {
+              createInlayText(editor, line, offset, index)?.let { inlays.add(it) }
+            }
+          }
+        }
+      } else if (suffixReplaceLength == 1) {
+        logger.info("suffixReplaceLength: $suffixReplaceLength")
+        logger.info("currentLineSuffix: $currentLineSuffix")
+        logger.info("textLines[0]: ${textLines[0]}")
+        logger.info("textLines.size: ${textLines.size}")
+        // Replace range contains one char
+        val replaceChar = currentLineSuffix[0]
+        // Insert part is substring of first line that before the char
+        // Append part is substring of first line that after the char
+        // If first line doesn't contain the char, insert part is full first line, append part is empty
+        val insertPart = if (textLines[0].startsWith(replaceChar)) {
+          ""
+        } else {
+          textLines[0].split(replaceChar).first()
+        }
+        val appendPart = if (insertPart.length < textLines[0].length) {
+          textLines[0].substring(insertPart.length + 1)
+        } else {
+          ""
+        }
+        if (insertPart.isNotEmpty()) {
+          createInlayText(editor, insertPart, offset, 0)?.let { inlays.add(it) }
+        }
+        if (appendPart.isNotEmpty()) {
+          createInlayText(editor, appendPart, offset + 1, 0)?.let { inlays.add(it) }
+        }
+        if (textLines.size > 1) {
+          if (currentLineSuffix.isNotEmpty()) {
+            val startOffset = if (insertPart.length < textLines[0].length) {
+              // First line contains the char
+              offset + 1
+            } else {
+              // First line doesn't contain the char
+              offset
+            }
+            logger.info("startOffset: $startOffset")
+            markupReplaceText(editor, startOffset, currentLineEndOffset).let { markups.add(it) }
+            textLines[textLines.lastIndex] += currentLineSuffix.substring(1)
+          }
+          textLines.forEachIndexed { index, line ->
+            if (index > 0) {
+              createInlayText(editor, line, offset, index)?.let { inlays.add(it) }
+            }
+          }
+        }
+      } else {
+        // Replace range contains multiple chars
+        // It's hard to match these chars in the insertion text, we just mark them up
+        createInlayText(editor, textLines[0], offset, 0)?.let { inlays.add(it) }
+        markupReplaceText(editor, offset, offset + suffixReplaceLength).let { markups.add(it) }
+        if (textLines.size > 1) {
+          if (currentLineSuffix.length > suffixReplaceLength) {
+            markupReplaceText(editor, offset + suffixReplaceLength, currentLineEndOffset).let { markups.add(it) }
+            textLines[textLines.lastIndex] += currentLineSuffix.substring(suffixReplaceLength)
+          }
+          textLines.forEachIndexed { index, line ->
+            if (index > 0) {
+              createInlayText(editor, line, offset, index)?.let { inlays.add(it) }
+            }
+          }
+        }
+      }
+
+      shownInlineCompletion = InlineCompletion(editor, offset, completion, inlays, markups)
     }
     val agentService = service<AgentService>()
     agentService.scope.launch {
@@ -66,10 +165,15 @@ class InlineCompletionService {
 
   fun accept() {
     shownInlineCompletion?.let {
-      logger.info("Accept inline completion at ${it.offset}: ${it.text}")
+      val choice = it.completion.choices.first()
+      logger.info("Accept inline completion at ${it.offset}: $choice")
+
+      val prefixReplaceLength = it.offset - choice.replaceRange.start
+      val text = choice.text.substring(prefixReplaceLength)
       WriteCommandAction.runWriteCommandAction(it.editor.project) {
-        it.editor.document.insertString(it.offset, it.text)
-        it.editor.caretModel.moveToOffset(it.offset + it.text.length)
+        it.editor.document.deleteString(it.offset, choice.replaceRange.end)
+        it.editor.document.insertString(it.offset, text)
+        it.editor.caretModel.moveToOffset(it.offset + text.length)
       }
       invokeLater {
         it.inlays.forEach(Disposer::dispose)
@@ -80,7 +184,7 @@ class InlineCompletionService {
           Agent.LogEventRequest(
             type = Agent.LogEventRequest.EventType.SELECT,
             completionId = it.completion.id,
-            choiceIndex = it.completion.choices.first().index
+            choiceIndex = choice.index
           )
         )
       }
@@ -92,26 +196,29 @@ class InlineCompletionService {
     shownInlineCompletion?.let {
       invokeLater {
         it.inlays.forEach(Disposer::dispose)
+        it.markups.forEach { markup ->
+          it.editor.markupModel.removeHighlighter(markup)
+        }
       }
       shownInlineCompletion = null
     }
   }
 
-  private fun createInlayLine(editor: Editor, offset: Int, line: String, index: Int): Inlay<*>? {
+  private fun createInlayText(editor: Editor, text: String, offset: Int, lineOffset: Int): Inlay<*>? {
     val renderer = object : EditorCustomElementRenderer {
       override fun calcWidthInPixels(inlay: Inlay<*>): Int {
-        return maxOf(getWidth(inlay.editor, line), 1)
+        return maxOf(getWidth(inlay.editor, text), 1)
       }
 
       override fun paint(inlay: Inlay<*>, graphics: Graphics, targetRect: Rectangle, textAttributes: TextAttributes) {
         graphics.font = getFont(inlay.editor)
         graphics.color = JBColor.GRAY
-        graphics.drawString(line, targetRect.x, targetRect.y + inlay.editor.ascent)
+        graphics.drawString(text, targetRect.x, targetRect.y + inlay.editor.ascent)
       }
 
       private fun getFont(editor: Editor): Font {
         return editor.colorsScheme.getFont(EditorFontType.PLAIN).let {
-          UIUtil.getFontWithFallbackIfNeeded(it, line).deriveFont(editor.colorsScheme.editorFontSize)
+          UIUtil.getFontWithFallbackIfNeeded(it, text).deriveFont(editor.colorsScheme.editorFontSize)
         }
       }
 
@@ -121,10 +228,20 @@ class InlineCompletionService {
         return metrics.stringWidth(line)
       }
     }
-    return if (index == 0) {
+    return if (lineOffset == 0) {
       editor.inlayModel.addInlineElement(offset, true, renderer)
     } else {
-      editor.inlayModel.addBlockElement(offset, true, false, -index, renderer)
+      editor.inlayModel.addBlockElement(offset, true, false, -lineOffset, renderer)
+    }
+  }
+
+  private fun markupReplaceText(editor: Editor, startOffset: Int, endOffset: Int): RangeHighlighter {
+    val textAttributes = TextAttributes().apply {
+      foregroundColor = JBColor.background()
+      backgroundColor = JBColor.background()
     }
+    return editor.markupModel.addRangeHighlighter(
+      startOffset, endOffset, HighlighterLayer.LAST + 1000, textAttributes, HighlighterTargetArea.EXACT_RANGE
+    )
   }
-}
+}