feat: display AI Fix ⚡️ [IDE-580]

src/main/kotlin/io/snyk/plugin/ui/jcef/ApplyFixHandler.kt

@@ -0,0 +1,215 @@
+package io.snyk.plugin.ui.jcef
+
+import com.intellij.openapi.application.ApplicationManager
+import com.intellij.openapi.command.WriteCommandAction
+import com.intellij.openapi.fileEditor.FileDocumentManager
+import com.intellij.openapi.project.Project
+import com.intellij.openapi.vfs.LocalFileSystem
+import com.intellij.openapi.vfs.VirtualFile
+import com.intellij.ui.jcef.JBCefBrowserBase
+import com.intellij.ui.jcef.JBCefJSQuery
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.withContext
+import org.cef.browser.CefBrowser
+import org.cef.browser.CefFrame
+import org.cef.handler.CefLoadHandlerAdapter
+
+data class DiffPatch(
+    val originalFile: String,
+    val fixedFile: String,
+    val hunks: List<Hunk>
+)
+
+data class Hunk(
+    val startLineOriginal: Int,
+    val numLinesOriginal: Int,
+    val startLineFixed: Int,
+    val numLinesFixed: Int,
+    val changes: List<Change>
+)
+
+sealed class Change {
+    data class Addition(val line: String) : Change()
+    data class Deletion(val line: String) : Change()
+    data class Context(val line: String) : Change()  // Unchanged line for context
+}
+
+class ApplyFixHandler(private val project: Project) {
+
+    fun generateApplyFixCommand(jbCefBrowser: JBCefBrowserBase): CefLoadHandlerAdapter {
+        val applyFixQuery = JBCefJSQuery.create(jbCefBrowser)
+
+        applyFixQuery.addHandler { value ->
+            val params = value.split(":")
+            val filePath = params[0]  // Path to the file that needs to be patched
+            val patch = params[1]      // The patch we received from LS
+
+            println("[applyFixHandler] Received request to apply fix on file: $filePath")
+            println("[applyFixHandler] Patch to apply: $patch")
+
+
+            // Avoid blocking the UI thread
+            CoroutineScope(Dispatchers.IO).launch {
+                try {
+                    val success = applyPatchAndSave(project, filePath, patch)
+
+                    val script = """
+                        window.receiveApplyFixResponse($success);
+                    """.trimIndent()
+
+                    withContext(Dispatchers.Main) {
+                        jbCefBrowser.cefBrowser.executeJavaScript(script, jbCefBrowser.cefBrowser.url, 0)
+                    }
+                } catch (e: Exception) {
+                    e.printStackTrace()
+                }
+            }
+            return@addHandler JBCefJSQuery.Response("success")
+        }
+
+        return object : CefLoadHandlerAdapter() {
+            override fun onLoadEnd(browser: CefBrowser, frame: CefFrame, httpStatusCode: Int) {
+                if (frame.isMain) {
+                    val script = """
+                    (function() {
+                        if (window.applyFixQuery) {
+                            return;
+                        }
+                        window.applyFixQuery = function(value) { ${applyFixQuery.inject("value")} };
+                    })();
+                    """.trimIndent()
+                    browser.executeJavaScript(script, browser.url, 0)
+                }
+            }
+        }
+    }
+
+    private fun applyPatchAndSave(project: Project, filePath: String, patch: String): Boolean {
+        val virtualFile = findVirtualFile(filePath) ?: run {
+            println("[applyPatchAndSave] Virtual file not found for path: $filePath")
+            return false
+        }
+
+        println("[applyPatchAndSave] Found virtual file: $virtualFile")
+
+        val fileContent = ApplicationManager.getApplication().runReadAction<String?> {
+            val document = FileDocumentManager.getInstance().getDocument(virtualFile)
+            document?.text
+        } ?: run {
+            println("[applyPatchAndSave] Document not found or is null for virtual file: $filePath")
+            return false
+        }
+
+        println("[applyPatchAndSave] Initial file content: $fileContent")
+
+        val diffPatch = parseDiff(patch)
+        val patchedContent = applyPatch(fileContent, diffPatch)
+
+        println("[applyPatchAndSave] Patched content that will be written: $patchedContent")
+
+        // Apply the patch inside a WriteCommandAction
+        WriteCommandAction.runWriteCommandAction(project) {
+            val document = FileDocumentManager.getInstance().getDocument(virtualFile)
+            document?.let {
+                it.setText(patchedContent)
+
+                println("[applyPatchAndSave] Content after applying patch: ${it.text}")
+
+                FileDocumentManager.getInstance().saveDocument(it)
+                println("[applyPatchAndSave] Patch applied successfully!")
+            } ?: run {
+                println("[applyPatchAndSave] Failed to find document for saving patched content.")
+            }
+        }
+
+        return true
+    }
+
+    fun applyPatch(fileContent: String, diffPatch: DiffPatch): String {
+        val lines = fileContent.lines().toMutableList()
+
+        for (hunk in diffPatch.hunks) {
+            var originalLineIndex = hunk.startLineOriginal - 1  // Convert to 0-based index
+
+            for (change in hunk.changes) {
+                when (change) {
+                    is Change.Addition -> {
+                        lines.add(originalLineIndex, change.line)
+                        originalLineIndex++
+                    }
+                    is Change.Deletion -> {
+                        if (originalLineIndex < lines.size && lines[originalLineIndex].trim() == change.line) {
+                            lines.removeAt(originalLineIndex)
+                        }
+                    }
+                    is Change.Context -> {
+                        originalLineIndex++  // Move past unchanged context lines
+                    }
+                }
+            }
+        }
+        return lines.joinToString("\n")
+    }
+
+    fun parseDiff(diff: String): DiffPatch {
+        val lines = diff.lines()
+        val originalFile = lines.first { it.startsWith("---") }.substringAfter("--- ")
+        val fixedFile = lines.first { it.startsWith("+++") }.substringAfter("+++ ")
+
+        val hunks = mutableListOf<Hunk>()
+        var currentHunk: Hunk? = null
+        val changes = mutableListOf<Change>()
+
+        for (line in lines) {
+            when {
+                line.startsWith("@@") -> {
+                    // Parse hunk header (e.g., @@ -4,9 +4,14 @@)
+                    val hunkHeader = line.substringAfter("@@ ").substringBefore(" @@").split(" ")
+                    val original = hunkHeader[0].substring(1).split(",")
+                    val fixed = hunkHeader[1].substring(1).split(",")
+
+                    val startLineOriginal = original[0].toInt()
+                    val numLinesOriginal = original.getOrNull(1)?.toInt() ?: 1
+                    val startLineFixed = fixed[0].toInt()
+                    val numLinesFixed = fixed.getOrNull(1)?.toInt() ?: 1
+
+                    if (currentHunk != null) {
+                        hunks.add(currentHunk.copy(changes = changes.toList()))
+                        changes.clear()
+                    }
+                    currentHunk = Hunk(
+                        startLineOriginal = startLineOriginal,
+                        numLinesOriginal = numLinesOriginal,
+                        startLineFixed = startLineFixed,
+                        numLinesFixed = numLinesFixed,
+                        changes = emptyList()
+                    )
+                }
+                line.startsWith("---") || line.startsWith("+++") -> {
+                    // Skip file metadata lines (--- and +++)
+                    continue
+                }
+                line.startsWith("-") -> changes.add(Change.Deletion(line.substring(1).trim()))
+                line.startsWith("+") -> changes.add(Change.Addition(line.substring(1).trim()))
+                else -> changes.add(Change.Context(line.trim()))
+            }
+        }
+
+        // Add the last hunk
+        if (currentHunk != null) {
+            hunks.add(currentHunk.copy(changes = changes.toList()))
+        }
+
+        return DiffPatch(
+            originalFile = originalFile,
+            fixedFile = fixedFile,
+            hunks = hunks
+        )
+    }
+
+    private fun findVirtualFile(filePath: String): VirtualFile? {
+        return LocalFileSystem.getInstance().findFileByPath(filePath)
+    }
+}

