diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/ResponseEditorPanel.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/ResponseEditorPanel.kt
index deaf48b4..96e09d8e 100644
--- a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/ResponseEditorPanel.kt
+++ b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/ResponseEditorPanel.kt
@@ -147,7 +147,7 @@ class ResponseEditorPanel(
                 }
 
                 if (!response.isNullOrBlank()) {
-                    stateManager.transitionToDiffState(originalCode, response, params.destination)
+                    stateManager.transitionToDiffState(originalCode, response, params.destination, params.source)
                 }
             } catch (e: Exception) {
                 logger.error("Failed to apply changes", e)
@@ -224,7 +224,8 @@ class ResponseEditorPanel(
             val containsText = currentText.contains(segment.search.trim())
 
             val newState = if (containsText) {
-                stateManager.createFromSegment(segment)
+                val finalSegment = createReplaceWaitingSegment(searchContent, replaceContent, virtualFile)
+                stateManager.createFromSegment(finalSegment, readOnly = false, eventSource = null, originalSuggestion = replaceContent)
             } else {
                 stateManager.transitionToFailedDiffState(
                     segment.search,
@@ -339,4 +340,4 @@ class ResponseEditorPanel(
             content
         }
     }
-}
\ No newline at end of file
+}
diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/state/EditorStateManager.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/state/EditorStateManager.kt
index 2404791e..a98c58f6 100644
--- a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/state/EditorStateManager.kt
+++ b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/state/EditorStateManager.kt
@@ -44,12 +44,21 @@ class EditorStateManager(private val project: Project) {
     }
 
     fun transitionToDiffState(originalCode: String, updatedCode: String, virtualFile: VirtualFile) {
+        transitionToDiffState(originalCode, updatedCode, virtualFile, null)
+    }
+
+    fun transitionToDiffState(
+        originalCode: String,
+        updatedCode: String,
+        virtualFile: VirtualFile,
+        originalSuggestion: String?
+    ) {
         val oldState = currentState ?: return
         val oldEditor = oldState.editor
         val language = oldState.segment.language.ifBlank { virtualFile.extension ?: "Text" }
         val segment = ReplaceWaiting(originalCode, updatedCode, language, virtualFile.path)
         val editor = EditorFactory.createEditor(project, segment)
-        val state = createDiffState(editor, segment)
+        val state = createDiffState(editor, segment, null, originalSuggestion)
 
         runInEdt {
             val headerComponent = state.createHeaderComponent(false)
diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/state/StandardDiffEditorState.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/state/StandardDiffEditorState.kt
index de1eef04..8999c1a5 100644
--- a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/state/StandardDiffEditorState.kt
+++ b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/state/StandardDiffEditorState.kt
@@ -28,6 +28,11 @@ class StandardDiffEditorState(
     eventSource: EventSource? = null,
     private val originalSuggestion: String? = null
 ) : DiffEditorState(editor, segment, project, diffViewer, virtualFile, eventSource) {
+    private var latestSuggestion: String? = when (segment) {
+        is SearchReplace -> segment.replace.takeIf { it.isNotBlank() }
+        is ReplaceWaiting -> segment.replace.takeIf { it.isNotBlank() }
+        else -> null
+    } ?: originalSuggestion?.takeIf { it.isNotBlank() }
 
     override fun applyAllChanges() {
         val before = diffViewer?.getDocument(Side.LEFT)?.text ?: return
@@ -50,6 +55,10 @@ class StandardDiffEditorState(
                 return
             }
 
+            if (replace.isNotBlank()) {
+                latestSuggestion = replace
+            }
+
             diffEditorManager.updateDiffContent(search, replace)
             (editor.permanentHeaderComponent as? DiffHeaderPanel)
                 ?.updateDiffStats(diffViewer?.diffChanges ?: emptyList())
@@ -67,11 +76,15 @@ class StandardDiffEditorState(
     override fun handleClose() {
         runInEdt {
             val responsePanel = editor.component.parent as? ResponseEditorPanel ?: return@runInEdt
-            val contentToKeep = originalSuggestion ?: when (segment) {
-                is SearchReplace -> segment.replace
-                is ReplaceWaiting -> segment.replace
-                else -> diffViewer?.getDocument(Side.RIGHT)?.text ?: ""
-            }
+            val contentToKeep =
+                originalSuggestion?.takeIf { it.isNotBlank() }
+                    ?: latestSuggestion?.takeIf { it.isNotBlank() }
+                    ?: when (segment) {
+                        is SearchReplace -> segment.replace
+                        is ReplaceWaiting -> segment.replace
+                        else -> null
+                    }?.takeIf { it.isNotBlank() }
+                    ?: (diffViewer?.getDocument(Side.RIGHT)?.text ?: "")
             responsePanel.replaceEditorWithSegment(
                 Code(contentToKeep, segment.language, segment.filePath)
             )