feat(intellij): auto trigger completion and add hotkeys to accept/dismiss completion. (#315)

clients/intellij/.idea/misc.xml

@@ -1,4 +1,3 @@
-<?xml version="1.0" encoding="UTF-8"?>
 <project version="4">
   <component name="ExternalStorageConfigurationManager" enabled="true" />
   <component name="ProjectRootManager" version="2" languageLevel="JDK_17" default="true" project-jdk-name="jbr-17" project-jdk-type="JavaSDK">

clients/intellij/src/main/kotlin/com/tabbyml/intellijtabby/actions/AcceptCompletion.kt

@@ -11,15 +11,15 @@ import com.tabbyml.intellijtabby.editor.InlineCompletionService
 class AcceptCompletion : AnAction() {
   override fun actionPerformed(e: AnActionEvent) {
     val inlineCompletionService = service<InlineCompletionService>()
-    val editor = e.getRequiredData(CommonDataKeys.EDITOR)
+    inlineCompletionService.accept()
-    inlineCompletionService.accept(editor)
   }
   
   override fun update(e: AnActionEvent) {
     val inlineCompletionService = service<InlineCompletionService>()
-    e.presentation.isEnabled = e.getData(CommonDataKeys.EDITOR) != null
+    e.presentation.isEnabled = e.project != null
-        && e.project != null
+        && e.getData(CommonDataKeys.EDITOR) != null
-        && inlineCompletionService.currentText != null
+        && inlineCompletionService.shownInlineCompletion != null
+        && e.getData(CommonDataKeys.EDITOR) == inlineCompletionService.shownInlineCompletion?.editor
   }
 
   override fun getActionUpdateThread(): ActionUpdateThread {

clients/intellij/src/main/kotlin/com/tabbyml/intellijtabby/actions/DismissCompletion.kt

@@ -16,9 +16,10 @@ class DismissCompletion : AnAction() {
 
   override fun update(e: AnActionEvent) {
     val inlineCompletionService = service<InlineCompletionService>()
-    e.presentation.isEnabled = e.getData(CommonDataKeys.EDITOR) != null
+    e.presentation.isEnabled = e.project != null
-        && e.project != null
+        && e.getData(CommonDataKeys.EDITOR) != null
-        && inlineCompletionService.currentText != null
+        && inlineCompletionService.shownInlineCompletion != null
+        && e.getData(CommonDataKeys.EDITOR) == inlineCompletionService.shownInlineCompletion?.editor
   }
 
   override fun getActionUpdateThread(): ActionUpdateThread {

clients/intellij/src/main/kotlin/com/tabbyml/intellijtabby/actions/TriggerCompletion.kt

@@ -4,7 +4,6 @@ import com.intellij.openapi.actionSystem.ActionUpdateThread
 import com.intellij.openapi.actionSystem.AnAction
 import com.intellij.openapi.actionSystem.AnActionEvent
 import com.intellij.openapi.actionSystem.CommonDataKeys
-import com.intellij.openapi.application.invokeLater
 import com.intellij.openapi.components.service
 import com.tabbyml.intellijtabby.agent.AgentService
 import com.tabbyml.intellijtabby.editor.InlineCompletionService
@@ -15,23 +14,17 @@ class TriggerCompletion : AnAction() {
     val agentService = service<AgentService>()
     val inlineCompletionService = service<InlineCompletionService>()
     val editor = e.getRequiredData(CommonDataKeys.EDITOR)
-    val file = e.getRequiredData(CommonDataKeys.PSI_FILE)
     val offset = editor.caretModel.primaryCaret.offset
 
     inlineCompletionService.dismiss()
-    agentService.getCompletion(editor, file, offset)?.thenAccept {
+    agentService.getCompletion(editor, offset)?.thenAccept {
-      invokeLater {
       inlineCompletionService.show(editor, offset, it)
     }
   }
-  }
 
   override fun update(e: AnActionEvent) {
-    val inlineCompletionService = service<InlineCompletionService>()
     e.presentation.isEnabled = e.project != null
         && e.getData(CommonDataKeys.EDITOR) != null
-        && e.getData(CommonDataKeys.PSI_FILE) != null
-        && inlineCompletionService.currentText == null
   }
 
   override fun getActionUpdateThread(): ActionUpdateThread {

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

@@ -84,14 +84,14 @@ class Agent : ProcessAdapter() {
     )
   }
 
-  fun getCompletions(request: CompletionRequest): CompletableFuture<CompletionResponse> {
+  fun getCompletions(request: CompletionRequest): CompletableFuture<CompletionResponse?> {
     return request("getCompletions", listOf(request))
   }
 
   private var requestId = 1
   private var ongoingRequest = mutableMapOf<Int, (response: String) -> Unit>()
 
-  private inline fun <reified T : Any> request(func: String, args: List<Any> = emptyList()): CompletableFuture<T> {
+  private inline fun <reified T : Any?> request(func: String, args: List<Any> = emptyList()): CompletableFuture<T> {
     val id = requestId++
     val data = listOf(id, mapOf("func" to func, "args" to args))
     val json = gson.toJson(data)

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

@@ -1,8 +1,10 @@
 package com.tabbyml.intellijtabby.agent
 
+import com.intellij.openapi.application.ReadAction
 import com.intellij.openapi.components.Service
 import com.intellij.openapi.diagnostic.Logger
 import com.intellij.openapi.editor.Editor
+import com.intellij.psi.PsiDocumentManager
 import com.intellij.psi.PsiFile
 import java.util.concurrent.CompletableFuture
 
@@ -23,9 +25,14 @@ class AgentService {
     }
   }
 
-  fun getCompletion(editor: Editor, file: PsiFile, offset: Int): CompletableFuture<Agent.CompletionResponse>? {
+  fun getCompletion(editor: Editor, offset: Int): CompletableFuture<Agent.CompletionResponse>? {
-    return agent.thenCompose {
+    return agent.thenCompose {agent ->
-      it?.getCompletions(Agent.CompletionRequest(
+      ReadAction.compute<PsiFile, Throwable> {
+        editor.project?.let { project ->
+          PsiDocumentManager.getInstance(project).getPsiFile(editor.document)
+        }
+      }?.let { file ->
+        agent?.getCompletions(Agent.CompletionRequest(
           file.virtualFile.path,
           file.language.id, // FIXME: map language id
           editor.document.text,
@@ -33,4 +40,5 @@ class AgentService {
         ))
       }
     }
+  }
 }

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

@@ -0,0 +1,62 @@
+package com.tabbyml.intellijtabby.editor
+
+import com.intellij.openapi.components.Service
+import com.intellij.openapi.components.service
+import com.intellij.openapi.diagnostic.Logger
+import com.intellij.openapi.editor.Editor
+import com.intellij.openapi.fileEditor.FileEditorManagerEvent
+import com.intellij.openapi.fileEditor.FileEditorManagerListener
+import com.intellij.openapi.project.Project
+import com.tabbyml.intellijtabby.agent.AgentService
+import java.util.*
+
+@Service
+class CompletionScheduler {
+  private val logger = Logger.getInstance(CompletionScheduler::class.java)
+
+  data class CompletionContext(val editor: Editor, val offset: Int, val timer: Timer)
+
+  private var project: Project? = null
+  var scheduled: CompletionContext? = null
+    private set
+
+
+  fun schedule(editor: Editor, offset: Int) {
+    clear()
+    val agentService = service<AgentService>()
+    val inlineCompletionService = service<InlineCompletionService>()
+    inlineCompletionService.dismiss()
+    val timer = Timer()
+    timer.schedule(object : TimerTask() {
+      override fun run() {
+        logger.info("Scheduled completion task running")
+        agentService.getCompletion(editor, offset)?.thenAccept {
+          inlineCompletionService.show(editor, offset, it)
+        }
+      }
+    }, 150)
+    scheduled = CompletionContext(editor, offset, timer)
+
+    if (project != editor.project) {
+      project = editor.project
+      editor.project?.messageBus?.connect()?.subscribe(
+        FileEditorManagerListener.FILE_EDITOR_MANAGER,
+        object: FileEditorManagerListener {
+          override fun selectionChanged(event: FileEditorManagerEvent) {
+            logger.info("FileEditorManagerListener selectionChanged.")
+            clear()
+          }
+        }
+      )
+    }
+  }
+
+  fun clear() {
+    scheduled?.let {
+      it.timer.cancel()
+      scheduled = null
+    }
+    val inlineCompletionService = service<InlineCompletionService>()
+    inlineCompletionService.dismiss()
+  }
+}

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

@@ -0,0 +1,37 @@
+package com.tabbyml.intellijtabby.editor
+
+import com.intellij.openapi.components.service
+import com.intellij.openapi.diagnostic.Logger
+import com.intellij.openapi.editor.event.*
+import com.intellij.openapi.fileEditor.FileEditorManager
+
+class EditorListener : EditorFactoryListener {
+  private val logger = Logger.getInstance(EditorListener::class.java)
+
+  override fun editorCreated(event: EditorFactoryEvent) {
+    val editor = event.editor
+    val editorManager = editor.project?.let { FileEditorManager.getInstance(it) } ?: return
+    val completionScheduler = service<CompletionScheduler>()
+
+    editor.caretModel.addCaretListener(object : CaretListener {
+      override fun caretPositionChanged(event: CaretEvent) {
+        if (editorManager.selectedTextEditor == editor) {
+          completionScheduler.scheduled?.let {
+            if (it.editor != editor || it.offset != editor.caretModel.primaryCaret.offset) {
+              completionScheduler.clear()
+            }
+          }
+        }
+      }
+    })
+
+    editor.document.addDocumentListener(object : DocumentListener {
+      override fun documentChanged(event: DocumentEvent) {
+        if (editorManager.selectedTextEditor == editor) {
+          val offset = event.offset + event.newFragment.length
+          completionScheduler.schedule(editor, offset)
+        }
+      }
+    })
+  }
+}

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

@@ -1,6 +1,6 @@
 package com.tabbyml.intellijtabby.editor
 
-import com.intellij.openapi.application.ReadAction
+import com.intellij.openapi.application.invokeLater
 import com.intellij.openapi.command.WriteCommandAction
 import com.intellij.openapi.components.Service
 import com.intellij.openapi.diagnostic.Logger
@@ -14,50 +14,56 @@ import com.tabbyml.intellijtabby.agent.Agent
 import java.awt.Graphics
 import java.awt.Rectangle
 
+
 @Service
 class InlineCompletionService {
   private val logger = Logger.getInstance(InlineCompletionService::class.java)
-  var currentText: String? = null
+
+  data class InlineCompletion(val editor: Editor, val text: String, val offset: Int, val inlays: List<Inlay<*>>)
+
+  var shownInlineCompletion: InlineCompletion? = null
     private set
-  var currentOffset: Int? = null
-    private set
-  private var currentInlays: MutableList<Inlay<*>> = mutableListOf()
 
   fun show(editor: Editor, offset: Int, completion: Agent.CompletionResponse) {
+    dismiss()
     if (completion.choices.isEmpty()) {
       return
     }
+    invokeLater {
       val text = completion.choices.first().text
       logger.info("Showing inline completion at $offset: $text")
       val lines = text.split("\n")
-    lines.forEachIndexed { index, line -> addInlayLine(editor, offset, line, index) }
+      val inlays = lines
-    currentText = text
+        .mapIndexed { index, line -> createInlayLine(editor, offset, line, index) }
-    currentOffset = offset
+        .filterNotNull()
+      shownInlineCompletion = InlineCompletion(editor, text, offset, inlays)
+    }
   }
 
-  fun accept(editor: Editor) {
+  fun accept() {
-    currentText?.let {
+    shownInlineCompletion?.let {
-      WriteCommandAction.runWriteCommandAction(editor.project) {
+      logger.info("Accept inline completion at ${it.offset}: ${it.text}")
-        editor.document.insertString(currentOffset!!, it)
+      WriteCommandAction.runWriteCommandAction(it.editor.project) {
-        editor.caretModel.moveToOffset(currentOffset!! + it.length)
+        it.editor.document.insertString(it.offset, it.text)
+        it.editor.caretModel.moveToOffset(it.offset + it.text.length)
       }
-      currentText = null
+      invokeLater {
-      currentOffset = null
+        it.inlays.forEach(Disposer::dispose)
-      currentInlays.forEach(Disposer::dispose)
+      }
-      currentInlays = mutableListOf()
+      shownInlineCompletion = null
     }
   }
 
   fun dismiss() {
-    currentText?.let {
+    shownInlineCompletion?.let {
-      currentText = null
+      invokeLater {
-      currentOffset = null
+        it.inlays.forEach(Disposer::dispose)
-      currentInlays.forEach(Disposer::dispose)
+      }
-      currentInlays = mutableListOf()
+      shownInlineCompletion = null
     }
   }
 
-  private fun addInlayLine(editor: Editor, offset: Int, line: String, index: Int) {
+  private fun createInlayLine(editor: Editor, offset: Int, line: String, index: Int): Inlay<*>? {
     val renderer = object : EditorCustomElementRenderer {
       override fun calcWidthInPixels(inlay: Inlay<*>): Int {
         // FIXME: Calc width?
@@ -69,13 +75,10 @@ class InlineCompletionService {
         graphics.drawString(line, targetRect.x, targetRect.y + inlay.editor.ascent)
       }
     }
-    val inlay = if (index == 0) {
+    return if (index == 0) {
       editor.inlayModel.addInlineElement(offset, true, renderer)
     } else {
       editor.inlayModel.addBlockElement(offset, true, false, -index, renderer)
     }
-    inlay?.let {
-      currentInlays.add(it)
-    }
   }
 }

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

@@ -25,26 +25,30 @@
     <!-- Extension points defined by the plugin.
          Read more: https://plugins.jetbrains.com/docs/intellij/plugin-extension-points.html -->
     <extensions defaultExtensionNs="com.intellij">
+        <editorFactoryListener implementation="com.tabbyml.intellijtabby.editor.EditorListener"/>
     </extensions>
 
     <actions>
+        <group id="Tabby.ToolsMenu" popup="true" text="Tabby" description="Tabby AI code assistant">
+            <add-to-group group-id="ToolsMenu" anchor="last"/>
             <action id="Tabby.TriggerCompletion"
                     class="com.tabbyml.intellijtabby.actions.TriggerCompletion"
-                text="TriggerCompletion"
+                    text="Trigger Completion"
                     description="Trigger completion at current position.">
-            <add-to-group group-id="EditorPopupMenu" anchor="first"/>
+                <keyboard-shortcut first-keystroke="alt BACK_SLASH" keymap="$default"/>
             </action>
             <action id="Tabby.AcceptCompletion"
                     class="com.tabbyml.intellijtabby.actions.AcceptCompletion"
-                text="AcceptCompletion"
+                    text="Accept Completion"
                     description="Trigger completion at current position.">
-            <add-to-group group-id="EditorPopupMenu" anchor="after" relative-to-action="Tabby.TriggerCompletion"/>
+                <keyboard-shortcut first-keystroke="TAB" keymap="$default"/>
             </action>
             <action id="Tabby.DismissCompletion"
                     class="com.tabbyml.intellijtabby.actions.DismissCompletion"
-                text="DismissCompletion"
+                    text="Dismiss Completion"
                     description="Trigger completion at current position.">
-            <add-to-group group-id="EditorPopupMenu" anchor="after" relative-to-action="Tabby.AcceptCompletion"/>
+                <keyboard-shortcut first-keystroke="ESCAPE" keymap="$default"/>
             </action>
+        </group>
     </actions>
 </idea-plugin>