From 4e54a2f90b5f5c8bba9b03083b4617d2e2272de1 Mon Sep 17 00:00:00 2001 From: Carl-Robert Linnupuu Date: Thu, 29 May 2025 12:15:57 +0100 Subject: [PATCH] refactor: sse message parsing --- ...{RetryListener.kt => AutoApplyListener.kt} | 0 .../chat/editor/state/EditorStateManager.kt | 16 +- .../codegpt/toolwindow/chat/parser/Segment.kt | 2 +- .../chat/parser/SseMessageParser.kt | 527 ++++++++++-------- .../chat/parser/SseMessageParserTest.kt | 228 ++++++++ 5 files changed, 541 insertions(+), 232 deletions(-) rename src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/{RetryListener.kt => AutoApplyListener.kt} (100%) create mode 100644 src/test/kotlin/ee/carlrobert/codegpt/toolwindow/chat/parser/SseMessageParserTest.kt diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/RetryListener.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/AutoApplyListener.kt similarity index 100% rename from src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/RetryListener.kt rename to src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/AutoApplyListener.kt 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 3068632d..5bebcb8a 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 @@ -7,18 +7,15 @@ import com.intellij.openapi.editor.ex.EditorEx import com.intellij.openapi.project.Project import com.intellij.openapi.vfs.LocalFileSystem import com.intellij.openapi.vfs.VirtualFile -import com.intellij.openapi.vfs.readText -import ee.carlrobert.codegpt.CodeGPTKeys import ee.carlrobert.codegpt.completions.AutoApplyParameters import ee.carlrobert.codegpt.completions.CompletionRequestService import ee.carlrobert.codegpt.toolwindow.chat.editor.ResponseEditorPanel import ee.carlrobert.codegpt.toolwindow.chat.editor.ResponseEditorPanel.Companion.RESPONSE_EDITOR_STATE_KEY -import ee.carlrobert.codegpt.toolwindow.chat.editor.RetryListener +import ee.carlrobert.codegpt.toolwindow.chat.editor.AutoApplyListener 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 ee.carlrobert.codegpt.toolwindow.chat.parser.SseMessageParser @Service(Service.Level.PROJECT) class EditorStateManager(private val project: Project) { @@ -42,22 +39,13 @@ class EditorStateManager(private val project: Project) { return state } - fun handleRetryForFailedSearch(replaceContent: String) { - val editor = currentState?.editor ?: return - val virtualFile = - CodeGPTKeys.TOOLWINDOW_EDITOR_FILE_DETAILS.get(editor)?.virtualFile ?: return - - getCodeEditsAsync(replaceContent, virtualFile, editor) - } - fun getCodeEditsAsync( content: String, virtualFile: VirtualFile, editor: EditorEx, ) { val params = AutoApplyParameters(content, virtualFile) - val messageParser = SseMessageParser() - val listener = RetryListener(project, messageParser, this) { newEditor -> + val listener = AutoApplyListener(project, this) { newEditor -> val responseEditorPanel = editor.component.parent as? ResponseEditorPanel ?: throw IllegalStateException("Expected parent to be ResponseEditorPanel") responseEditorPanel.replaceEditor(editor, newEditor) diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/parser/Segment.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/parser/Segment.kt index 62385545..ed0ddd85 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/parser/Segment.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/parser/Segment.kt @@ -39,4 +39,4 @@ data class SearchReplace( val replace: String, override val language: String, override val filePath: String? -) : Segment(search, language, filePath) \ No newline at end of file +) : Segment(replace, language, filePath) \ No newline at end of file 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 5e351526..8568c60c 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 @@ -1,239 +1,332 @@ package ee.carlrobert.codegpt.toolwindow.chat.parser -enum class State { OUTSIDE, CODE_HEADER_WAITING, IN_CODE, IN_SEARCH, IN_REPLACE, IN_THINKING } - class SseMessageParser : MessageParser { - var state = State.OUTSIDE - private val buffer = StringBuilder() - private val parsedSegments = mutableListOf() + private companion object { + const val CODE_FENCE = "```" + const val THINK_START = "" + const val THINK_END = "" + const val SEARCH_MARKER = "<<<<<<< SEARCH" + const val SEPARATOR_MARKER = "=======" + const val REPLACE_MARKER = ">>>>>>> REPLACE" + const val NEWLINE = "\n" + const val HEADER_DELIMITER = ":" + const val HEADER_PARTS_LIMIT = 2 + } - private var currentCodeHeader: CodeHeader? = null - private val codeBuilder = StringBuilder() - private val headerBuilder = StringBuilder() - private val searchBuilder = StringBuilder() - private val replaceBuilder = StringBuilder() - private val thinkingBuilder = StringBuilder() + private var parserState: ParserState = ParserState.Outside + private val buffer = StringBuilder() fun clear() { - state = State.OUTSIDE + parserState = ParserState.Outside buffer.clear() - parsedSegments.clear() - currentCodeHeader = null - codeBuilder.clear() - headerBuilder.clear() - searchBuilder.clear() - replaceBuilder.clear() - thinkingBuilder.clear() } - /** - * Parse incoming partial text and return any completed segments. - * Leftover text remains in buffer until more input arrives. - */ override fun parse(input: String): List { buffer.append(input) - val output = mutableListOf() + val segments = mutableListOf() - loop@ while (true) { - when (state) { - State.OUTSIDE -> { - val fenceIdx = buffer.indexOf("```") - val thinkStartIdx = buffer.indexOf("") + while (processNextSegment(segments)) { + // Continue processing until no more complete segments can be extracted + } - when { - fenceIdx != -1 && (thinkStartIdx == -1 || fenceIdx < thinkStartIdx) -> { - if (fenceIdx > 0) { - output += Text(buffer.substring(0, fenceIdx)) - } - buffer.delete(0, fenceIdx + 3) - state = State.CODE_HEADER_WAITING - headerBuilder.clear() - continue@loop - } + segments.addAll(getPendingSegments()) - thinkStartIdx != -1 -> { - if (thinkStartIdx > 0) { - output += Text(buffer.substring(0, thinkStartIdx)) - } - buffer.delete(0, thinkStartIdx + "".length) - state = State.IN_THINKING - thinkingBuilder.clear() - continue@loop - } + return segments + } - else -> break@loop - } + private fun processNextSegment(segments: MutableList): Boolean { + return when (val state = parserState) { + is ParserState.Outside -> processOutsideState(segments) + is ParserState.CodeHeaderWaiting -> processCodeHeaderState(segments, state) + is ParserState.InCode -> processInCodeState(segments, state) + is ParserState.InSearch -> processInSearchState(segments, state) + is ParserState.InReplace -> processInReplaceState(segments, state) + is ParserState.InThinking -> processInThinkingState(segments, state) + } + } + + private fun processOutsideState(segments: MutableList): Boolean { + val fenceIdx = buffer.indexOf(CODE_FENCE) + val thinkStartIdx = buffer.indexOf(THINK_START) + + return when { + shouldProcessCodeFence(fenceIdx, thinkStartIdx) -> { + extractTextBeforeIndex(fenceIdx)?.let { segments.add(it) } + consumeFromBuffer(fenceIdx + CODE_FENCE.length) + parserState = ParserState.CodeHeaderWaiting() + true + } + + thinkStartIdx != -1 -> { + extractTextBeforeIndex(thinkStartIdx)?.let { segments.add(it) } + consumeFromBuffer(thinkStartIdx + THINK_START.length) + parserState = ParserState.InThinking() + true + } + + else -> false + } + } + + private fun processCodeHeaderState( + segments: MutableList, + state: ParserState.CodeHeaderWaiting + ): Boolean { + val nlIdx = buffer.indexOf(NEWLINE) + if (nlIdx < 0) return false + + val headerLine = buffer.substring(0, nlIdx).trim() + consumeFromBuffer(nlIdx + 1) + + val updatedHeader = state.content + headerLine + val header = parseCodeHeader(updatedHeader) + + return if (header != null) { + segments.add(header) + parserState = ParserState.InCode(header) + true + } else { + segments.add(CodeHeaderWaiting(updatedHeader)) + parserState = ParserState.CodeHeaderWaiting(updatedHeader) + false + } + } + + private fun processInCodeState( + segments: MutableList, + state: ParserState.InCode + ): Boolean { + val nlIdx = buffer.indexOf(NEWLINE) + if (nlIdx < 0) return false + + val line = buffer.substring(0, nlIdx) + consumeFromBuffer(nlIdx + 1) + + return when { + line.trim() == CODE_FENCE -> { + if (state.content.isNotEmpty()) { + segments.add(Code(state.content, state.header.language, state.header.filePath)) } + segments.add(CodeEnd("")) + parserState = ParserState.Outside + true + } - State.CODE_HEADER_WAITING -> { - val nlIdx = buffer.indexOf("\n") - if (nlIdx < 0) break@loop - val headerLine = buffer.substring(0, nlIdx).trim() - buffer.delete(0, nlIdx + 1) - headerBuilder.append(headerLine) + line.trimStart().startsWith(SEARCH_MARKER) -> { + segments.add(SearchWaiting("", state.header.language, state.header.filePath)) + parserState = ParserState.InSearch(state.header, "") + true + } - val headerText = headerBuilder.toString() - val parts = headerText.split(":", limit = 2) - - val language = parts.getOrNull(0) ?: "" - val fileName = parts.getOrNull(1) - - if (parts.size > 0) { - currentCodeHeader = CodeHeader(language, fileName) - output += currentCodeHeader!! - state = State.IN_CODE - codeBuilder.clear() - headerBuilder.clear() - } else { - output += CodeHeaderWaiting(headerText) - } - } - - State.IN_CODE -> { - val idx = buffer.indexOf("\n") - if (idx < 0) break@loop - val line = buffer.substring(0, idx) - buffer.delete(0, idx + 1) - when { - line.trim() == "```" -> { - if (codeBuilder.isNotEmpty()) { - output += Code( - codeBuilder.toString(), - currentCodeHeader!!.language, - currentCodeHeader!!.filePath - ) - } - output += CodeEnd("") - state = State.OUTSIDE - } - - line.trimStart().startsWith("<<<<<<< SEARCH") -> { - state = State.IN_SEARCH - searchBuilder.clear() - output += SearchWaiting( - "", - currentCodeHeader!!.language, - currentCodeHeader!!.filePath - ) - } - - else -> codeBuilder.appendLine(line) - } - } - - State.IN_SEARCH -> { - val idx = buffer.indexOf("\n") - if (idx < 0) break@loop - val line = buffer.substring(0, idx) - buffer.delete(0, idx + 1) - if (line.trim() == "=======") { - state = State.IN_REPLACE - replaceBuilder.clear() - output += ReplaceWaiting( - searchBuilder.toString(), - "", - currentCodeHeader!!.language, - currentCodeHeader!!.filePath - ) - } else { - searchBuilder.appendLine(line) - output += SearchWaiting( - searchBuilder.toString(), - currentCodeHeader!!.language, - currentCodeHeader!!.filePath - ) - } - } - - State.IN_REPLACE -> { - val idx = buffer.indexOf("\n") - if (idx < 0) break@loop - val line = buffer.substring(0, idx) - buffer.delete(0, idx + 1) - if (line.trim().startsWith(">>>>>>> REPLACE")) { - output += SearchReplace( - search = searchBuilder.toString(), - replace = replaceBuilder.toString(), - language = currentCodeHeader!!.language, - filePath = currentCodeHeader!!.filePath - ) - state = State.IN_CODE - } else { - replaceBuilder.appendLine(line) - output += ReplaceWaiting( - searchBuilder.toString(), - replaceBuilder.toString(), - currentCodeHeader!!.language, - currentCodeHeader!!.filePath - ) - } - } - - State.IN_THINKING -> { - val thinkEndIdx = buffer.indexOf("") - if (thinkEndIdx < 0) { - if (buffer.isNotEmpty()) { - thinkingBuilder.append(buffer) - output += Thinking(thinkingBuilder.toString()) - buffer.clear() - } - break@loop - } - - thinkingBuilder.append(buffer.substring(0, thinkEndIdx)) - output += Thinking(thinkingBuilder.toString()) - buffer.delete(0, thinkEndIdx + "".length) - state = State.OUTSIDE - thinkingBuilder.clear() - continue@loop - } + else -> { + val newContent = + if (state.content.isEmpty()) line else state.content + NEWLINE + line + parserState = ParserState.InCode(state.header, newContent) + false } } - - when (state) { - State.OUTSIDE -> - if (buffer.isNotBlank()) - output += Text(buffer.toString()) - - State.CODE_HEADER_WAITING -> - if (headerBuilder.isNotBlank()) - output += CodeHeaderWaiting(headerBuilder.toString()) - - State.IN_CODE -> - if (codeBuilder.isNotBlank()) - output += Code( - codeBuilder.toString(), - currentCodeHeader!!.language, - currentCodeHeader!!.filePath - ) - - State.IN_SEARCH -> - if (searchBuilder.isNotBlank()) - output += SearchWaiting( - searchBuilder.toString(), - currentCodeHeader!!.language, - currentCodeHeader!!.filePath - ) - - State.IN_REPLACE -> - if (replaceBuilder.isNotBlank()) - output += ReplaceWaiting( - searchBuilder.toString(), - replaceBuilder.toString(), - currentCodeHeader!!.language, - currentCodeHeader!!.filePath - ) - - State.IN_THINKING -> - if (thinkingBuilder.isNotBlank() || buffer.isNotBlank()) { - thinkingBuilder.append(buffer) - buffer.clear() - output += Thinking(thinkingBuilder.toString()) - } - } - - parsedSegments.addAll(output) - return output } -} \ No newline at end of file + + private fun processInSearchState( + segments: MutableList, + state: ParserState.InSearch + ): Boolean { + val nlIdx = buffer.indexOf(NEWLINE) + if (nlIdx < 0) return false + + val line = buffer.substring(0, nlIdx) + consumeFromBuffer(nlIdx + 1) + + return if (line.trim() == SEPARATOR_MARKER) { + segments.add( + ReplaceWaiting( + state.searchContent, + "", + state.header.language, + state.header.filePath + ) + ) + parserState = ParserState.InReplace(state.header, state.searchContent, "") + true + } else { + val newSearch = + if (state.searchContent.isEmpty()) line else state.searchContent + NEWLINE + line + segments.add(SearchWaiting(newSearch, state.header.language, state.header.filePath)) + parserState = ParserState.InSearch(state.header, newSearch) + false + } + } + + private fun processInReplaceState( + segments: MutableList, + state: ParserState.InReplace + ): Boolean { + val nlIdx = buffer.indexOf(NEWLINE) + if (nlIdx < 0) return false + + val line = buffer.substring(0, nlIdx) + consumeFromBuffer(nlIdx + 1) + + return when { + line.trim().startsWith(REPLACE_MARKER) -> { + segments.add( + SearchReplace( + search = state.searchContent, + replace = state.replaceContent, + language = state.header.language, + filePath = state.header.filePath + ) + ) + parserState = ParserState.InCode(state.header) + true + } + line.trim() == CODE_FENCE -> { + // Invalid search/replace block - missing REPLACE marker + // Mark done + segments.add(CodeEnd("")) + parserState = ParserState.Outside + true + } + else -> { + val newReplace = + if (state.replaceContent.isEmpty()) line else state.replaceContent + NEWLINE + line + segments.add( + ReplaceWaiting( + state.searchContent, + newReplace, + state.header.language, + state.header.filePath + ) + ) + parserState = ParserState.InReplace(state.header, state.searchContent, newReplace) + false + } + } + } + + private fun processInThinkingState( + segments: MutableList, + state: ParserState.InThinking + ): Boolean { + val thinkEndIdx = buffer.indexOf(THINK_END) + + return if (thinkEndIdx >= 0) { + val thinkingContent = state.content + buffer.substring(0, thinkEndIdx) + segments.add(Thinking(thinkingContent)) + consumeFromBuffer(thinkEndIdx + THINK_END.length) + parserState = ParserState.Outside + true + } else { + if (buffer.isNotEmpty()) { + val newContent = state.content + buffer.toString() + segments.add(Thinking(newContent)) + buffer.clear() + parserState = ParserState.InThinking(newContent) + } + false + } + } + + private fun getPendingSegments(): List { + return when (val state = parserState) { + is ParserState.Outside -> { + if (buffer.isNotBlank()) listOf(Text(buffer.toString())) + else emptyList() + } + + is ParserState.CodeHeaderWaiting -> { + if (state.content.isNotBlank()) listOf(CodeHeaderWaiting(state.content)) + else emptyList() + } + + is ParserState.InCode -> { + if (state.content.isNotBlank()) { + listOf(Code(state.content, state.header.language, state.header.filePath)) + } else emptyList() + } + + is ParserState.InSearch -> { + if (state.searchContent.isNotBlank()) { + listOf( + SearchWaiting( + state.searchContent, + state.header.language, + state.header.filePath + ) + ) + } else emptyList() + } + + is ParserState.InReplace -> { + if (state.replaceContent.isNotBlank()) { + listOf( + ReplaceWaiting( + state.searchContent, + state.replaceContent, + state.header.language, + state.header.filePath + ) + ) + } else emptyList() + } + + is ParserState.InThinking -> { + val fullContent = state.content + buffer.toString() + buffer.clear() + if (fullContent.isNotBlank()) listOf(Thinking(fullContent)) + else emptyList() + } + } + } + + private fun shouldProcessCodeFence(fenceIdx: Int, thinkStartIdx: Int): Boolean { + return fenceIdx != -1 && (thinkStartIdx == -1 || fenceIdx < thinkStartIdx) + } + + private fun extractTextBeforeIndex(index: Int): Text? { + return if (index > 0) Text(buffer.substring(0, index)) else null + } + + private fun consumeFromBuffer(length: Int) { + buffer.delete(0, length) + } + + private fun parseCodeHeader(headerText: String): CodeHeader? { + val parts = headerText.split(HEADER_DELIMITER, limit = HEADER_PARTS_LIMIT) + return if (parts.isNotEmpty()) { + CodeHeader( + language = parts.getOrNull(0) ?: "", + filePath = parts.getOrNull(1) + ) + } else null + } + + private sealed class ParserState { + object Outside : ParserState() + + data class CodeHeaderWaiting( + val content: String = "" + ) : ParserState() + + data class InCode( + val header: CodeHeader, + val content: String = "" + ) : ParserState() + + data class InSearch( + val header: CodeHeader, + val searchContent: String = "" + ) : ParserState() + + data class InReplace( + val header: CodeHeader, + val searchContent: String, + val replaceContent: String = "" + ) : ParserState() + + data class InThinking( + val content: String = "" + ) : ParserState() + } +} diff --git a/src/test/kotlin/ee/carlrobert/codegpt/toolwindow/chat/parser/SseMessageParserTest.kt b/src/test/kotlin/ee/carlrobert/codegpt/toolwindow/chat/parser/SseMessageParserTest.kt new file mode 100644 index 00000000..b3743bc1 --- /dev/null +++ b/src/test/kotlin/ee/carlrobert/codegpt/toolwindow/chat/parser/SseMessageParserTest.kt @@ -0,0 +1,228 @@ +package ee.carlrobert.codegpt.toolwindow.chat.parser + +import org.assertj.core.api.Assertions.assertThat +import org.junit.Test +import kotlin.random.Random + +class SseMessageParserStreamTest { + + /** + * Simulates streaming by breaking input into random chunks + */ + private fun simulateStreaming( + parser: SseMessageParser, + input: String, + minChunkSize: Int = 1, + maxChunkSize: Int = 10, + seed: Long? = null + ): List { + val random = seed?.let { Random(it) } ?: Random + val allSegments = mutableListOf() + var position = 0 + + while (position < input.length) { + val chunkSize = random.nextInt(minChunkSize, maxChunkSize + 1) + .coerceAtMost(input.length - position) + val chunk = input.substring(position, position + chunkSize) + + val segments = parser.parse(chunk) + allSegments.addAll(segments) + + position += chunkSize + + // Simulate network delay + Thread.sleep(random.nextLong(5, 20)) + } + + return allSegments + } + + @Test + fun shouldHandleStreamedCodeBlock() { + val parser = SseMessageParser() + val input = """ + Here is some code: + ```kotlin:MyFile.kt + fun main() { + println("Hello, World!") + } + ``` + Done! + """.trimIndent() + + val segments = simulateStreaming(parser, input, seed = 42) + + val segmentTypes = segments.map { it::class.simpleName } + assertThat(segmentTypes).contains("Text", "CodeHeader", "Code", "CodeEnd") + val codeSegments = segments.filterIsInstance() + assertThat(codeSegments).isNotEmpty + assertThat(codeSegments.last().content).contains("println(\"Hello, World!\")") + } + + @Test + fun shouldHandleStreamedSearchReplace() { + val parser = SseMessageParser() + val input = """ + ```kotlin:MyFile.kt + <<<<<<< SEARCH + fun oldFunction() { + return "old" + } + ======= + fun newFunction() { + return "new" + } + >>>>>>> REPLACE + ``` + """.trimIndent() + + val segments = simulateStreaming(parser, input, minChunkSize = 3, maxChunkSize = 15, seed = 123) + + val searchReplaceSegments = segments.filterIsInstance() + assertThat(searchReplaceSegments).hasSize(1) + val sr = searchReplaceSegments[0] + assertThat(sr.search).contains("oldFunction") + assertThat(sr.replace).contains("newFunction") + assertThat(sr.language).isEqualTo("kotlin") + assertThat(sr.filePath).isEqualTo("MyFile.kt") + } + + @Test + fun shouldHandleSearchReplaceWithInvalidEnding() { + val parser = SseMessageParser() + val input = """ + Here's some text. + + ```kotlin:MyFile.kt + <<<<<<< SEARCH + fun oldFunction() { + return "old" + } + ======= + fun newFunction() { + return "new" + } + ``` + + Here's some other text. + """.trimIndent() + + val segments = simulateStreaming(parser, input, minChunkSize = 3, maxChunkSize = 15, seed = 123) + + val searchReplaceSegments = segments.filterIsInstance() + assertThat(searchReplaceSegments).isEmpty() + val codeEndSegments = segments.filterIsInstance() + assertThat(codeEndSegments).hasSize(1) + val textSegments = segments.filterIsInstance() + assertThat(textSegments.any { it.content.contains("Here's some other text") }).isTrue + val replaceWaitingSegments = segments.filterIsInstance() + assertThat(replaceWaitingSegments.size).isLessThan(10) + } + + @Test + fun shouldHandleStreamedThinkingBlock() { + val parser = SseMessageParser() + val input = """ + Let me analyze this... + + First, I need to understand the requirements. + Then, I'll design a solution. + Finally, I'll implement it. + + Here's my solution: + """.trimIndent() + + val segments = simulateStreaming(parser, input, minChunkSize = 2, maxChunkSize = 8, seed = 456) + + val thinkingSegments = segments.filterIsInstance() + assertThat(thinkingSegments).isNotEmpty + val finalThinking = thinkingSegments.last() + assertThat(finalThinking.content).contains("understand the requirements") + assertThat(finalThinking.content).contains("design a solution") + assertThat(finalThinking.content).contains("implement it") + } + + @Test + fun shouldHandleMultipleCodeBlocksStreamed() { + val parser = SseMessageParser() + val input = """ + First: + ```java + System.out.println("1"); + ``` + Second: + ```python + print("2") + ``` + Third: + ```javascript + console.log("3"); + ``` + """.trimIndent() + + val segments = simulateStreaming(parser, input, minChunkSize = 5, maxChunkSize = 20, seed = 789) + + val codeHeaders = segments.filterIsInstance() + assertThat(codeHeaders).hasSize(3) + assertThat(codeHeaders.map { it.language }).containsExactly("java", "python", "javascript") + val codeSegments = segments.filterIsInstance() + assertThat(codeSegments.any { it.content.contains("System.out.println") }).isTrue + assertThat(codeSegments.any { it.content.contains("print(\"2\")") }).isTrue + assertThat(codeSegments.any { it.content.contains("console.log") }).isTrue + } + + @Test + fun shouldHandleMixedContentStreamed() { + val parser = SseMessageParser() + val input = """ + Starting analysis... + + Processing request... + + Here's the code: + ```kotlin:Solution.kt + <<<<<<< SEARCH + val old = 1 + ======= + val new = 2 + >>>>>>> REPLACE + ``` + And a simple block: + ```python + print("done") + ``` + """.trimIndent() + + val segments = simulateStreaming(parser, input, minChunkSize = 3, maxChunkSize = 12, seed = 999) + + val segmentTypeSet = segments.map { it::class.simpleName }.toSet() + assertThat(segmentTypeSet).contains( + "Text", "Thinking", "CodeHeader", "SearchWaiting", + "ReplaceWaiting", "SearchReplace", "Code", "CodeEnd" + ) + } + + @Test + fun shouldHandleRandomChunkingConsistently() { + val input = """ + ```kotlin:Test.kt + class Test { + fun method() { + println("Hello") + } + } + ``` + """.trimIndent() + + repeat(5) { iteration -> + val parser = SseMessageParser() + val segments = simulateStreaming(parser, input, seed = iteration.toLong()) + + val codeSegments = segments.filterIsInstance() + assertThat(codeSegments).isNotEmpty + val finalCode = codeSegments.last().content + assertThat(finalCode).contains("class Test") + assertThat(finalCode).contains("println(\"Hello\")") + } + } +} \ No newline at end of file