fix: CopilotUndoManager service for tracking actions on save (#124)

* CopilotUndoManager service for tracking actions on save

* unit test and fixes

* project metadata

* spotless

* run tests on gh actions

* Update validation.yml

use gradel wrapper from project

* use test logger

.github/workflows/validation.yml

@@ -19,10 +19,10 @@ jobs:
           java-version: '21'
           distribution: 'temurin'
           cache: gradle
-      - name: Setup Gradle
-        uses: gradle/actions/setup-gradle@v3
       - name: Check Spotless
         run: ./gradlew spotlessCheck
+      - name: Run Unit Tests
+        run: ./gradlew test
 
   buildPlugin:
     runs-on: ubuntu-latest
@@ -35,8 +35,6 @@ jobs:
           java-version: '21'
           distribution: 'temurin'
           cache: gradle
-      - name: Setup Gradle
-        uses: gradle/actions/setup-gradle@v3
       - name: Build plugin
         run: ./gradlew --no-configuration-cache buildPlugin
       - uses: actions/upload-artifact@v4.3.3
@@ -62,7 +60,5 @@ jobs:
           java-version: '21'
           distribution: 'temurin'
           cache: gradle
-      - name: Setup Gradle
-        uses: gradle/actions/setup-gradle@v3
       - name: Verify plugin
         run: ./gradlew --no-configuration-cache verifyPlugin

build.gradle.kts

@@ -1,8 +1,12 @@
+import org.jetbrains.intellij.platform.gradle.TestFrameworkType
+
 plugins {
   id("java")
   id("org.jetbrains.kotlin.jvm") version "1.9.21"
   id("org.jetbrains.intellij.platform") version "2.0.0"
   id("com.diffplug.spotless") version "7.0.0.BETA2"
+
+  id("com.adarshr.test-logger") version "4.0.0"
 }
 
 group = "com.vaadin"
@@ -48,7 +52,12 @@ dependencies {
     pluginVerifier()
     zipSigner()
     instrumentationTools()
+
+    testFramework(TestFrameworkType.Platform)
   }
+
+  testImplementation(kotlin("test"))
+  testImplementation("junit:junit:4.13.2")
 }
 
 configure<com.diffplug.gradle.spotless.SpotlessExtension> {
@@ -84,4 +93,6 @@ tasks {
     channels.set(listOf(publishChannel))
     token.set(System.getenv("PUBLISH_TOKEN"))
   }
+
+  test { useJUnitPlatform() }
 }

src/main/kotlin/com/vaadin/plugin/copilot/handler/AbstractHandler.kt

@@ -5,10 +5,13 @@ import com.intellij.openapi.editor.Document
 import com.intellij.openapi.fileEditor.FileDocumentManager
 import com.intellij.openapi.fileEditor.FileEditor
 import com.intellij.openapi.fileEditor.FileEditorManager
+import com.intellij.openapi.fileEditor.OpenFileDescriptor
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.vfs.VirtualFile
+import com.intellij.openapi.vfs.findDocument
 import com.intellij.psi.PsiDocumentManager
 import com.vaadin.plugin.copilot.CopilotPluginUtil
+import com.vaadin.plugin.copilot.service.CopilotUndoManager
 import io.netty.handler.codec.http.HttpResponseStatus
 import java.io.File
 
@@ -58,10 +61,30 @@ abstract class AbstractHandler(val project: Project) : Handler {
         return FileEditorWrapper(editors.first(), project, false)
     }
 
+    open fun postSave(vfsFile: VirtualFile) {
+        LOG.info("File $vfsFile contents saved")
+        notifyUndoManager(vfsFile)
+        commitAndFlush(vfsFile.findDocument())
+        openFileInEditor(vfsFile)
+    }
+
+    private fun openFileInEditor(vfsFile: VirtualFile) {
+        val openFileDescriptor = OpenFileDescriptor(project, vfsFile)
+        FileEditorManager.getInstance(project).openTextEditor(openFileDescriptor, false)
+    }
+
+    fun notifyUndoManager(vfsFile: VirtualFile) {
+        getCopilotUndoManager().fileWritten(vfsFile)
+    }
+
     fun commitAndFlush(vfsDoc: Document?) {
         if (vfsDoc != null) {
             PsiDocumentManager.getInstance(project).commitDocument(vfsDoc)
             FileDocumentManager.getInstance().saveDocuments(vfsDoc::equals)
         }
     }