src/main/kotlin/io/snyk/plugin/ui/jcef/GenerateAIFixHandler.kt

@@ -0,0 +1,78 @@
+package io.snyk.plugin.ui.jcef
+
+import com.google.gson.Gson
+import com.intellij.openapi.project.Project
+import com.intellij.ui.jcef.JBCefBrowserBase
+import com.intellij.ui.jcef.JBCefJSQuery
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.withContext
+import org.cef.browser.CefBrowser
+import org.cef.browser.CefFrame
+import org.cef.handler.CefLoadHandlerAdapter
+import snyk.common.lsp.LanguageServerWrapper
+
+class GenerateAIFixHandler(private val project: Project) {
+
+    fun generateAIFixCommand(jbCefBrowser: JBCefBrowserBase): CefLoadHandlerAdapter {
+        val aiFixQuery = JBCefJSQuery.create(jbCefBrowser)
+
+        aiFixQuery.addHandler { value ->
+            val params = value.split(":")
+            val folderURI = params[0]
+            val fileURI = params[1]
+            val issueID = params[2]
+
+            // Avoids blocking the UI thread
+            CoroutineScope(Dispatchers.IO).launch {
+                try {
+                    val responseDiff: List<LanguageServerWrapper.Fix> =
+                        LanguageServerWrapper.getInstance().sendCodeFixDiffsCommand(folderURI, fileURI, issueID)
+
+                    val script = """
+                        window.receiveAIFixResponse(${Gson().toJson(responseDiff)});
+                    """.trimIndent()
+
+                    withContext(Dispatchers.Main) {
+                        jbCefBrowser.cefBrowser.executeJavaScript(script, jbCefBrowser.cefBrowser.url, 0)
+                        JBCefJSQuery.Response("success")
+                    }
+                } catch (e: Exception) {
+                    e.printStackTrace()
+                }
+            }
+            return@addHandler JBCefJSQuery.Response("success")
+        }
+
+        return object : CefLoadHandlerAdapter() {
+            override fun onLoadEnd(browser: CefBrowser, frame: CefFrame, httpStatusCode: Int) {
+                if (frame.isMain) {
+                    val script = """
+                    (function() {
+                        if (window.aiFixQuery) {
+                            return;
+                        }
+                        window.aiFixQuery = function(value) { ${aiFixQuery.inject("value")} };
+
+                        const aiFixButton = document.getElementById('generate-ai-fix');
+                        const retryFixButton = document.getElementById('retry-generate-fix');
+                        const issueId = aiFixButton.getAttribute('issue-id');
+                        const folderPath = aiFixButton.getAttribute('folder-path');
+                        const filePath = aiFixButton.getAttribute('file-path');
+
+                        aiFixButton.addEventListener('click', () => {
+                            window.aiFixQuery(folderPath + ":" + filePath + ":" + issueId);
+                        });
+
+                        retryFixButton.addEventListener('click', () => {
+                            window.aiFixQuery(folderPath + ":" + filePath + ":" + issueId);
+                        });
+                    })();
+                    """
+                    browser.executeJavaScript(script, browser.url, 0)
+                }
+            }
+        }
+    }
+}

