From 3fb1016549a41ad87f565cc6a9f459c3348282ec Mon Sep 17 00:00:00 2001
From: "Kirill.Karnaukhov" <kirill.karnaukhov@jetbrains.com>
Date: Thu, 10 Oct 2024 14:39:40 +0200
Subject: [PATCH] [ai-completion] LLM-11794: call AI completion on insertion of
 templates

GitOrigin-RevId: fc92fb2ae37b974a4836b364021958bcdad87ed7
---
 .../completion/InlineCompletionEvent.kt       | 10 ++++++
 .../listeners/InlineCompletionListeners.kt    | 33 +++++++++++++++----
 2 files changed, 37 insertions(+), 6 deletions(-)

diff --git a/platform/platform-impl/src/com/intellij/codeInsight/inline/completion/InlineCompletionEvent.kt b/platform/platform-impl/src/com/intellij/codeInsight/inline/completion/InlineCompletionEvent.kt
index a20745d37265..ba0829ae56a1 100644
--- a/platform/platform-impl/src/com/intellij/codeInsight/inline/completion/InlineCompletionEvent.kt
+++ b/platform/platform-impl/src/com/intellij/codeInsight/inline/completion/InlineCompletionEvent.kt
@@ -232,6 +232,16 @@ interface InlineCompletionEvent {
     }
   }
 
+  /**
+   * Triggered by insertion of templates, like Live Templates.
+   */
+  @ApiStatus.Experimental
+  class TemplateInserted internal constructor(val editor: Editor) : InlineCompletionEvent {
+    override fun toRequest(): InlineCompletionRequest? {
+      return getRequest(event = this, editor = editor)
+    }
+  }
+
   @ApiStatus.Internal
   @ApiStatus.Experimental
   sealed class PartialAccept @ApiStatus.Internal constructor(val editor: Editor) : InlineCompletionEvent {
diff --git a/platform/platform-impl/src/com/intellij/codeInsight/inline/completion/listeners/InlineCompletionListeners.kt b/platform/platform-impl/src/com/intellij/codeInsight/inline/completion/listeners/InlineCompletionListeners.kt
index b8709e06e007..449170a16779 100644
--- a/platform/platform-impl/src/com/intellij/codeInsight/inline/completion/listeners/InlineCompletionListeners.kt
+++ b/platform/platform-impl/src/com/intellij/codeInsight/inline/completion/listeners/InlineCompletionListeners.kt
@@ -4,6 +4,7 @@ package com.intellij.codeInsight.inline.completion.listeners
 import com.intellij.codeInsight.completion.CompletionUtil
 import com.intellij.codeInsight.editorActions.TypedHandlerDelegate
 import com.intellij.codeInsight.inline.completion.InlineCompletion
+import com.intellij.codeInsight.inline.completion.InlineCompletionEvent
 import com.intellij.codeInsight.inline.completion.TypingEvent
 import com.intellij.codeInsight.inline.completion.logs.InlineCompletionUsageTracker.ShownEvents.FinishType
 import com.intellij.codeInsight.inline.completion.session.InlineCompletionContext
@@ -14,8 +15,7 @@ import com.intellij.codeInsight.template.TemplateManagerListener
 import com.intellij.codeInsight.template.impl.TemplateState
 import com.intellij.codeWithMe.ClientId
 import com.intellij.codeWithMe.isCurrent
-import com.intellij.openapi.actionSystem.CommonDataKeys
-import com.intellij.openapi.actionSystem.DataContext
+import com.intellij.openapi.actionSystem.*
 import com.intellij.openapi.actionSystem.ex.AnActionListener
 import com.intellij.openapi.diagnostic.thisLogger
 import com.intellij.openapi.editor.ClientEditorManager
@@ -27,6 +27,7 @@ import com.intellij.openapi.util.Key
 import com.intellij.openapi.util.TextRange
 import com.intellij.openapi.util.removeUserData
 import com.intellij.psi.PsiFile
+import com.intellij.util.application
 import com.intellij.util.concurrency.annotations.RequiresEdt
 import org.jetbrains.annotations.ApiStatus
 import java.awt.event.FocusEvent
@@ -225,7 +226,7 @@ private class InlineCompletionTemplateListener : TemplateManagerListener {
   override fun templateStarted(state: TemplateState) {
     state.editor?.let { editor ->
       start(editor)
-      state.addTemplateStateListener(TemplateStateListener { finish(editor) })
+      state.addTemplateStateListener(TemplateStateListener(editor, ::finish))
     }
   }
 
@@ -237,13 +238,33 @@ private class InlineCompletionTemplateListener : TemplateManagerListener {
     editor.removeUserData(TEMPLATE_IN_PROGRESS_KEY)
   }
 
-  private class TemplateStateListener(private val finish: () -> Unit) : TemplateEditingListener {
-    override fun templateFinished(template: Template, brokenOff: Boolean) = finish()
-    override fun templateCancelled(template: Template?) = finish()
+  private class TemplateStateListener(private val editor: Editor, private val finish: (Editor) -> Unit) : TemplateEditingListener {
+
+    override fun templateFinished(template: Template, brokenOff: Boolean) {
+      finish(editor)
+      if (!brokenOff && template.isSuitableToInvokeEvent()) {
+        application.invokeLater {
+          if (!editor.isDisposed) {
+            InlineCompletion.getHandlerOrNull(editor)?.invokeEvent(InlineCompletionEvent.TemplateInserted(editor))
+          }
+        }
+      }
+    }
+
+    override fun templateCancelled(template: Template?) {
+      finish(editor)
+    }
 
     override fun currentVariableChanged(templateState: TemplateState, template: Template?, oldIndex: Int, newIndex: Int) = Unit
+
     override fun waitingForInput(template: Template?) = Unit
+
     override fun beforeTemplateFinished(state: TemplateState, template: Template?) = Unit
+
+    private fun Template.isSuitableToInvokeEvent(): Boolean {
+      // Usually, 'key' is responsible for the name of the live template
+      return key != ""
+    }
   }
 
   companion object {