+
+    fun getCopilotUndoManager(): CopilotUndoManager {
+        return project.getService(CopilotUndoManager::class.java)
+    }
 }

src/main/kotlin/com/vaadin/plugin/copilot/handler/RedoHandler.kt

@@ -1,36 +1,26 @@
 package com.vaadin.plugin.copilot.handler
 
-import com.intellij.openapi.application.runInEdt
-import com.intellij.openapi.application.runWriteAction
-import com.intellij.openapi.command.impl.UndoManagerImpl
+import com.intellij.openapi.command.undo.UndoManager
+import com.intellij.openapi.fileEditor.FileEditor
 import com.intellij.openapi.project.Project
-import com.intellij.openapi.vfs.findDocument
-import com.vaadin.plugin.actions.VaadinCompileOnSaveAction
+import com.intellij.openapi.vfs.VirtualFile
 
 class RedoHandler(project: Project, data: Map<String, Any>) : UndoHandler(project, data) {
 
-    private val copilotActionPrefix = "_Redo Vaadin Copilot"
+    override fun getOpsCount(vfsFile: VirtualFile): Int {
+        return getCopilotUndoManager().getRedoCount(vfsFile)
+    }
+
+    override fun before(vfsFile: VirtualFile) {
+        getCopilotUndoManager().redoStart(vfsFile)
+    }
+
+    override fun runManagerAction(undoManager: UndoManager, editor: FileEditor) {
+        undoManager.redo(editor)
+    }
 
-    override fun run(): HandlerResponse {
-        for (vfsFile in vfsFiles) {
-            runInEdt {
-                getEditorWrapper(vfsFile).use { wrapper ->
-                    val editor = wrapper.getFileEditor()
-                    val undoManager = UndoManagerImpl.getInstance(project)
-                    runWriteAction {
-                        if (undoManager.isRedoAvailable(editor)) {
-                            val redo = undoManager.getRedoActionNameAndDescription(editor).first
-                            if (redo.startsWith(copilotActionPrefix)) {
-                                undoManager.redo(editor)
-                                commitAndFlush(vfsFile.findDocument())
-                                LOG.info("$redo performed on ${vfsFile.path}")
-                                VaadinCompileOnSaveAction().compile(project, vfsFile)
-                            }
-                        }
-                    }
-                }
-            }
-        }
-        return RESPONSE_OK
+    override fun after(vfsFile: VirtualFile) {
+        getCopilotUndoManager().redoDone(vfsFile)
+        LOG.info("$vfsFile redo performed")
     }
 }

src/main/kotlin/com/vaadin/plugin/copilot/handler/UndoHandler.kt

@@ -3,51 +3,84 @@ package com.vaadin.plugin.copilot.handler
 import com.intellij.openapi.application.runInEdt
 import com.intellij.openapi.application.runWriteAction
 import com.intellij.openapi.command.impl.UndoManagerImpl
+import com.intellij.openapi.command.undo.UndoManager
+import com.intellij.openapi.fileEditor.FileEditor
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.vfs.VfsUtil
 import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.openapi.vfs.findDocument
-import com.vaadin.plugin.actions.VaadinCompileOnSaveAction
 import java.io.File
 
 open class UndoHandler(project: Project, data: Map<String, Any>) : AbstractHandler(project) {
 
-    private val copilotActionPrefix = "_Undo Vaadin Copilot"
-
-    protected val vfsFiles: ArrayList<VirtualFile> = ArrayList()
+    private val vfsFiles: ArrayList<VirtualFile> = ArrayList()
 
     init {
         val paths = data["files"] as Collection<String>
         for (path in paths) {
             val file = File(path)
             if (isFileInsideProject(project, file)) {
-                VfsUtil.findFileByIoFile(file, true)?.let { vfsFiles.add(it) }
+                var vfsFile = VfsUtil.findFileByIoFile(file, true)
+                if (vfsFile != null) {
+                    vfsFiles.add(vfsFile)
+                } else {
+                    // if we want to undo file removal we need to create empty virtual file to write
+                    // content
+                    runWriteAction {
+                        val parent = VfsUtil.createDirectories(file.parent)
+                        vfsFile = parent.createChildData(this, file.name)
+                        vfsFiles.add(vfsFile!!)
+                    }
+                }
             } else {
-                LOG.warn("File $file is not a part of a project")
+                LOG.warn("File $path is not a part of a project")
             }
         }
     }
 
     override fun run(): HandlerResponse {
         for (vfsFile in vfsFiles) {
+            val count = getOpsCount(vfsFile)
+            if (count == 0) {
+                continue
+            }
+
             runInEdt {
                 getEditorWrapper(vfsFile).use { wrapper ->
                     val undoManager = UndoManagerImpl.getInstance(project)
                     val editor = wrapper.getFileEditor()
                     runWriteAction {
-                        if (undoManager.isUndoAvailable(editor)) {
-                            val undo = undoManager.getUndoActionNameAndDescription(editor).first
-                            if (undo.startsWith(copilotActionPrefix)) {
-                                undoManager.undo(editor)
-                                commitAndFlush(vfsFile.findDocument())
-                                LOG.info("$undo performed on ${vfsFile.path}")
-                                VaadinCompileOnSaveAction().compile(project, vfsFile)
+                        try {
+                            before(vfsFile)
+                            var i = 0
+                            while (i++ < count) {
+                                runManagerAction(undoManager, editor)
                             }
+                            commitAndFlush(vfsFile.findDocument())
+                        } finally {
+                            after(vfsFile)
                         }
                     }
                 }
             }
         }
         return RESPONSE_OK
     }
+
+    open fun getOpsCount(vfsFile: VirtualFile): Int {
+        return getCopilotUndoManager().getUndoCount(vfsFile)
+    }
+
+    open fun before(vfsFile: VirtualFile) {
+        getCopilotUndoManager().undoStart(vfsFile)
+    }
+
+    open fun runManagerAction(undoManager: UndoManager, editor: FileEditor) {
+        undoManager.undo(editor)
+    }
+
+    open fun after(vfsFile: VirtualFile) {
+        getCopilotUndoManager().undoDone(vfsFile)
+        LOG.info("$vfsFile undo performed")
+    }
 }

src/main/kotlin/com/vaadin/plugin/copilot/handler/WriteBase64FileHandler.kt

@@ -2,24 +2,12 @@ package com.vaadin.plugin.copilot.handler
 
 import com.intellij.openapi.editor.Document
 import com.intellij.openapi.project.Project
-import com.intellij.openapi.vfs.VfsUtil
 import com.intellij.openapi.vfs.VirtualFile
-import com.intellij.psi.PsiFile
-import com.intellij.psi.PsiManager
-import java.io.File
-import java.nio.file.Files
 import java.util.Base64
 
 class WriteBase64FileHandler(project: Project, data: Map<String, Any>) : WriteFileHandler(project, data) {
 
     override fun doWrite(vfsFile: VirtualFile?, doc: Document?, content: String) {
         vfsFile?.setBinaryContent(Base64.getDecoder().decode(content))
     }
-
-    override fun doCreate(ioFile: File, content: String): PsiFile {
-        Files.createFile(ioFile.toPath())
-        val vfsFile = VfsUtil.findFileByIoFile(ioFile, true)
-        doWrite(vfsFile!!, null, content)
-        return PsiManager.getInstance(project).findFile(vfsFile)!!
-    }
 }