src/main/kotlin/io/snyk/plugin/ui/jcef/ThemeBasedStylingGenerator.kt

@@ -3,8 +3,10 @@ package io.snyk.plugin.ui.jcef
 import com.intellij.openapi.editor.colors.ColorKey
 import com.intellij.openapi.editor.colors.EditorColors
 import com.intellij.openapi.editor.colors.EditorColorsManager
+import com.intellij.ui.JBColor
 import com.intellij.ui.jcef.JBCefBrowserBase
 import com.intellij.util.ui.JBUI
+import com.intellij.util.ui.PlatformColors
 import com.intellij.util.ui.UIUtil
 import org.cef.browser.CefBrowser
 import org.cef.browser.CefFrame
@@ -63,8 +65,8 @@ class ThemeBasedStylingGenerator {
                     val textColor = toCssHex(JBUI.CurrentTheme.Label.foreground())
                     val linkColor = toCssHex(JBUI.CurrentTheme.Link.Foreground.ENABLED)
                     val dataFlowColor = toCssHex(baseColor)
+                    val editorColor = toCssHex(baseColor)
                     val borderColor = toCssHex(JBUI.CurrentTheme.CustomFrameDecorations.separatorForeground())
-                    val editorColor = toCssHex(UIUtil.getTextFieldBackground())
                     val labelColor = toCssHex(JBUI.CurrentTheme.Label.foreground())
 
                     val globalScheme = EditorColorsManager.getInstance().globalScheme
@@ -92,7 +94,11 @@ class ThemeBasedStylingGenerator {
                                 '--border-color': "$borderColor",
                                 '--editor-color': "$editorColor",
                                 '--label-color': "'$labelColor'",
-                                '--vulnerability-overview-pre-background-color': "${toCssHex(codeTagBgColor)}",
+                                '--container-background-color': "${toCssHex(codeTagBgColor)}",
+                                '--generated-ai-fix-button-background-color': "$linkColor",
+                                '--dark-button-border-default': "$borderColor",
+                                '--dark-button-default': "$linkColor",
+
                             };
                             for (let [property, value] of Object.entries(properties)) {
                                 document.documentElement.style.setProperty(property, value);

src/main/kotlin/io/snyk/plugin/ui/jcef/Utils.kt

@@ -19,7 +19,7 @@ object JCEFUtils {
         val cefClient = JBCefApp.getInstance().createClient()
         cefClient.setProperty("JS_QUERY_POOL_SIZE", 1)
         val jbCefBrowser =
-            JBCefBrowserBuilder().setClient(cefClient).setEnableOpenDevToolsMenuItem(false)
+            JBCefBrowserBuilder().setClient(cefClient).setEnableOpenDevToolsMenuItem(true)
                 .setMouseWheelEventEnable(true).build()
         jbCefBrowser.setOpenLinksInExternalBrowser(true)