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 a558d6c0..130bb96f 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 @@ -153,11 +153,11 @@ public class ChatMessageResponseBody extends JPanel { } public void addToolStatusPanel(JComponent component) { - currentlyProcessedTextPane = null; - currentlyProcessedEditorPanel = null; - currentlyProcessedMermaidPanel = null; - streamOutputParser.clear(); + finishCurrentStreamingSection(); + streamOutputParser.startNewVisualSection(); contentPanel.add(component); + contentPanel.revalidate(); + contentPanel.repaint(); } public void updateMessage(String partialMessage) { @@ -379,6 +379,12 @@ public class ChatMessageResponseBody extends JPanel { } } + private void finishCurrentStreamingSection() { + currentlyProcessedTextPane = null; + currentlyProcessedEditorPanel = null; + currentlyProcessedMermaidPanel = null; + } + private void processCode(Segment item) { if (isMermaidCode(item)) { processMermaid(item.getContent()); 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 ba10638c..96529002 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 @@ -26,6 +26,18 @@ class SseMessageParser : MessageParser { buffer.clear() } + fun startNewVisualSection() { + buffer.clear() + parserState = when (val state = parserState) { + is ParserState.Outside -> ParserState.Outside + is ParserState.CodeHeaderWaiting -> state.copy(content = "") + is ParserState.InCode -> state.copy(content = "") + is ParserState.InSearch -> state.copy(searchContent = "") + is ParserState.InReplace -> state.copy(searchContent = "", replaceContent = "") + is ParserState.InThinking -> ParserState.InThinking() + } + } + override fun parse(input: String): List { val segments = mutableListOf() var position = 0 diff --git a/src/test/kotlin/ee/carlrobert/codegpt/toolwindow/chat/ui/ChatMessageResponseBodyTest.kt b/src/test/kotlin/ee/carlrobert/codegpt/toolwindow/chat/ui/ChatMessageResponseBodyTest.kt new file mode 100644 index 00000000..e1a60007 --- /dev/null +++ b/src/test/kotlin/ee/carlrobert/codegpt/toolwindow/chat/ui/ChatMessageResponseBodyTest.kt @@ -0,0 +1,86 @@ +package ee.carlrobert.codegpt.toolwindow.chat.ui + +import com.intellij.openapi.application.ApplicationManager +import com.intellij.openapi.editor.EditorFactory +import com.intellij.openapi.util.Disposer +import ee.carlrobert.codegpt.toolwindow.chat.editor.ResponseEditorPanel +import org.assertj.core.api.Assertions.assertThat +import testsupport.IntegrationTest +import java.awt.Component +import java.awt.Container +import javax.swing.JPanel +import javax.swing.JTextPane + +class ChatMessageResponseBodyTest : IntegrationTest() { + + fun testToolPanelStartsNewTextPaneForSubsequentStreamedText() { + ApplicationManager.getApplication().invokeAndWait { + val disposable = Disposer.newDisposable() + try { + val responseBody = ChatMessageResponseBody(project, false, disposable) + + responseBody.updateMessage("Hello") + responseBody.addToolStatusPanel(JPanel()) + responseBody.updateMessage(" world") + + val textPanes = findComponents(responseBody, JTextPane::class.java) + assertThat(textPanes).hasSize(2) + assertThat(textPanes.first().text).contains("Hello") + assertThat(textPanes.last().text).contains("world") + assertThat(textPanes.last().text).doesNotContain("Hello") + } finally { + Disposer.dispose(disposable) + } + } + } + + fun testToolPanelStartsNewCodePanelForSubsequentStreamedCode() { + ApplicationManager.getApplication().invokeAndWait { + val disposable = Disposer.newDisposable() + try { + val responseBody = ChatMessageResponseBody(project, false, disposable) + + responseBody.updateMessage("```kotlin\nfun main() {\n") + responseBody.addToolStatusPanel(JPanel()) + responseBody.updateMessage(" println(\"x\")\n}\n```") + + val editorPanels = findComponents(responseBody, ResponseEditorPanel::class.java) + assertThat(editorPanels).hasSize(2) + assertThat(editorPanels.first().getEditor()?.document?.text) + .contains("fun main() {") + .doesNotContain("println(\"x\")") + assertThat(editorPanels.last().getEditor()?.document?.text) + .contains("println(\"x\")") + .contains("}") + .doesNotContain("fun main() {") + + val textPanes = findComponents(responseBody, JTextPane::class.java) + assertThat(textPanes).isEmpty() + + editorPanels.forEach { panel -> + panel.getEditor()?.let { editor -> + EditorFactory.getInstance().releaseEditor(editor) + } + } + } finally { + Disposer.dispose(disposable) + } + } + } + + private fun findComponents(root: Component, type: Class): List { + val matches = mutableListOf() + + fun visit(component: Component) { + if (type.isInstance(component)) { + matches.add(type.cast(component)) + } + if (component is Container) { + component.components.forEach(::visit) + } + } + + visit(root) + return matches + } +}