diff --git a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ui/ChatMessageResponseBody.java b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ui/ChatMessageResponseBody.java
index e9ec4e82..812d024f 100644
--- a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ui/ChatMessageResponseBody.java
+++ b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ui/ChatMessageResponseBody.java
@@ -137,7 +137,7 @@ public class ChatMessageResponseBody extends JPanel {
   public ChatMessageResponseBody withResponse(@NotNull String response) {
     try {
       for (var item : new CompleteMessageParser().parse(response)) {
-        processResponse(item, false, false);
+        processResponse(item, false);
         currentlyProcessedTextPane = null;
         currentlyProcessedEditorPanel = null;
       }
@@ -158,7 +158,7 @@ public class ChatMessageResponseBody extends JPanel {
 
     var parsedResponse = streamOutputParser.parse(partialMessage);
     for (Segment item : parsedResponse) {
-      processResponse(item, true, true);
+      processResponse(item, true);
     }
   }
 
@@ -287,7 +287,7 @@ public class ChatMessageResponseBody extends JPanel {
         .orElse(null);
   }
 
-  private void processResponse(Segment item, boolean caretVisible, boolean partialResponse) {
+  private void processResponse(Segment item, boolean caretVisible) {
     if (item instanceof Thinking) {
       processThinkingOutput(item.getContent());
       return;
diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/diff/DiffEditorManager.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/diff/DiffEditorManager.kt
index 0ca5ae1f..7f645eec 100644
--- a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/diff/DiffEditorManager.kt
+++ b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/diff/DiffEditorManager.kt
@@ -5,7 +5,6 @@ import com.intellij.diff.tools.fragmented.UnifiedDiffViewer
 import com.intellij.diff.util.DiffUtil
 import com.intellij.diff.util.Side
 import com.intellij.openapi.application.runInEdt
-import com.intellij.openapi.application.runWriteAction
 import com.intellij.openapi.diff.DiffBundle
 import com.intellij.openapi.editor.Document
 import com.intellij.openapi.editor.LogicalPosition
@@ -14,8 +13,7 @@ import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.text.StringUtil
 import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.openapi.vfs.readText
-import com.intellij.util.concurrency.annotations.RequiresEdt
-import ee.carlrobert.codegpt.toolwindow.chat.editor.diff.DiffManagerUtil.replaceContent
+import com.intellij.util.application
 
 class DiffEditorManager(
     private val project: Project,
@@ -27,14 +25,20 @@ class DiffEditorManager(
         val currentText = virtualFile?.readText() ?: return
         val document = diffViewer.getDocument(Side.RIGHT)
 
-        runInEdt {
-            document.replaceContent(
-                project,
-                currentText.replaceLast(searchContent.trim(), replaceContent.trim())
-            )
-
-            diffViewer.rediff()
-            scrollToLastChange(diffViewer)
+        application.executeOnPooledThread {
+            runInEdt {
+                ensureDocumentWritable(project, document)
+                if (DiffUtil.executeWriteCommand(document, project, "Updating document") {
+                        document.setText(
+                            StringUtil.convertLineSeparators(
+                                currentText.replaceLast(searchContent.trim(), replaceContent.trim())
+                            )
+                        )
+                    }) {
+                    diffViewer.rediff()
+                    scrollToLastChange(diffViewer)
+                }
+            }
         }
     }
 
@@ -46,7 +50,7 @@ class DiffEditorManager(
 
     fun applyAllChanges(): List<UnifiedDiffChange> {
         val document = diffViewer.getDocument(Side.LEFT)
-        DiffManagerUtil.ensureDocumentWritable(project, document)
+        ensureDocumentWritable(project, document)
 
         val allChanges = mutableListOf<UnifiedDiffChange>()
 
@@ -81,19 +85,8 @@ class DiffEditorManager(
             ScrollType.CENTER
         )
     }
-}
 
-object DiffManagerUtil {
-
-    @RequiresEdt
-    fun Document.replaceContent(project: Project, replaceContent: String) {
-        ensureDocumentWritable(project, this)
-        DiffUtil.executeWriteCommand(this, project, "Updating document") {
-            setText(StringUtil.convertLineSeparators(replaceContent))
-        }
-    }
-
-    fun ensureDocumentWritable(project: Project, document: Document) {
+    private fun ensureDocumentWritable(project: Project, document: Document) {
         if (!document.isWritable) {
             DiffUtil.makeWritable(project, document)
             document.setReadOnly(false)
diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/parser/SseMessageParser.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/parser/SseMessageParser.kt
index 62432525..5b2d65b9 100644
--- a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/parser/SseMessageParser.kt
+++ b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/parser/SseMessageParser.kt
@@ -120,6 +120,10 @@ class SseMessageParser : MessageParser {
             }
 
             line.trimStart().startsWith(SEARCH_MARKER) -> {
+                // Emit accumulated code content before transitioning
+                if (state.content.isNotEmpty()) {
+                    segments.add(Code(state.content, state.header.language, state.header.filePath))
+                }
                 segments.add(SearchWaiting("", state.header.language, state.header.filePath))
                 parserState = ParserState.InSearch(state.header, "")
                 true
@@ -129,7 +133,7 @@ class SseMessageParser : MessageParser {
                 val newContent =
                     if (state.content.isEmpty()) line else state.content + NEWLINE + line
                 parserState = ParserState.InCode(state.header, newContent)
-                false
+                true
             }
         }
     }