feat: introduce whole diff edit tool for ProxyAI provider

src/main/kotlin/ee/carlrobert/codegpt/agent/AgentFactory.kt

@@ -130,7 +130,8 @@ object AgentFactory {
             sessionId,
             approveToolCall,
             selected,
-            approvalBashHandler(approveToolCall)
+            approvalBashHandler(approveToolCall),
+            provider
         )
 
         return AIAgent.Companion(
@@ -295,7 +296,8 @@ object AgentFactory {
                 sessionId,
                 approveToolCall,
                 selectedTools,
-                approvalBashHandler(approveToolCall)
+                approvalBashHandler(approveToolCall),
+                provider
             ),
             installFeatures = {
                 installFeatures()
@@ -363,7 +365,8 @@ object AgentFactory {
                 sessionId,
                 approveToolCall = approveToolCall,
                 selected = selectedTools,
-                bashConfirmationHandler = { ShellCommandConfirmation.Approved }
+                bashConfirmationHandler = { ShellCommandConfirmation.Approved },
+                provider = provider
             ),
             installFeatures = {
                 installFeatures()
@@ -520,21 +523,30 @@ object AgentFactory {
         sessionId: String,
         approveToolCall: (suspend (name: String, details: String) -> Boolean)?,
         selected: Set<SubagentTool>,
-        bashConfirmationHandler: BashCommandConfirmationHandler
+        bashConfirmationHandler: BashCommandConfirmationHandler,
+        provider: ServiceType
     ): ToolRegistry {
         return ToolRegistry.Companion {
             if (SubagentTool.READ in selected) tool(ReadTool(project))
             if (SubagentTool.EDIT in selected) {
-                tool(
-                    ConfirmingEditTool(EditTool(project)) { name, details ->
-                        approveToolCall?.invoke(name, details) ?: true
-                    }
-                )
+                if (provider == ServiceType.PROXYAI) {
+                    tool(
+                        ConfirmingProxyAIEditTool(ProxyAIEditTool(project), project) { request ->
+                            approveToolCall?.invoke("Edit", request.details) ?: false
+                        }
+                    )
+                } else {
+                    tool(
+                        ConfirmingEditTool(EditTool(project)) { name, details ->
+                            approveToolCall?.invoke(name, details) ?: false
+                        }
+                    )
+                }
             }
             if (SubagentTool.WRITE in selected) {
                 tool(
                     ConfirmingWriteTool(WriteTool(project)) { name, details ->
-                        approveToolCall?.invoke(name, details) ?: true
+                        approveToolCall?.invoke(name, details) ?: false
                     }
                 )
             }

src/main/kotlin/ee/carlrobert/codegpt/agent/ProxyAIAgent.kt

@@ -83,7 +83,7 @@ object ProxyAIAgent {
         val projectInstructions = searchForInstructions(project.basePath)
         val executor = AgentFactory.createExecutor(provider, events)
         val pendingMessageQueue = pendingMessages.getOrPut(sessionId) { ArrayDeque() }
-        val toolRegistry = createToolRegistry(project, events, sessionId)
+        val toolRegistry = createToolRegistry(project, events, sessionId, provider)
         val agentModel = service<ModelSelectionService>().getAgentModel()
         val agent = AIAgent(
             promptExecutor = executor,
@@ -225,26 +225,38 @@ object ProxyAIAgent {
     private fun createToolRegistry(
         project: Project,
         events: AgentEvents,
-        sessionId: String
+        sessionId: String,
+        provider: ServiceType
     ): ToolRegistry {
         return ToolRegistry {
             tool(ReadTool(project))
-            tool(
-                ConfirmingEditTool(EditTool(project)) { name, details ->
-                    try {
-                        events.approveToolCall(
-                            ToolApprovalRequest(
-                                if (name.equals("Edit", true)
-                                ) ToolApprovalType.EDIT else ToolApprovalType.GENERIC,
-                                "Allow $name?",
-                                details
-                            )
+            val approveHandler: suspend (String, String) -> Boolean = { name, details ->
+                try {
+                    events.approveToolCall(
+                        ToolApprovalRequest(
+                            if (name.equals("Edit", true)
+                            ) ToolApprovalType.EDIT else ToolApprovalType.GENERIC,
+                            "Allow $name?",
+                            details
                         )
-                    } catch (_: Exception) {
-                        false
-                    }
+                    )
+                } catch (_: Exception) {
+                    false
                 }
-            )
+            }
+            if (provider == ServiceType.PROXYAI) {
+                tool(
+                    ConfirmingProxyAIEditTool(ProxyAIEditTool(project), project) { request ->
+                        try {
+                            events.approveToolCall(request)
+                        } catch (_: Exception) {
+                            false
+                        }
+                    }
+                )
+            } else {
+                tool(ConfirmingEditTool(EditTool(project), approveHandler))
+            }
             tool(
                 ConfirmingWriteTool(WriteTool(project)) { name, details ->
                     try {

src/main/kotlin/ee/carlrobert/codegpt/agent/tools/ConfirmingTools.kt

@@ -1,6 +1,12 @@
 package ee.carlrobert.codegpt.agent.tools
 
 import ai.koog.agents.core.tools.Tool
+import com.intellij.openapi.project.Project
+import ee.carlrobert.codegpt.toolwindow.agent.ui.approval.ToolApprovalRequest
+import ee.carlrobert.codegpt.toolwindow.agent.ui.approval.ToolApprovalType
+import ee.carlrobert.codegpt.toolwindow.agent.ui.approval.ProxyAIEditPayload
+import ee.carlrobert.codegpt.util.EditorUtil
+import ee.carlrobert.codegpt.util.file.FileUtil
 
 class ConfirmingWriteTool(
     private val delegate: WriteTool,
@@ -52,3 +58,63 @@ class ConfirmingEditTool(
         return delegate.execute(args)
     }
 }
+
+class ConfirmingProxyAIEditTool(
+    private val delegate: ProxyAIEditTool,
+    private val project: Project,
+    private val approve: suspend (request: ToolApprovalRequest) -> Boolean
+) : Tool<ProxyAIEditTool.Args, ProxyAIEditTool.Result>(
+    argsSerializer = ProxyAIEditTool.Args.serializer(),
+    resultSerializer = ProxyAIEditTool.Result.serializer(),
+    name = delegate.name,
+    description = delegate.descriptor.description
+) {
+
+    override suspend fun execute(args: ProxyAIEditTool.Args): ProxyAIEditTool.Result {
+        val result = delegate.execute(args)
+        
+        if (result !is ProxyAIEditTool.Result.Success) {
+            return result
+        }
+        
+        val payload = ProxyAIEditPayload(
+            filePath = result.filePath,
+            updateSnippet = args.updateSnippet,
+            originalContent = result.originalCode,
+            updatedContent = result.updatedCode
+        )
+        
+        val approved = approve(
+            ToolApprovalRequest(
+                ToolApprovalType.EDIT,
+                "Allow Edit?",
+                args.shortDescription,
+                payload
+            )
+        )
+        
+        if (!approved) {
+            return ProxyAIEditTool.Result.Error(
+                filePath = args.filePath,
+                error = "User rejected edit operation"
+            )
+        }
+        
+        val normalizedPath = result.filePath.replace("\\", "/")
+        val virtualFile = FileUtil.findVirtualFile(normalizedPath)
+            ?: return ProxyAIEditTool.Result.Error(
+                filePath = result.filePath,
+                error = "File not found in IntelliJ VFS: ${result.filePath}"
+            )
+
+        val written = EditorUtil.writeDocumentContent(project, virtualFile, result.updatedCode)
+        if (!written) {
+            return ProxyAIEditTool.Result.Error(
+                filePath = args.filePath,
+                error = "Failed to write changes to document"
+            )
+        }
+        
+        return result
+    }
+}

src/main/kotlin/ee/carlrobert/codegpt/agent/tools/ProxyAIEditTool.kt

@@ -0,0 +1,184 @@
+package ee.carlrobert.codegpt.agent.tools
+
+import ai.koog.agents.core.tools.Tool
+import ai.koog.agents.core.tools.annotations.LLMDescription
+import com.intellij.openapi.application.runReadAction
+import com.intellij.openapi.fileEditor.FileDocumentManager
+import com.intellij.openapi.project.Project
+import ee.carlrobert.codegpt.completions.CompletionClientProvider
+import ee.carlrobert.codegpt.settings.ProxyAISettingsService
+import ee.carlrobert.codegpt.settings.service.FeatureType
+import ee.carlrobert.codegpt.settings.service.ModelSelectionService
+import ee.carlrobert.codegpt.tokens.truncateToolResult
+import ee.carlrobert.codegpt.util.file.FileUtil
+import ee.carlrobert.llm.client.codegpt.request.AutoApplyRequest
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.withContext
+import kotlinx.serialization.SerialName
+import kotlinx.serialization.Serializable
+
+class ProxyAIEditTool(private val project: Project) :
+    Tool<ProxyAIEditTool.Args, ProxyAIEditTool.Result>(
+        argsSerializer = Args.serializer(),
+        resultSerializer = Result.serializer(),
+        name = "Edit",
+        description = """
+        Generates code changes using the ProxyAI auto-apply endpoint.
+
+        Provide an update snippet that describes the edits with surrounding context.
+        If there are multiple changes in the same file, include them ALL in a single update_snippet
+        (do NOT make multiple Edit tool calls for one file).
+        The update snippet must be included in update_snippet using the following format:
+
+        <|update_snippet|>
+        // ... existing code ...
+        [UPDATED CODE SNIPPET 1]
+        // ... existing code ...
+        [UPDATED CODE SNIPPET 2]
+        // ... existing code ...
+        <|/update_snippet|>
+
+        Example (multiple edits in one file):
+        <|update_snippet|>
+        function add(a, b) {
+          // ... existing code ...
+          return a + b + 1;
+        }
+        // ... existing code ...
+        const label = "Sum: " + add(1, 2);
+        // ... existing code ...
+        <|/update_snippet|>
+
+        Parameters:
+        - file_path: Absolute path to the file
+        - update_snippet: The update snippet in the required format above
+        - short_description: Short description of the requested edit.
+
+        Notes:
+        - The file must exist and be writable
+        - The tool will fail if no changes are produced
+    """.trimIndent()
+    ) {
+
+    @Serializable
+    data class Args(
+        @property:LLMDescription(
+            "The absolute path to the file to edit. Must be an absolute path, not a relative path."
+        )
+        @SerialName("file_path")
+        val filePath: String,
+        @property:LLMDescription(
+            "Update snippet containing the edits, formatted with <|update_snippet|> markers."
+        )
+        @SerialName("update_snippet")
+        val updateSnippet: String,
+        @property:LLMDescription(
+            "Short description of the edit."
+        )
+        @SerialName("short_description")
+        val shortDescription: String
+    )
+
+    @Serializable
+    sealed class Result {
+        @Serializable
+        data class Success(
+            val filePath: String,
+            val originalCode: String,
+            val updatedCode: String,
+        ) : Result()
+
+        @Serializable
+        data class Error(
+            val filePath: String,
+            val error: String
+        ) : Result()
+    }
+
+    override suspend fun execute(args: Args): Result {
+        return try {
+            val svc = project.getService(ProxyAISettingsService::class.java)
+            if (svc.isPathIgnored(args.filePath)) {
+                return Result.Error(
+                    filePath = args.filePath,
+                    error = ".proxyai ignore rules block editing this path"
+                )
+            }
+
+            FileUtil.validateFileForEdit(args.filePath).getOrElse { error ->
+                return Result.Error(
+                    filePath = args.filePath,
+                    error = error.message ?: "File validation failed"
+                )
+            }
+
+            if (args.updateSnippet.isBlank()) {
+                return Result.Error(
+                    filePath = args.filePath,
+                    error = "update_snippet is empty or missing"
+                )
+            }
+
+            val normalizedPath = args.filePath.replace("\\", "/")
+            val virtualFile = FileUtil.findVirtualFile(normalizedPath)
+                ?: return Result.Error(
+                    filePath = args.filePath,
+                    error = "File not found in IntelliJ VFS: ${args.filePath}"
+                )
+
+            val document = withContext(Dispatchers.Default) {
+                runReadAction { FileDocumentManager.getInstance().getDocument(virtualFile) }
+            } ?: return Result.Error(
+                filePath = args.filePath,
+                error = "Cannot get document for file: ${args.filePath}"
+            )
+
+            val original = document.text
+
+            val model =
+                ModelSelectionService.getInstance().getModelForFeature(FeatureType.AUTO_APPLY)
+            val updated = withContext(Dispatchers.IO) {
+                CompletionClientProvider.getCodeGPTClient()
+                    .applyChanges(AutoApplyRequest(model, original, args.updateSnippet))
+                    .mergedCode
+            }
+
+            when {
+                updated.isNullOrBlank() -> Result.Error(
+                    filePath = args.filePath,
+                    error = "Auto-apply did not return updated content"
+                )
+
+                updated == original -> Result.Error(
+                    filePath = args.filePath,
+                    error = "No changes produced by auto-apply"
+                )
+
+                else -> Result.Success(
+                    filePath = args.filePath,
+                    originalCode = original,
+                    updatedCode = updated,
+                )
+            }
+        } catch (e: Exception) {
+            Result.Error(
+                filePath = args.filePath,
+                error = "Failed to edit file: ${e.message}"
+            )
+        }
+    }
+
+    override fun encodeResultToString(result: Result): String = when (result) {
+        is Result.Success -> {
+            buildString {
+                appendLine("Successfully generated edits for '${result.filePath}'")
+            }
+                .trimEnd()
+                .truncateToolResult()
+        }
+
+        is Result.Error -> {
+            ("Error editing file '${result.filePath}': ${result.error}").truncateToolResult()
+        }
+    }
+}

src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/AgentApprovalManager.kt

@@ -30,10 +30,11 @@ import com.intellij.util.ui.JBUI
 import com.intellij.util.ui.components.BorderLayoutPanel
 import ee.carlrobert.codegpt.CodeGPTBundle
 import ee.carlrobert.codegpt.Icons
-import ee.carlrobert.codegpt.agent.tools.EditTool
+import ee.carlrobert.codegpt.agent.tools.EditArgsSnapshot
 import ee.carlrobert.codegpt.agent.tools.WriteTool
 import ee.carlrobert.codegpt.toolwindow.agent.ui.renderer.applyStringReplacement
 import ee.carlrobert.codegpt.toolwindow.agent.ui.renderer.getFileContentWithFallback
+import ee.carlrobert.codegpt.util.UpdateSnippetUtil
 import kotlinx.coroutines.CompletableDeferred
 import java.awt.BorderLayout
 import java.awt.FlowLayout
@@ -85,34 +86,47 @@ class AgentApprovalManager(
         attachApprovalActionsAndShowDiff(request, decision)
     }
 
-    fun openEditApprovalDiff(args: EditTool.Args, decision: CompletableDeferred<Boolean>) {
+    fun openEditApprovalDiff(
+        args: EditArgsSnapshot,
+        decision: CompletableDeferred<Boolean>,
+        proposedContent: String? = null
+    ) {
         val path = try {
             Paths.get(args.filePath).normalize().toString()
         } catch (_: Exception) {
             args.filePath
         }
 
-        val vf = LocalFileSystem.getInstance().refreshAndFindFileByPath(path)
-        val factory = DiffContentFactory.getInstance()
+        ApplicationManager.getApplication().executeOnPooledThread {
+            val vf = LocalFileSystem.getInstance().refreshAndFindFileByPath(path)
+            val factory = DiffContentFactory.getInstance()
 
-        val current = getFileContentWithFallback(path)
-        val proposed =
-            applyStringReplacement(current, args.oldString, args.newString, args.replaceAll)
-
-        val left = if (vf != null) factory.create(project, vf) else factory.create(project, current)
-        val rightDoc =
-            EditorFactory.getInstance().createDocument(convertLineSeparators(proposed)).apply {
-                setReadOnly(true)
+            val current = getFileContentWithFallback(path)
+            val proposed = proposedContent ?: run {
+                val rawSnippet = if (args.newString.isNotBlank()) args.newString else args.oldString
+                if (UpdateSnippetUtil.containsMarkers(rawSnippet)) {
+                    current
+                } else {
+                    applyStringReplacement(current, args.oldString, args.newString, args.replaceAll)
+                }
             }
-        val right = if (vf != null) factory.create(project, rightDoc, vf)
-        else factory.create(project, rightDoc, FileTypes.PLAIN_TEXT)
 
-        val request = SimpleDiffRequest(
-            "Edit File",
-            listOf(left, right),
-            listOf("Current", "Proposed")
-        )
-        attachApprovalActionsAndShowDiff(request, decision)
+            val left = if (vf != null) factory.create(project, vf) else factory.create(project, current)
+            val rightDoc =
+                EditorFactory.getInstance().createDocument(convertLineSeparators(proposed)).apply {
+                    setReadOnly(true)
+                }
+            val right = if (vf != null) factory.create(project, rightDoc, vf)
+            else factory.create(project, rightDoc, FileTypes.PLAIN_TEXT)
+
+            val request = SimpleDiffRequest(
+                "Edit File",
+                listOf(left, right),
+                listOf("Current", "Proposed")
+            )
+
+            runInEdt { attachApprovalActionsAndShowDiff(request, decision) }
+        }
     }
 
     private fun attachApprovalActionsAndShowDiff(
@@ -287,6 +301,7 @@ class AgentApprovalManager(
             val fileEditor = manager.getSelectedEditor(diffFile)
                 ?: manager.getEditors(diffFile).firstOrNull()
             val target = fileEditor?.component ?: return@runInEdt
+            if (!target.isShowing) return@runInEdt
 
             val popup = JBPopupFactory.getInstance()
                 .createComponentPopupBuilder(panel, null)
@@ -314,6 +329,7 @@ class AgentApprovalManager(
 
                 private fun relocate() {
                     if (popup.isDisposed) return
+                    if (!target.isShowing) return
                     val loc = target.locationOnScreen
                     val nx = loc.x + target.width - size.width - margin
                     val ny = loc.y + target.height - size.height - margin

src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/AgentEventHandler.kt

@@ -16,8 +16,19 @@ import ee.carlrobert.codegpt.conversations.message.TokenUsage
 import ee.carlrobert.codegpt.settings.service.ServiceType
 import ee.carlrobert.codegpt.toolwindow.agent.ui.*
 import ee.carlrobert.codegpt.toolwindow.agent.ui.approval.*
+import ee.carlrobert.codegpt.toolwindow.agent.ui.descriptor.Badge
+import ee.carlrobert.codegpt.toolwindow.agent.ui.descriptor.ToolKind
+import ee.carlrobert.codegpt.toolwindow.agent.ui.renderer.ChangeColors
+import ee.carlrobert.codegpt.toolwindow.agent.ui.renderer.applyStringReplacement
+import ee.carlrobert.codegpt.toolwindow.agent.ui.renderer.diffBadgeText
+import ee.carlrobert.codegpt.toolwindow.agent.ui.renderer.getFileContentWithFallback
+import ee.carlrobert.codegpt.toolwindow.agent.ui.renderer.lineDiffStats
 import ee.carlrobert.codegpt.toolwindow.chat.ui.ChatMessageResponseBody
 import ee.carlrobert.codegpt.toolwindow.chat.ui.ChatToolWindowScrollablePanel
+import ee.carlrobert.codegpt.util.UpdateSnippetUtil
+import com.intellij.openapi.application.runReadAction
+import com.intellij.openapi.fileEditor.FileDocumentManager
+import com.intellij.openapi.vfs.LocalFileSystem
 import ee.carlrobert.codegpt.ui.textarea.UserInputPanel
 import kotlinx.coroutines.*
 import java.awt.Component
@@ -60,7 +71,7 @@ class AgentEventHandler(
     private var lastWriteArgs: WriteTool.Args? = null
 
     @Volatile
-    private var lastEditArgs: EditTool.Args? = null
+    private var lastEditArgs: EditArgsSnapshot? = null
 
     private val approvalQueue: ArrayDeque<ApprovalRequest> = ArrayDeque()
 
@@ -241,16 +252,33 @@ class AgentEventHandler(
 
         if (isWrite || isEdit) {
             val deferred = CompletableDeferred<Boolean>()
+            val resolvedRequest = if (isEdit && request.payload == null) {
+                val payload = lastEditArgs?.let { args ->
+                    EditPayload(
+                        filePath = args.filePath,
+                        oldString = args.oldString,
+                        newString = args.newString,
+                        replaceAll = args.replaceAll,
+                        proposedContent = null
+                    )
+                }
+                if (payload != null) request.copy(payload = payload) else request
+            } else {
+                request
+            }
             runInEdt {
-                approvalQueue.addLast(ApprovalRequest(request, deferred))
+                approvalQueue.addLast(ApprovalRequest(resolvedRequest, deferred))
                 maybeShowNextApproval()
             }
 
             lastWriteArgs?.let {
                 if (isWrite) agentApprovalManager.openWriteApprovalDiff(it, deferred)
             }
-            lastEditArgs?.let {
-                if (isEdit) agentApprovalManager.openEditApprovalDiff(it, deferred)
+            lastEditArgs?.let { args ->
+                if (isEdit) {
+                    val proposed = (resolvedRequest.payload as? EditPayload)?.proposedContent
+                    agentApprovalManager.openEditApprovalDiff(args, deferred, proposed)
+                }
             }
             return deferred.await()
         }
@@ -277,7 +305,7 @@ class AgentEventHandler(
     override fun onTextReceived(text: String) {
         runInEdt {
             val cleanedText =
-                text.replace(Regex("<think>.*?</think>", RegexOption.DOT_MATCHES_ALL), "")
+                text.replace(Regex("<tool_call>.*?<tool_call>", RegexOption.DOT_MATCHES_ALL), "")
             currentResponseBody?.updateMessage(cleanedText)
             scrollablePanel.update()
             scrollablePanel.scrollToBottom()
@@ -312,19 +340,12 @@ class AgentEventHandler(
 
             else -> {
                 when (args) {
-                    is EditTool.Args -> {
-                        lastEditArgs = args
-                        val originalContent = runCatching {
-                            java.io.File(args.filePath).readText()
-                        }.getOrNull() ?: ""
-                        project.service<RollbackService>()
-                            .trackEdit(sessionId, args.filePath, args, originalContent)
+                    is EditTool.Args, is ProxyAIEditTool.Args, is EditArgsSnapshot -> {
+                        trackEditOperation(args)
                     }
 
                     is WriteTool.Args -> {
-                        lastWriteArgs = args
-                        project.service<RollbackService>()
-                            .trackWrite(sessionId, args.filePath, args)
+                        trackWriteOperation(args)
                     }
                 }
 
@@ -415,9 +436,10 @@ class AgentEventHandler(
                     RunEntry.WriteEntry(cid, parentId, args, null)
                 }
 
-                is EditTool.Args -> {
-                    lastEditArgs = args
-                    RunEntry.EditEntry(cid, parentId, args, null)
+                is EditTool.Args, is ProxyAIEditTool.Args, is EditArgsSnapshot -> {
+                    val snapshot = snapshotFromEditArgs(args) ?: return@runInEdt
+                    lastEditArgs = snapshot
+                    RunEntry.EditEntry(cid, parentId, snapshot, null)
                 }
 
                 is TaskTool.Args -> RunEntry.TaskEntry(cid, parentId, args, null)
@@ -545,6 +567,10 @@ class AgentEventHandler(
             }
         }
 
+        if (next.model.type == ToolApprovalType.EDIT) {
+            updateEditToolCardPreview(next.model)
+        }
+
         runCatching {
             project.service<AgentToolWindowContentManager>()
                 .setTabStatus(sessionId, AgentToolWindowTabbedPane.TabStatus.APPROVAL)
@@ -588,6 +614,56 @@ class AgentEventHandler(
         approvalContainer.repaint()
     }
 
+    private fun updateEditToolCardPreview(request: ToolApprovalRequest) {
+        val payload = request.payload ?: return
+        val (path, before, after) = when (payload) {
+            is ProxyAIEditPayload -> Triple(
+                payload.filePath,
+                payload.originalContent,
+                payload.updatedContent
+            )
+            is EditPayload -> {
+                val rawSnippet = payload.newString.ifBlank { payload.oldString }
+                if (UpdateSnippetUtil.containsMarkers(rawSnippet)) return
+                val currentContent = getFileContentWithFallback(payload.filePath)
+                val proposed = payload.proposedContent ?: applyStringReplacement(
+                    currentContent,
+                    payload.oldString,
+                    payload.newString,
+                    payload.replaceAll
+                )
+                Triple(payload.filePath, currentContent, proposed)
+            }
+            else -> return
+        }
+
+        val (inserted, deleted, changed) = lineDiffStats(before, after)
+        val texts = diffBadgeText(inserted, deleted, changed)
+        val diffBadges = listOf(
+            Badge(texts.inserted, ChangeColors.inserted),
+            Badge(texts.deleted, ChangeColors.deleted),
+            Badge(texts.changed, ChangeColors.modified)
+        )
+
+        val card = mainToolCards.values.firstOrNull { candidate ->
+            val descriptor = candidate.getDescriptor()
+            descriptor.kind == ToolKind.EDIT && descriptor.fileLink?.path == path
+        } ?: return
+
+        card.updateDescriptor { descriptor ->
+            val nonDiffBadges = descriptor.secondaryBadges.filterNot { isDiffBadge(it) }
+            descriptor.copy(
+                secondaryBadges = nonDiffBadges + diffBadges,
+                summary = null
+            )
+        }
+    }
+
+    private fun isDiffBadge(badge: Badge): Boolean {
+        val text = badge.text
+        return text.startsWith("[+") || text.startsWith("[-") || text.startsWith("[~")
+    }
+
     private fun maybeShowNextQuestion() {
         if (currentApproval != null || currentQuestion != null) return
         val next = questionQueue.pollFirst() ?: return
@@ -691,4 +767,26 @@ class AgentEventHandler(
         approvalQueue.clear()
         subagentViewHolders.clear()
     }
+
+    private fun trackEditOperation(args: Any?) {
+        val snapshot = snapshotFromEditArgs(args) ?: return
+        lastEditArgs = snapshot
+        val normalizedPath = snapshot.filePath.replace("\\", "/")
+        val originalContent = runCatching {
+            val vf = LocalFileSystem.getInstance().findFileByPath(normalizedPath)
+            val documentText = vf?.let { file ->
+                runReadAction { FileDocumentManager.getInstance().getDocument(file)?.text }
+            }
+            documentText ?: java.io.File(normalizedPath).readText()
+        }.getOrNull() ?: ""
+        project.service<RollbackService>()
+            .trackEdit(sessionId, normalizedPath, snapshot, originalContent)
+    }
+
+    private fun trackWriteOperation(args: WriteTool.Args) {
+        lastWriteArgs = args
+        val normalizedPath = args.filePath.replace("\\", "/")
+        project.service<RollbackService>()
+            .trackWrite(sessionId, normalizedPath, args)
+    }
 }

src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/ui/approval/EditApprovalPanel.kt

@@ -17,14 +17,18 @@ import com.intellij.util.ui.JBUI
 import com.intellij.util.ui.components.BorderLayoutPanel
 import ee.carlrobert.codegpt.toolwindow.agent.ui.renderer.ChangeColors
 import ee.carlrobert.codegpt.toolwindow.agent.ui.renderer.applyStringReplacement
+import ee.carlrobert.codegpt.toolwindow.agent.ui.renderer.diffBadgeText
 import ee.carlrobert.codegpt.toolwindow.agent.ui.renderer.getFileContentWithFallback
 import ee.carlrobert.codegpt.toolwindow.agent.ui.renderer.lineDiffStats
 import ee.carlrobert.codegpt.toolwindow.chat.editor.ResponseEditorPanel
 import ee.carlrobert.codegpt.toolwindow.chat.parser.ReplaceWaiting
+import ee.carlrobert.codegpt.util.UpdateSnippetUtil
 import java.awt.BorderLayout
+import java.awt.FlowLayout
 import java.nio.file.Paths
 import javax.swing.BorderFactory
 import javax.swing.JComponent
+import javax.swing.JPanel
 import ee.carlrobert.codegpt.toolwindow.chat.editor.factory.EditorFactory as ChatEditorFactory
 
 class EditApprovalPanel(
@@ -48,9 +52,12 @@ class EditApprovalPanel(
             )
         )
 
-        val payload = request.payload as? EditPayload
-        val filePath = payload?.filePath ?: ""
-        val diffComponent = createInlineDiffComponent(payload)
+        val filePath = when (val payload = request.payload) {
+            is EditPayload -> payload.filePath
+            is ProxyAIEditPayload -> payload.filePath
+            else -> ""
+        }
+        val diffComponent = createInlineDiffComponent()
 
         add(
             panel {
@@ -75,15 +82,7 @@ class EditApprovalPanel(
                         cell(link).gap(RightGap.SMALL)
                         diffCounts?.let { (ins, del, changed) ->
                             if (ins + del + changed > 0) {
-                                cell(colorLabel("+${ins}", ChangeColors.inserted).apply {
-                                    font = JBUI.Fonts.smallFont()
-                                })
-                                cell(colorLabel("-${del}", ChangeColors.deleted).apply {
-                                    font = JBUI.Fonts.smallFont()
-                                })
-                                cell(colorLabel("~${changed}", ChangeColors.modified).apply {
-                                    font = JBUI.Fonts.smallFont()
-                                })
+                                cell(compactDiffPanel(ins, del, changed))
                             }
                         }
                     }
@@ -113,27 +112,61 @@ class EditApprovalPanel(
         )
     }
 
-    private fun createInlineDiffComponent(payload: EditPayload?): JComponent? {
-        if (payload == null) return null
-        val path = try {
-            Paths.get(payload.filePath).normalize().toString()
-        } catch (_: Exception) {
-            payload.filePath
+    private fun createInlineDiffComponent(): JComponent? {
+        val payload = request.payload ?: return null
+        
+        val (path, current, proposed) = when (payload) {
+            is EditPayload -> {
+                val normalizedPath = try {
+                    Paths.get(payload.filePath).normalize().toString()
+                } catch (_: Exception) {
+                    payload.filePath
+                }
+                val currentContent = getFileContentWithFallback(normalizedPath)
+                val proposedContent = if (!payload.proposedContent.isNullOrBlank()) {
+                    payload.proposedContent
+                } else {
+                    val rawSnippet = if (payload.newString.isNotBlank()) payload.newString else payload.oldString
+                    if (UpdateSnippetUtil.containsMarkers(rawSnippet)) {
+                        return JBLabel("Preview not available").apply {
+                            foreground = JBUI.CurrentTheme.Label.disabledForeground()
+                            border = JBUI.Borders.empty(8)
+                        }.let { label -> BorderLayoutPanel().apply { addToCenter(label) } }
+                    }
+                    applyStringReplacement(
+                        currentContent,
+                        payload.oldString,
+                        payload.newString,
+                        payload.replaceAll
+                    )
+                }
+                Triple(normalizedPath, currentContent, proposedContent)
+            }
+            is ProxyAIEditPayload -> {
+                val normalizedPath = try {
+                    Paths.get(payload.filePath).normalize().toString()
+                } catch (_: Exception) {
+                    payload.filePath
+                }
+                Triple(normalizedPath, payload.originalContent, payload.updatedContent)
+            }
+            else -> return null
         }
-        val vfs = LocalFileSystem.getInstance().refreshAndFindFileByPath(path)
-
-        val current = getFileContentWithFallback(path)
-        val proposed = applyStringReplacement(
-            current,
-            payload.oldString,
-            payload.newString,
-            payload.replaceAll
-        )
 
         val (insRaw, delRaw, changed) = lineDiffStats(current, proposed)
         diffCounts = Triple(insRaw, delRaw, changed)
 
+        val vfs = LocalFileSystem.getInstance().refreshAndFindFileByPath(path)
         val language = vfs?.extension ?: "text"
+        return buildDiffEditor(current, proposed, language, path)
+    }
+
+    private fun buildDiffEditor(
+        current: String,
+        proposed: String,
+        language: String,
+        path: String
+    ): JComponent {
         val segment = ReplaceWaiting(current, proposed, language, path)
         val editor = ChatEditorFactory.createEditor(project, segment)
         ResponseEditorPanel.RESPONSE_EDITOR_DIFF_VIEWER_KEY.get(editor)?.let { viewer ->
@@ -180,4 +213,20 @@ class EditApprovalPanel(
 
     private fun colorLabel(text: String, color: JBColor): JBLabel =
         JBLabel(text).apply { foreground = color }
+
+    private fun compactDiffPanel(inserted: Int, deleted: Int, changed: Int): JComponent {
+        val texts = diffBadgeText(inserted, deleted, changed)
+        return JPanel(FlowLayout(FlowLayout.LEFT, 0, 0)).apply {
+            isOpaque = false
+            add(colorLabel(texts.inserted, ChangeColors.inserted).apply {
+                font = JBUI.Fonts.smallFont()
+            })
+            add(colorLabel(texts.deleted, ChangeColors.deleted).apply {
+                font = JBUI.Fonts.smallFont()
+            })
+            add(colorLabel(texts.changed, ChangeColors.modified).apply {
+                font = JBUI.Fonts.smallFont()
+            })
+        }
+    }
 }

src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/ui/approval/ToolApprovalPayload.kt

@@ -16,6 +16,13 @@ data class EditPayload(
     val filePath: String,
     val oldString: String,
     val newString: String,
-    val replaceAll: Boolean
+    val replaceAll: Boolean,
+    val proposedContent: String? = null
 ) : ToolApprovalPayload
 
+data class ProxyAIEditPayload(
+    val filePath: String,
+    val updateSnippet: String,
+    val originalContent: String,
+    val updatedContent: String
+) : ToolApprovalPayload