From 6ae4187f6b14a9f3d43de42a79f490bba3ca5e71 Mon Sep 17 00:00:00 2001 From: Carl-Robert Linnupuu Date: Thu, 29 Jan 2026 00:41:36 +0000 Subject: [PATCH] fix: rollback changes tracking --- .../codegpt/agent/rollback/RollbackService.kt | 44 ++++++++++++++----- .../codegpt/agent/tools/EditArgsSnapshot.kt | 38 ++++++++++++++++ .../ee/carlrobert/codegpt/util/EditorUtil.kt | 16 ++++++- .../carlrobert/codegpt/util/file/FileUtil.kt | 30 ++++++++++--- .../rollback/RollbackServiceTrackingTest.kt | 4 +- 5 files changed, 111 insertions(+), 21 deletions(-) create mode 100644 src/main/kotlin/ee/carlrobert/codegpt/agent/tools/EditArgsSnapshot.kt diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/rollback/RollbackService.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/rollback/RollbackService.kt index fc9d26aa..a7569e64 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/rollback/RollbackService.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/rollback/RollbackService.kt @@ -3,16 +3,18 @@ package ee.carlrobert.codegpt.agent.rollback import com.intellij.history.Label import com.intellij.history.LocalHistory import com.intellij.openapi.application.ModalityState +import com.intellij.openapi.application.runReadAction import com.intellij.openapi.application.runInEdt import com.intellij.openapi.application.runWriteAction import com.intellij.openapi.components.Service +import com.intellij.openapi.fileEditor.FileDocumentManager import com.intellij.openapi.fileTypes.FileTypeManager import com.intellij.openapi.project.Project import com.intellij.openapi.vfs.LocalFileSystem import com.intellij.openapi.vfs.VfsUtilCore import com.intellij.openapi.vfs.VirtualFileManager import ee.carlrobert.codegpt.settings.ProxyAISettingsService -import ee.carlrobert.codegpt.agent.tools.EditTool +import ee.carlrobert.codegpt.agent.tools.EditArgsSnapshot import ee.carlrobert.codegpt.agent.tools.WriteTool import kotlinx.coroutines.Dispatchers import kotlinx.coroutines.withContext @@ -51,9 +53,15 @@ class RollbackService(private val project: Project) { * Track an EditTool operation directly from the agent. * This captures the original file content before the edit is applied. */ - fun trackEdit(sessionId: String, filePath: String, args: EditTool.Args, originalContent: String) { + fun trackEdit( + sessionId: String, + filePath: String, + args: EditArgsSnapshot, + originalContent: String + ) { val tracker = activeRuns[sessionId] ?: return - tracker.recordExplicitEdit(filePath, args, originalContent) + val normalizedPath = filePath.replace("\\", "/") + tracker.recordExplicitEdit(normalizedPath, args, originalContent) } /** @@ -62,7 +70,8 @@ class RollbackService(private val project: Project) { */ fun trackWrite(sessionId: String, filePath: String, args: WriteTool.Args) { val tracker = activeRuns[sessionId] ?: return - tracker.recordExplicitWrite(filePath, args) + val normalizedPath = filePath.replace("\\", "/") + tracker.recordExplicitWrite(normalizedPath, args) } fun startSession(sessionId: String) { @@ -94,16 +103,23 @@ class RollbackService(private val project: Project) { } fun getDiffData(sessionId: String, path: String): RollbackDiffData? { - if (!isTrackable(path)) return null val snapshot = snapshots[sessionId] ?: return null val change = snapshot.changes[path] ?: return null + if (change.kind != ChangeKind.DELETED && !isTrackable(path)) return null val beforeText = when (change.kind) { ChangeKind.ADDED -> "" - else -> decodeLabelContent( - snapshot.labelRef, - change.originalPath ?: path, - change.originalContent - ) + else -> { + val original = change.originalContent + if (original != null && original.isNotEmpty()) { + decodeContent(original) + } else { + decodeLabelContent( + snapshot.labelRef, + change.originalPath ?: path, + change.originalContent + ) + } + } } val afterText = when (change.kind) { ChangeKind.DELETED -> "" @@ -212,6 +228,10 @@ class RollbackService(private val project: Project) { private fun readCurrentText(path: String): String { val vf = LocalFileSystem.getInstance().refreshAndFindFileByPath(path) ?: return "" + val docText = runReadAction { + FileDocumentManager.getInstance().getDocument(vf)?.text + } + if (docText != null) return docText return runCatching { VfsUtilCore.loadText(vf) }.getOrDefault("") } @@ -348,7 +368,7 @@ class RollbackService(private val project: Project) { val labelRef: Label, val changes: MutableMap = ConcurrentHashMap() ) { - fun recordExplicitEdit(filePath: String, args: EditTool.Args, originalContent: String) { + fun recordExplicitEdit(filePath: String, args: EditArgsSnapshot, originalContent: String) { val existing = changes[filePath] if (existing?.kind == ChangeKind.ADDED) return if (existing?.kind == ChangeKind.MOVED) return @@ -466,4 +486,4 @@ enum class ChangeKind { sealed class RollbackResult { data class Success(val message: String) : RollbackResult() data class Failure(val message: String) : RollbackResult() -} \ No newline at end of file +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/EditArgsSnapshot.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/EditArgsSnapshot.kt new file mode 100644 index 00000000..6c60d8af --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/EditArgsSnapshot.kt @@ -0,0 +1,38 @@ +package ee.carlrobert.codegpt.agent.tools + +data class EditArgsSnapshot( + val filePath: String, + val oldString: String, + val newString: String, + val replaceAll: Boolean, + val shortDescription: String +) + +fun EditTool.Args.toSnapshot(): EditArgsSnapshot { + return EditArgsSnapshot( + filePath = filePath, + oldString = oldString, + newString = newString, + replaceAll = replaceAll, + shortDescription = shortDescription + ) +} + +fun ProxyAIEditTool.Args.toSnapshot(): EditArgsSnapshot { + return EditArgsSnapshot( + filePath = filePath, + oldString = "", + newString = updateSnippet, + replaceAll = false, + shortDescription = shortDescription + ) +} + +fun snapshotFromEditArgs(args: Any?): EditArgsSnapshot? { + return when (args) { + is EditTool.Args -> args.toSnapshot() + is ProxyAIEditTool.Args -> args.toSnapshot() + is EditArgsSnapshot -> args + else -> null + } +} \ No newline at end of file diff --git a/src/main/kotlin/ee/carlrobert/codegpt/util/EditorUtil.kt b/src/main/kotlin/ee/carlrobert/codegpt/util/EditorUtil.kt index 0cfb4086..82738366 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/util/EditorUtil.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/util/EditorUtil.kt @@ -210,4 +210,18 @@ object EditorUtil { } } } -} \ No newline at end of file + + @JvmStatic + fun writeDocumentContent(project: Project, virtualFile: VirtualFile, content: String): Boolean { + val document = runReadAction { + FileDocumentManager.getInstance().getDocument(virtualFile) + } ?: return false + + WriteCommandAction.runWriteCommandAction(project) { + document.setText(content) + FileDocumentManager.getInstance().saveDocument(document) + } + + return true + } +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/util/file/FileUtil.kt b/src/main/kotlin/ee/carlrobert/codegpt/util/file/FileUtil.kt index 6cade673..c4cc5468 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/util/file/FileUtil.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/util/file/FileUtil.kt @@ -9,6 +9,7 @@ import com.intellij.openapi.util.io.FileUtil.createDirectory import com.intellij.openapi.vfs.JarFileSystem import com.intellij.openapi.vfs.LocalFileSystem import com.intellij.openapi.vfs.VfsUtilCore +import com.intellij.openapi.vfs.VirtualFileManager import com.intellij.openapi.vfs.VirtualFile import ee.carlrobert.codegpt.settings.service.llama.LlamaSettings.getLlamaModelsPath import java.io.File @@ -42,7 +43,6 @@ object FileUtil { } } - @JvmStatic fun readContent(virtualFile: VirtualFile): String { try { return VfsUtilCore.loadText(virtualFile) @@ -113,7 +113,6 @@ object FileUtil { }.takeIf { it } ?: throw RuntimeException("Failed to create directory: $directoryPath") } - @JvmStatic fun getFileExtension(filename: String?): String { val pattern = Pattern.compile("[^.]+$") val matcher = filename?.let { pattern.matcher(it) } @@ -124,7 +123,6 @@ object FileUtil { return "" } - @JvmStatic fun findLanguageExtensionMapping(language: String? = ""): Map.Entry { val defaultValue = mapOf("Text" to ".txt").entries.first() val mapper = ObjectMapper() @@ -169,7 +167,6 @@ object FileUtil { } } - @JvmStatic fun getImageMediaType(fileName: String?): String { return when (val fileExtension = getFileExtension(fileName)) { "png" -> "image/png" @@ -178,7 +175,6 @@ object FileUtil { } } - @JvmStatic fun getResourceContent(filePath: String?): String { try { Objects.requireNonNull(filePath?.let { FileUtil::class.java.getResourceAsStream(it) }) @@ -216,7 +212,6 @@ object FileUtil { return value.toString() } - @JvmStatic fun findFirstExtension( languageFileExtensionMappings: List, language: String? = "" @@ -254,4 +249,27 @@ object FileUtil { null } } + + fun findVirtualFile(normalizedPath: String): VirtualFile? { + return VirtualFileManager.getInstance().refreshAndFindFileByUrl("file://$normalizedPath") + ?: LocalFileSystem.getInstance().findFileByIoFile(File(normalizedPath)) + } + + fun validateFileForEdit(filePath: String): Result { + val normalizedPath = filePath.replace("\\", "/") + val file = File(normalizedPath) + + return when { + !file.exists() -> Result.failure( + IllegalArgumentException("File not found: $filePath (File does not exist on filesystem)") + ) + !file.isFile -> Result.failure( + IllegalArgumentException("Path is not a file: $filePath") + ) + !file.canWrite() -> Result.failure( + IllegalArgumentException("File is not writable: $filePath") + ) + else -> Result.success(file) + } + } } \ No newline at end of file diff --git a/src/test/kotlin/ee/carlrobert/codegpt/agent/rollback/RollbackServiceTrackingTest.kt b/src/test/kotlin/ee/carlrobert/codegpt/agent/rollback/RollbackServiceTrackingTest.kt index abfb5cdf..7af63259 100644 --- a/src/test/kotlin/ee/carlrobert/codegpt/agent/rollback/RollbackServiceTrackingTest.kt +++ b/src/test/kotlin/ee/carlrobert/codegpt/agent/rollback/RollbackServiceTrackingTest.kt @@ -1,7 +1,7 @@ package ee.carlrobert.codegpt.agent.rollback import com.intellij.openapi.vfs.LocalFileSystem -import ee.carlrobert.codegpt.agent.tools.EditTool +import ee.carlrobert.codegpt.agent.tools.EditArgsSnapshot import ee.carlrobert.codegpt.agent.tools.WriteTool import org.assertj.core.api.Assertions.assertThat import testsupport.IntegrationTest @@ -28,7 +28,7 @@ class RollbackServiceTrackingTest : IntegrationTest() { rollbackService.trackEdit( sessionId = sessionId, filePath = filePath, - args = EditTool.Args(filePath, "before", "after", "test", false), + args = EditArgsSnapshot(filePath, "before", "after", false, "test"), originalContent = "before" ) val snapshot = rollbackService.finishSession(sessionId)