diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/AutoApplyListener.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/AutoApplyListener.kt
index eea5b4a7..9d3e2f1a 100644
--- a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/AutoApplyListener.kt
+++ b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/AutoApplyListener.kt
@@ -17,23 +17,29 @@ class AutoApplyListener(
     private val project: Project,
     private val stateManager: EditorStateManager,
     private val virtualFile: VirtualFile,
+    private val originalSuggestion: String,
     private val onEditorReplaced: (EditorEx, EditorEx) -> Unit
 ) : CompletionEventListener<String> {
 
     private val logger = logger<AutoApplyListener>()
     private var editorReplaced: Boolean = false
     private val messageParser = SseMessageParser()
+    private var eventSource: EventSource? = null
 
     override fun onOpen() {
         CompletionProgressNotifier.update(project, true)
     }
 
     override fun onMessage(message: String, eventSource: EventSource?) {
-        processMessageSegments(message, eventSource)
+        if (this.eventSource == null && eventSource != null) {
+            this.eventSource = eventSource
+        }
+        processMessageSegments(message)
     }
 
     override fun onError(error: ErrorDetails?, ex: Throwable?) {
-        logger.error("Something went wrong while retrying diff-based editing", ex)
+        logger.error("Something went wrong while applying the changes", ex)
+        ErrorHandler.handleError(error, ex)
         handleComplete()
     }
 
@@ -45,10 +51,7 @@ class AutoApplyListener(
         handleComplete()
     }
 
-    private fun processMessageSegments(
-        message: String,
-        eventSource: EventSource?
-    ) {
+    private fun processMessageSegments(message: String) {
         val segments = messageParser.parse(message)
         for (segment in segments) {
             when (segment) {
@@ -85,7 +88,7 @@ class AutoApplyListener(
         val containsText = currentText.contains(segment.search.trim())
 
         val newState = if (containsText) {
-            stateManager.createFromSegment(segment)
+            stateManager.createFromSegment(segment, false, eventSource, originalSuggestion)
         } else {
             stateManager.transitionToFailedDiffState(
                 segment.search,
@@ -101,5 +104,7 @@ class AutoApplyListener(
         val editor = stateManager.getCurrentState()?.editor ?: return
         (editor.permanentHeaderComponent as? DiffHeaderPanel)?.handleDone()
         CompletionProgressNotifier.update(project, false)
+        eventSource = null
     }
+
 }
\ No newline at end of file
diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/ErrorHandler.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/ErrorHandler.kt
new file mode 100644
index 00000000..7a78c9df
--- /dev/null
+++ b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/ErrorHandler.kt
@@ -0,0 +1,21 @@
+package ee.carlrobert.codegpt.toolwindow.chat.editor
+
+import com.intellij.notification.NotificationType
+import ee.carlrobert.codegpt.ui.OverlayUtil
+import ee.carlrobert.llm.client.openai.completion.ErrorDetails
+
+object ErrorHandler {
+    
+    fun handleError(error: ErrorDetails?, ex: Throwable?) {
+        val errorMessage = formatErrorMessage(error, ex)
+        OverlayUtil.showNotification(errorMessage, NotificationType.ERROR)
+    }
+    
+    fun formatErrorMessage(error: ErrorDetails?, ex: Throwable?): String {
+        return when {
+            error?.code == "insufficient_quota" -> "You exceeded your current quota, please check your plan and billing details."
+            ex?.message != null -> "Error: ${ex.message}"
+            else -> "An unknown error occurred while applying changes."
+        }
+    }
+}
\ No newline at end of file
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 e469484c..b94aac4c 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
@@ -2,7 +2,9 @@ package ee.carlrobert.codegpt.toolwindow.chat.editor
 
 import com.intellij.diff.tools.fragmented.UnifiedDiffViewer
 import com.intellij.openapi.Disposable
+import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.application.runInEdt
+import com.intellij.openapi.diagnostic.thisLogger
 import com.intellij.openapi.editor.EditorKind
 import com.intellij.openapi.editor.LogicalPosition
 import com.intellij.openapi.editor.ScrollType
@@ -10,11 +12,9 @@ import com.intellij.openapi.editor.event.BulkAwareDocumentListener
 import com.intellij.openapi.editor.event.DocumentEvent
 import com.intellij.openapi.editor.ex.EditorEx
 import com.intellij.openapi.project.Project
-import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.util.Disposer
 import com.intellij.openapi.util.Key
 import com.intellij.openapi.vfs.VirtualFile
-import com.intellij.openapi.diagnostic.thisLogger
 import com.intellij.openapi.vfs.readText
 import com.intellij.util.ui.JBUI
 import com.intellij.util.ui.components.BorderLayoutPanel
@@ -25,6 +25,7 @@ import ee.carlrobert.codegpt.toolwindow.chat.editor.diff.DiffSyncManager
 import ee.carlrobert.codegpt.toolwindow.chat.editor.factory.ComponentFactory
 import ee.carlrobert.codegpt.toolwindow.chat.editor.factory.ComponentFactory.EXPANDED_KEY
 import ee.carlrobert.codegpt.toolwindow.chat.editor.factory.ComponentFactory.MIN_LINES_FOR_EXPAND
+import ee.carlrobert.codegpt.toolwindow.chat.editor.header.DefaultHeaderPanel
 import ee.carlrobert.codegpt.toolwindow.chat.editor.header.DiffHeaderPanel
 import ee.carlrobert.codegpt.toolwindow.chat.editor.state.EditorState
 import ee.carlrobert.codegpt.toolwindow.chat.editor.state.EditorStateManager
@@ -98,17 +99,39 @@ class ResponseEditorPanel(
         }
     }
 
-    fun applyCodeAsync(content: String, virtualFile: VirtualFile, editor: EditorEx) {
-        CompletionRequestService.getInstance().autoApplyAsync(
+    fun replaceEditorWithSegment(segment: Segment) {
+        val oldEditor = stateManager.getCurrentState()?.editor ?: return
+        val newState = stateManager.createFromSegment(segment)
+        replaceEditor(oldEditor, newState.editor)
+    }
+
+    fun removeCurrentEditor() {
+        runInEdt {
+            removeAll()
+            stateManager.clearCurrentState()
+            revalidate()
+            repaint()
+        }
+    }
+
+    fun applyCodeAsync(content: String, virtualFile: VirtualFile, editor: EditorEx, headerPanel: DefaultHeaderPanel? = null) {
+        val eventSource = CompletionRequestService.getInstance().autoApplyAsync(
             AutoApplyParameters(content, virtualFile),
-            AutoApplyListener(project, stateManager, virtualFile) { oldEditor, newEditor ->
+            AutoApplyListener(project, stateManager, virtualFile, content) { oldEditor, newEditor ->
                 val responseEditorPanel = editor.component.parent as? ResponseEditorPanel
                     ?: throw IllegalStateException("Expected parent to be ResponseEditorPanel")
                 responseEditorPanel.replaceEditor(oldEditor, newEditor)
             })
+        
+        val panel = headerPanel ?: (editor.permanentHeaderComponent as? DefaultHeaderPanel)
+        panel?.setLoading(eventSource)
     }
-    
-    internal fun createReplaceWaitingSegment(searchContent: String, replaceContent: String, virtualFile: VirtualFile): ReplaceWaiting {
+
+    internal fun createReplaceWaitingSegment(
+        searchContent: String,
+        replaceContent: String,
+        virtualFile: VirtualFile
+    ): ReplaceWaiting {
         return ReplaceWaiting(
             search = searchContent,
             replace = replaceContent,
@@ -116,26 +139,30 @@ class ResponseEditorPanel(
             filePath = virtualFile.path
         )
     }
-    
-    fun createDiffEditorForDirectApply(searchContent: String, replaceContent: String, virtualFile: VirtualFile) {
+
+    fun createDiffEditorForDirectApply(
+        searchContent: String,
+        replaceContent: String,
+        virtualFile: VirtualFile
+    ) {
         try {
             val segment = createReplaceWaitingSegment(searchContent, "", virtualFile)
-            
+
             val oldEditor = stateManager.getCurrentState()?.editor
             if (oldEditor == null) {
                 logger.warn("No current editor state found for direct apply")
                 return
             }
-            
+
             val currentText = try {
                 virtualFile.readText()
             } catch (e: Exception) {
                 logger.error("Failed to read file content for direct apply", e)
                 return
             }
-            
+
             val containsText = currentText.contains(segment.search.trim())
-            
+
             val newState = if (containsText) {
                 stateManager.createFromSegment(segment)
             } else {
@@ -148,15 +175,16 @@ class ResponseEditorPanel(
                     return
                 }
             }
-            
+
             replaceEditor(oldEditor, newState.editor)
-            
-            val finalSegment = createReplaceWaitingSegment(searchContent, replaceContent, virtualFile)
+
+            val finalSegment =
+                createReplaceWaitingSegment(searchContent, replaceContent, virtualFile)
             newState.updateContent(finalSegment)
-            
+
             val currentEditor = newState.editor
             val headerPanel = currentEditor.permanentHeaderComponent as? DiffHeaderPanel
-            
+
             ApplicationManager.getApplication().invokeLater {
                 if (!project.isDisposed) {
                     headerPanel?.handleDone()
diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/SearchReplaceHandler.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/SearchReplaceHandler.kt
index 9b7d997b..91134de5 100644
--- a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/SearchReplaceHandler.kt
+++ b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/SearchReplaceHandler.kt
@@ -5,6 +5,7 @@ import com.intellij.openapi.editor.EditorKind
 import com.intellij.openapi.editor.ex.EditorEx
 import com.intellij.openapi.vfs.readText
 import ee.carlrobert.codegpt.CodeGPTKeys
+import ee.carlrobert.codegpt.codecompletions.CompletionProgressNotifier
 import ee.carlrobert.codegpt.toolwindow.chat.editor.ResponseEditorPanel.Companion.RESPONSE_EDITOR_DIFF_VIEWER_VALUE_PAIR_KEY
 import ee.carlrobert.codegpt.toolwindow.chat.editor.header.DiffHeaderPanel
 import ee.carlrobert.codegpt.toolwindow.chat.editor.state.EditorStateManager
@@ -25,11 +26,11 @@ class SearchReplaceHandler(
     private var searchFailed = false
 
     fun handleSearchReplace(item: SearchReplace) {
+        handleReplace(item, item.filePath, item.search, item.replace)
+
         val editor = stateManager.getCurrentState()?.editor ?: return
         (editor.permanentHeaderComponent as? DiffHeaderPanel)?.handleDone()
-
         RESPONSE_EDITOR_DIFF_VIEWER_VALUE_PAIR_KEY.set(editor, Pair(item.search, item.replace))
-        handleReplace(item, item.filePath, item.search, item.replace)
     }
 
     fun handleReplace(item: ReplaceWaiting) {
diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/header/DefaultHeaderPanel.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/header/DefaultHeaderPanel.kt
index 999c3c42..b5062b40 100644
--- a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/header/DefaultHeaderPanel.kt
+++ b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/header/DefaultHeaderPanel.kt
@@ -4,29 +4,32 @@ import com.intellij.icons.AllIcons
 import com.intellij.openapi.actionSystem.*
 import com.intellij.openapi.actionSystem.toolbarLayout.ToolbarLayoutStrategy
 import com.intellij.openapi.application.runInEdt
+import com.intellij.openapi.diagnostic.thisLogger
 import com.intellij.openapi.editor.ex.EditorEx
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.ui.JBMenuItem
 import com.intellij.openapi.ui.JBPopupMenu
 import com.intellij.openapi.vfs.readText
-import com.intellij.ui.AnimatedIcon
-import com.intellij.ui.components.JBLabel
 import com.intellij.util.ui.JBUI
 import ee.carlrobert.codegpt.CodeGPTBundle
 import ee.carlrobert.codegpt.toolwindow.chat.editor.ResponseEditorPanel
 import ee.carlrobert.codegpt.toolwindow.chat.editor.actions.*
 import ee.carlrobert.codegpt.util.EditorUtil
 import ee.carlrobert.codegpt.util.StringUtil
+import okhttp3.sse.EventSource
 import javax.swing.JPanel
 
 class DefaultHeaderPanel(config: HeaderConfig) : HeaderPanel(config) {
 
-    private val loadingLabel: JBLabel by lazy {
-        JBLabel(
-            CodeGPTBundle.get("toolwindow.chat.editor.diff.thinking"),
-            AnimatedIcon.Default(),
-            JBLabel.LEFT
-        )
+    companion object {
+        private val logger = thisLogger()
+    }
+
+    private var currentEventSource: EventSource? = null
+    private val loadingPanel = LoadingPanel(
+        CodeGPTBundle.get("toolwindow.chat.editor.diff.thinking")
+    ) {
+        handleDone()
     }
 
     init {
@@ -35,21 +38,28 @@ class DefaultHeaderPanel(config: HeaderConfig) : HeaderPanel(config) {
 
     override fun initializeRightPanel(rightPanel: JPanel) {
         if (config.loading) {
-            rightPanel.add(loadingLabel)
+            rightPanel.add(loadingPanel)
         } else {
             rightPanel.add(createHeaderActions().component)
         }
     }
 
-    fun setLoading(label: String = "Loading...") {
+    fun setLoading(
+        eventSource: EventSource? = null,
+        label: String = CodeGPTBundle.get("toolwindow.chat.editor.diff.applying")
+    ) {
+        currentEventSource = eventSource
+        loadingPanel.setText(label)
+        loadingPanel.setEventSource(eventSource)
+
         runInEdt {
-            loadingLabel.text = label
-            setRightPanelComponent(loadingLabel)
+            setRightPanelComponent(loadingPanel)
         }
     }
 
     fun handleDone() {
         runInEdt {
+            currentEventSource = null
             setRightPanelComponent(createHeaderActions().component)
         }
     }
@@ -72,23 +82,27 @@ class DefaultHeaderPanel(config: HeaderConfig) : HeaderPanel(config) {
     }
 
     private fun handleApply(project: Project, editor: EditorEx) {
-        val file = virtualFile
-            ?: EditorUtil.getSelectedEditor(project)?.virtualFile
-            ?: throw IllegalStateException("Virtual file is null")
-
-        val directApplyThreshold = 0.85
-        val coefficient = StringUtil.getDiceCoefficient(editor.document.text, file.readText())
-        if (coefficient > directApplyThreshold) {
+        try {
+            val file = virtualFile
+                ?: EditorUtil.getSelectedEditor(project)?.virtualFile
+                ?: throw IllegalStateException("Could not find file")
             val responseEditorPanel = editor.component.parent as? ResponseEditorPanel
-                ?: throw IllegalStateException("Could not find corresponding ResponseEditorPanel")
-            responseEditorPanel.createDiffEditorForDirectApply(file.readText(), editor.document.text, file)
-            return
-        }
+                ?: throw IllegalStateException("Could not find editor panel")
 
-        val responseEditorPanel = editor.component.parent as? ResponseEditorPanel
-            ?: throw IllegalStateException("Could not find corresponding ResponseEditorPanel")
-        responseEditorPanel.applyCodeAsync(editor.document.text, file, editor)
-        setLoading("Editing...")
+            val directApplyThreshold = 0.85
+            val coefficient = StringUtil.getDiceCoefficient(editor.document.text, file.readText())
+            if (coefficient > directApplyThreshold) {
+                responseEditorPanel.createDiffEditorForDirectApply(
+                    file.readText(),
+                    editor.document.text,
+                    file
+                )
+                return
+            }
+            responseEditorPanel.applyCodeAsync(editor.document.text, file, editor, this)
+        } catch (e: Exception) {
+            logger.error(e.message, e)
+        }
     }
 
     private fun createToolbar(actionGroup: ActionGroup): ActionToolbar {
diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/header/DiffHeaderPanel.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/header/DiffHeaderPanel.kt
index 87a94aba..a6bfeaf5 100644
--- a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/header/DiffHeaderPanel.kt
+++ b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/header/DiffHeaderPanel.kt
@@ -1,46 +1,67 @@
 package ee.carlrobert.codegpt.toolwindow.chat.editor.header
 
 import com.intellij.diff.tools.fragmented.UnifiedDiffChange
+import com.intellij.icons.AllIcons
+import com.intellij.icons.AllIcons.Actions.Close
+import com.intellij.openapi.actionSystem.AnAction
+import com.intellij.openapi.actionSystem.AnActionEvent
 import com.intellij.openapi.application.runInEdt
-import com.intellij.openapi.editor.Document
-import com.intellij.ui.AnimatedIcon
 import com.intellij.ui.components.ActionLink
-import com.intellij.ui.components.JBLabel
 import ee.carlrobert.codegpt.CodeGPTBundle
 import ee.carlrobert.codegpt.toolwindow.chat.editor.ResponseEditorPanel
 import ee.carlrobert.codegpt.toolwindow.chat.editor.diff.DiffAcceptedPanel
+import ee.carlrobert.codegpt.ui.IconActionButton
+import okhttp3.sse.EventSource
 import java.awt.BorderLayout
-import javax.swing.BoxLayout
-import javax.swing.JPanel
+import javax.swing.*
 
 interface DiffHeaderActions {
     fun onAcceptAll()
     fun onOpenDiff()
+    fun onClose()
 }
 
 class DiffHeaderPanel(
     config: HeaderConfig,
     retry: Boolean,
-    private val actions: DiffHeaderActions
+    private val actions: DiffHeaderActions,
+    private var eventSource: EventSource? = null
 ) : HeaderPanel(config) {
 
-    private val loadingLabel: JBLabel = JBLabel(
-        if (retry) CodeGPTBundle.get("toolwindow.chat.editor.diff.retrying")
-        else CodeGPTBundle.get("toolwindow.chat.editor.diff.reading"),
-        AnimatedIcon.Default(),
-        JBLabel.LEFT
-    )
+    private val loadingPanel = LoadingPanel(
+        when {
+            retry -> CodeGPTBundle.get("toolwindow.chat.editor.diff.retrying")
+            eventSource != null -> CodeGPTBundle.get("toolwindow.chat.editor.diff.applying")
+            else -> CodeGPTBundle.get("toolwindow.chat.editor.diff.editing")
+        },
+        eventSource
+    ) {
+        handleDone()
+    }
 
     private val actionLinksPanel = JPanel().apply {
         layout = BoxLayout(this, BoxLayout.X_AXIS)
         isVisible = false
         add(ActionLink("View Diff") { actions.onOpenDiff() })
-        add(separator())
+        add(Box.createHorizontalStrut(6))
         add(ActionLink(CodeGPTBundle.get("shared.acceptAll")) { actions.onAcceptAll() })
+        add(separator())
+        add(IconActionButton(
+            object : AnAction("Close", "Close the diff view", Close) {
+                override fun actionPerformed(e: AnActionEvent) {
+                    actions.onClose()
+                }
+            },
+            "close-diff"
+        ))
     }
 
     init {
         setupUI()
+        runInEdt {
+            loadingPanel.isVisible = true
+            loadingPanel.setEventSource(eventSource)
+        }
     }
 
     override fun initializeRightPanel(rightPanel: JPanel) {
@@ -48,22 +69,24 @@ class DiffHeaderPanel(
 
         rightPanel.apply {
             add(actionLinksPanel)
-            add(loadingLabel)
+            add(loadingPanel)
         }
     }
 
     fun handleDone() {
+        eventSource = null
         runInEdt {
             actionLinksPanel.isVisible = true
-            loadingLabel.isVisible = false
+            loadingPanel.isVisible = false
             revalidate()
             repaint()
         }
     }
 
     fun handleChangesApplied(before: String, after: String, patches: List<UnifiedDiffChange>) {
+        eventSource = null
         actionLinksPanel.isVisible = false
-        loadingLabel.isVisible = false
+        loadingPanel.isVisible = false
 
         virtualFile?.let {
             val diffAcceptedPanel = DiffAcceptedPanel(config.project, it, before, after, patches)
@@ -81,8 +104,9 @@ class DiffHeaderPanel(
 
     fun editing() {
         runInEdt {
-            loadingLabel.text = CodeGPTBundle.get("toolwindow.chat.editor.diff.editing")
-            loadingLabel.isVisible = true
+            loadingPanel.setText(CodeGPTBundle.get("toolwindow.chat.editor.diff.editing"))
+            loadingPanel.isVisible = true
+            loadingPanel.showStopButton(eventSource != null)
         }
     }
 }
diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/header/HeaderPanel.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/header/HeaderPanel.kt
index 05dfb976..b1077f11 100644
--- a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/header/HeaderPanel.kt
+++ b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/header/HeaderPanel.kt
@@ -46,6 +46,7 @@ abstract class HeaderPanel(protected val config: HeaderConfig) : BorderLayoutPan
 
     private val statsComponent = SimpleColoredComponent().apply {
         font = JBUI.Fonts.smallFont()
+        preferredSize = Dimension(preferredSize.width, 16)
     }
     private val errorLabel = JBLabel(AllIcons.General.Error).apply {
         isVisible = config.error != null
@@ -53,7 +54,6 @@ abstract class HeaderPanel(protected val config: HeaderConfig) : BorderLayoutPan
     }
     private val rightPanel = JPanel().apply {
         layout = BoxLayout(this, BoxLayout.X_AXIS)
-        alignmentY = 0.5f
         isOpaque = false
     }
 
@@ -72,14 +72,8 @@ abstract class HeaderPanel(protected val config: HeaderConfig) : BorderLayoutPan
     protected fun setupUI() {
         setupPanelAppearance()
         addToLeft(createLeftPanel(virtualFile))
-        rightPanel.removeAll()
         initializeRightPanel(rightPanel)
-
-        val rightCenteringPanel = JPanel(BorderLayout()).apply {
-            isOpaque = false
-            add(rightPanel, BorderLayout.CENTER)
-        }
-        addToRight(rightCenteringPanel)
+        addToRight(rightPanel)
     }
 
     protected fun setRightPanelComponent(component: JComponent?) {
diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/header/LoadingPanel.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/header/LoadingPanel.kt
new file mode 100644
index 00000000..37ea0b69
--- /dev/null
+++ b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/header/LoadingPanel.kt
@@ -0,0 +1,60 @@
+package ee.carlrobert.codegpt.toolwindow.chat.editor.header
+
+import com.intellij.icons.AllIcons
+import com.intellij.openapi.actionSystem.AnAction
+import com.intellij.openapi.actionSystem.AnActionEvent
+import com.intellij.ui.*
+import com.intellij.ui.components.JBLabel
+import ee.carlrobert.codegpt.ui.IconActionButton
+import okhttp3.sse.EventSource
+import javax.swing.BoxLayout
+import javax.swing.JPanel
+
+class LoadingPanel(
+    initialText: String,
+    private var eventSource: EventSource? = null,
+    private val onCancel: (() -> Unit)? = null
+) : JPanel() {
+
+    private val loadingLabel = JBLabel(initialText, AnimatedIcon.Default(), JBLabel.LEFT)
+    
+    private val stopButton = IconActionButton(
+        object : AnAction("Stop", "Stop the current operation", AllIcons.Actions.Suspend) {
+            override fun actionPerformed(e: AnActionEvent) {
+                eventSource?.cancel()
+                onCancel?.invoke()
+            }
+        },
+        "stop-operation"
+    ).apply {
+        isVisible = eventSource != null
+    }
+    
+    init {
+        layout = BoxLayout(this, BoxLayout.X_AXIS)
+        add(loadingLabel)
+        add(SeparatorComponent(
+            ColorUtil.fromHex("#48494b"),
+            SeparatorOrientation.VERTICAL
+        ).apply {
+            setVGap(4)
+            setHGap(6)
+        })
+        add(stopButton)
+    }
+    
+    fun setText(text: String) {
+        loadingLabel.text = text
+    }
+    
+    fun setEventSource(source: EventSource?) {
+        eventSource = source
+        stopButton.isVisible = source != null
+        revalidate()
+        repaint()
+    }
+    
+    fun showStopButton(show: Boolean) {
+        stopButton.isVisible = show && eventSource != null
+    }
+}
\ No newline at end of file
diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/state/DiffEditorState.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/state/DiffEditorState.kt
index 83515996..88eac704 100644
--- a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/state/DiffEditorState.kt
+++ b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/state/DiffEditorState.kt
@@ -29,6 +29,7 @@ import ee.carlrobert.codegpt.toolwindow.chat.editor.header.DiffHeaderPanel
 import ee.carlrobert.codegpt.toolwindow.chat.editor.header.HeaderConfig
 import ee.carlrobert.codegpt.toolwindow.chat.parser.Segment
 import ee.carlrobert.codegpt.util.file.FileUtil
+import okhttp3.sse.EventSource
 import java.util.*
 import javax.swing.Icon
 import javax.swing.JButton
@@ -39,7 +40,8 @@ abstract class DiffEditorState(
     override val segment: Segment,
     override val project: Project,
     val diffViewer: UnifiedDiffViewer?,
-    val virtualFile: VirtualFile?
+    val virtualFile: VirtualFile?,
+    private val eventSource: EventSource? = null
 ) : EditorState {
 
     companion object {
@@ -91,6 +93,10 @@ abstract class DiffEditorState(
             override fun onOpenDiff() {
                 openDiff()
             }
+
+            override fun onClose() {
+                handleClose()
+            }
         }
 
         return DiffHeaderPanel(
@@ -102,12 +108,15 @@ abstract class DiffEditorState(
                 false
             ),
             readOnly,
-            actions
+            actions,
+            eventSource
         )
     }
 
     abstract fun applyAllChanges()
 
+    abstract fun handleClose()
+
     private fun openDiff() {
         if (virtualFile == null) {
             throw IllegalStateException("Virtual file is null")
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 76980d17..2cb5e904 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
@@ -12,21 +12,24 @@ import ee.carlrobert.codegpt.toolwindow.chat.editor.diff.DiffEditorManager
 import ee.carlrobert.codegpt.toolwindow.chat.editor.factory.EditorFactory
 import ee.carlrobert.codegpt.toolwindow.chat.parser.Code
 import ee.carlrobert.codegpt.toolwindow.chat.parser.Segment
+import okhttp3.sse.EventSource
 
 class EditorStateManager(private val project: Project) {
 
     private var currentState: EditorState? = null
     private var diffEditorManager: DiffEditorManager? = null
 
-    fun createFromSegment(segment: Segment, readOnly: Boolean = false): EditorState {
+    fun createFromSegment(segment: Segment, readOnly: Boolean = false, eventSource: EventSource? = null, originalSuggestion: String? = null): EditorState {
         val editor = EditorFactory.createEditor(project, segment)
         val state = if (editor.editorKind == EditorKind.DIFF) {
-            createDiffState(editor, segment)
+            createDiffState(editor, segment, eventSource, originalSuggestion)
         } else {
             RegularEditorState(editor, segment, project)
         }
+        
         runInEdt {
-            EditorFactory.configureEditor(editor, state.createHeaderComponent(readOnly))
+            val headerComponent = state.createHeaderComponent(readOnly)
+            EditorFactory.configureEditor(editor, headerComponent)
         }
 
         RESPONSE_EDITOR_STATE_KEY.set(editor, state)
@@ -55,7 +58,8 @@ class EditorStateManager(private val project: Project) {
             FailedDiffEditorState(newEditor, newSegment, project, searchContent, replaceContent)
 
         runInEdt {
-            EditorFactory.configureEditor(newEditor, newState.createHeaderComponent(false))
+            val headerComponent = newState.createHeaderComponent(false)
+            EditorFactory.configureEditor(newEditor, headerComponent)
         }
 
         this.currentState = newState
@@ -67,19 +71,27 @@ class EditorStateManager(private val project: Project) {
         return currentState
     }
 
-    private fun createDiffState(editor: EditorEx, segment: Segment): EditorState {
+    fun clearCurrentState() {
+        currentState = null
+    }
+
+    private fun createDiffState(editor: EditorEx, segment: Segment, eventSource: EventSource? = null, originalSuggestion: String? = null): EditorState {
         val virtualFile = getVirtualFile(segment.filePath)
         val diffViewer = ResponseEditorPanel.RESPONSE_EDITOR_DIFF_VIEWER_KEY.get(editor)
         val diffEditorManager = DiffEditorManager(project, diffViewer, virtualFile)
         this.diffEditorManager = diffEditorManager
-        return StandardDiffEditorState(
+        
+        val state = StandardDiffEditorState(
             editor,
             segment,
             project,
             diffViewer,
             virtualFile,
-            diffEditorManager
+            diffEditorManager,
+            eventSource,
+            originalSuggestion
         )
+        return state
     }
 
     private fun getVirtualFile(filePath: String?): VirtualFile? {
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 d0b8a9b7..de1eef04 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
@@ -8,13 +8,15 @@ import com.intellij.openapi.editor.EditorKind
 import com.intellij.openapi.editor.ex.EditorEx
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.vfs.VirtualFile
-import com.intellij.openapi.vfs.readText
 import com.intellij.util.application
+import ee.carlrobert.codegpt.toolwindow.chat.editor.ResponseEditorPanel
 import ee.carlrobert.codegpt.toolwindow.chat.editor.diff.DiffEditorManager
 import ee.carlrobert.codegpt.toolwindow.chat.editor.header.DiffHeaderPanel
+import ee.carlrobert.codegpt.toolwindow.chat.parser.Code
 import ee.carlrobert.codegpt.toolwindow.chat.parser.ReplaceWaiting
 import ee.carlrobert.codegpt.toolwindow.chat.parser.SearchReplace
 import ee.carlrobert.codegpt.toolwindow.chat.parser.Segment
+import okhttp3.sse.EventSource
 
 class StandardDiffEditorState(
     editor: EditorEx,
@@ -22,9 +24,11 @@ class StandardDiffEditorState(
     project: Project,
     diffViewer: UnifiedDiffViewer?,
     virtualFile: VirtualFile?,
-    private val diffEditorManager: DiffEditorManager
-) : DiffEditorState(editor, segment, project, diffViewer, virtualFile) {
-    
+    private val diffEditorManager: DiffEditorManager,
+    eventSource: EventSource? = null,
+    private val originalSuggestion: String? = null
+) : DiffEditorState(editor, segment, project, diffViewer, virtualFile, eventSource) {
+
     override fun applyAllChanges() {
         val before = diffViewer?.getDocument(Side.LEFT)?.text ?: return
         val after = diffViewer.getDocument(Side.RIGHT).text
@@ -45,7 +49,7 @@ class StandardDiffEditorState(
             } else {
                 return
             }
-            
+
             diffEditorManager.updateDiffContent(search, replace)
             (editor.permanentHeaderComponent as? DiffHeaderPanel)
                 ?.updateDiffStats(diffViewer?.diffChanges ?: emptyList())
@@ -59,4 +63,18 @@ 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 ?: ""
+            }
+            responsePanel.replaceEditorWithSegment(
+                Code(contentToKeep, segment.language, segment.filePath)
+            )
+        }
+    }
 }
diff --git a/src/main/resources/messages/codegpt.properties b/src/main/resources/messages/codegpt.properties
index c6298986..eefb7692 100644
--- a/src/main/resources/messages/codegpt.properties
+++ b/src/main/resources/messages/codegpt.properties
@@ -190,10 +190,10 @@ toolwindow.chat.editor.action.autoApply.description=Apply suggested changes auto
 toolwindow.chat.editor.action.autoApply.noActiveFile=Active file not found
 toolwindow.chat.editor.action.autoApply.fileTooLarge=Active file too large to process
 toolwindow.chat.editor.action.autoApply.reject=Reject All
-toolwindow.chat.editor.diff.reading=Reading...
-toolwindow.chat.editor.diff.thinking=Thinking...
-toolwindow.chat.editor.diff.editing=Editing...
-toolwindow.chat.editor.diff.retrying=Retrying...
+toolwindow.chat.editor.diff.applying=Applying
+toolwindow.chat.editor.diff.thinking=Thinking
+toolwindow.chat.editor.diff.editing=Editing
+toolwindow.chat.editor.diff.retrying=Retrying
 toolwindow.chat.editor.action.autoApply.error=Something went wrong while applying changes. {0}
 toolwindow.chat.editor.action.autoApply.taskTitle=Apply changes
 toolwindow.chat.editor.action.autoApply.loadingMessage=ProxyAI: Applying changes