diff --git a/ai-bridge/services/claude/message-sender.js b/ai-bridge/services/claude/message-sender.js index 70c78ba6..821449e4 100644 --- a/ai-bridge/services/claude/message-sender.js +++ b/ai-bridge/services/claude/message-sender.js @@ -37,7 +37,12 @@ import { import { createPreToolUseHook } from './permission-mode.js'; import { loadMcpServersConfigAsRecord } from './mcp-status/config-loader.js'; import { setActiveQueryResult } from './message-session-registry.js'; -import { normalizeStreamDelta, resolveSnapshotDelta, resetTurnBlockState } from './stream-delta-normalizer.js'; +import { + normalizeStreamDelta, + resolveSnapshotDelta, + resetTurnBlockState, + prepareAssistantSnapshotBlock, +} from './stream-delta-normalizer.js'; import { generateSessionTitle } from '../session-title-service.js'; import { getClaudeCliPathOverride } from '../../utils/claude-cli-path.js'; @@ -214,13 +219,39 @@ function processStreamMessage(msg, state, logPrefix) { } } } - if (event.type === 'content_block_start' && event.content_block?.type === 'thinking') { - console.log('[THINKING_START]'); + if (event.type === 'content_block_start' && event.content_block) { + const hasPreviousBlock = + (state.textBlockContentByIndex instanceof Map && state.textBlockContentByIndex.size > 0) + || (state.thinkingBlockContentByIndex instanceof Map + && state.thinkingBlockContentByIndex.size > 0); + if (hasPreviousBlock) { + resetTurnBlockState(state); + process.stdout.write('[BLOCK_RESET]\n'); + } } } return; } + // Assistant messages can represent one normalized content block rather than + // one cumulative response. Prepare the boundary before emitting [MESSAGE] so + // Java never merges a new thinking block into the previous one. + if (state.streamingEnabled && msg.type === 'assistant') { + const content = msg.message?.content; + const blocks = Array.isArray(content) + ? content + : typeof content === 'string' ? [{ type: 'text', text: content }] : []; + for (let i = 0; i < blocks.length; i += 1) { + const block = blocks[i]; + if (block.type === 'text' || block.type === 'thinking') { + const blockText = block.type === 'thinking' ? block.thinking || block.text || '' : block.text || ''; + if (prepareAssistantSnapshotBlock(state, block.type, i, blockText, msg)) { + process.stdout.write('[BLOCK_RESET]\n'); + } + } + } + } + // Determine whether to output the full [MESSAGE] tag let shouldOutput = true; if (state.streamingEnabled && msg.type === 'assistant') { diff --git a/ai-bridge/services/claude/stream-delta-normalizer.js b/ai-bridge/services/claude/stream-delta-normalizer.js index 189a08b8..774bf6ed 100644 --- a/ai-bridge/services/claude/stream-delta-normalizer.js +++ b/ai-bridge/services/claude/stream-delta-normalizer.js @@ -21,6 +21,19 @@ function modeKey(kind, blockIndex) { return `${kind}:${blockIndex}`; } +function getPayloadMap(turnState, fieldName) { + if (!(turnState[fieldName] instanceof Map)) { + turnState[fieldName] = new Map(); + } + return turnState[fieldName]; +} + +function isPrefixRelated(existing, incoming) { + return existing === incoming + || existing.startsWith(incoming) + || incoming.startsWith(existing); +} + function computeNovelDelta(previous, incoming, mode, origin) { if (!incoming) { return { novel: '', next: previous, mode }; @@ -133,9 +146,26 @@ export function normalizeStreamDelta(turnState, kind, index, incoming, origin = const modeMap = getModeMap(turnState); const mKey = modeKey(kind, blockIndex); const mode = modeMap.get(mKey); + const streamPayloads = getPayloadMap(turnState, 'lastStreamDeltaByKey'); + const snapshotPayloads = getPayloadMap(turnState, 'lastSnapshotPayloadByKey'); - const result = computeNovelDelta(previous, text, mode, origin); + // The SDK may yield one assistant message per content block after the live + // stream delta has already been delivered. Those messages carry the block + // fragment again rather than a cumulative snapshot. Treat an exact replay + // from either channel as already rendered; otherwise the Java layer sees the + // same thinking fragment twice. + const replayedSnapshot = origin === 'snapshot' + && (streamPayloads.get(mKey) === text || snapshotPayloads.get(mKey) === text); + const result = replayedSnapshot + ? { novel: '', next: previous, mode } + : computeNovelDelta(previous, text, mode, origin); blockMap.set(blockIndex, result.next); + if (origin === 'stream' && text) { + streamPayloads.set(mKey, text); + } + if (origin === 'snapshot' && text) { + snapshotPayloads.set(mKey, text); + } if (result.mode && result.mode !== mode) { modeMap.set(mKey, result.mode); } @@ -167,6 +197,68 @@ export function resolveSnapshotDelta(turnState, kind, index, snapshot) { return { delta, hadPrevious }; } +/** + * Detect a new assistant content block carried by a non-streaming snapshot. + * + * Claude Code normalizes one assistant message per content block, so several + * messages can share a response ID while carrying unrelated thinking text. + * The live stream path has explicit message/content-block boundaries; this + * helper supplies the same boundary for providers that only yield snapshots. + * + * @returns {boolean} whether the caller must emit a BLOCK_RESET marker first + */ +export function prepareAssistantSnapshotBlock(turnState, kind, index, incoming, message) { + if (!turnState.streamingEnabled || typeof incoming !== 'string' || incoming.length === 0) { + return false; + } + + const messageId = message?.message?.id; + if (typeof messageId !== 'string' || messageId.length === 0) { + return false; + } + + const messageUuid = typeof message?.uuid === 'string' + ? message.uuid + : typeof message?.message?.uuid === 'string' ? message.message.uuid : null; + const blockIndex = getBlockIndex(index); + const key = modeKey(kind, blockIndex); + const blockMap = getBlockMap( + turnState, + kind === 'thinking' ? 'thinkingBlockContentByIndex' : 'textBlockContentByIndex', + ); + const previous = blockMap.get(blockIndex) || ''; + const previousMessageId = turnState.lastAssistantMessageId; + const previousBlockUuid = turnState.lastAssistantBlockUuid; + const sameResponse = previousMessageId === messageId; + const streamPayloads = getPayloadMap(turnState, 'lastStreamDeltaByKey'); + const snapshotPayloads = getPayloadMap(turnState, 'lastSnapshotPayloadByKey'); + const isReplay = streamPayloads.get(key) === incoming || snapshotPayloads.get(key) === incoming; + const mode = turnState.blockStreamModeByKey instanceof Map + ? turnState.blockStreamModeByKey.get(key) + : undefined; + const uuidChanged = Boolean(messageUuid && previousBlockUuid && messageUuid !== previousBlockUuid); + + let startsNewBlock = false; + if (previous && previousMessageId && previousMessageId !== messageId) { + startsNewBlock = true; + } else if (previous && sameResponse && !isReplay) { + // A UUID change identifies the next normalized content-block message even + // when two thinking summaries happen to share a prefix. A divergent + // unkeyed snapshot is otherwise a new block unless cumulative snapshot mode + // explicitly tells us it is a correction of the current block. + const related = isPrefixRelated(previous, incoming); + const isSnapshotCorrection = mode === 'snapshot' && !uuidChanged; + startsNewBlock = uuidChanged || (!related && !isSnapshotCorrection); + } + + if (startsNewBlock) { + resetTurnBlockState(turnState); + } + turnState.lastAssistantMessageId = messageId; + turnState.lastAssistantBlockUuid = messageUuid; + return startsNewBlock; +} + /** * Reset all per-block streaming bookkeeping at an assistant-turn boundary. * @@ -187,4 +279,8 @@ export function resetTurnBlockState(turnState) { turnState.textBlockContentByIndex = new Map(); turnState.thinkingBlockContentByIndex = new Map(); turnState.blockStreamModeByKey = new Map(); + turnState.lastStreamDeltaByKey = new Map(); + turnState.lastSnapshotPayloadByKey = new Map(); + turnState.lastAssistantMessageId = null; + turnState.lastAssistantBlockUuid = null; } diff --git a/ai-bridge/services/claude/stream-event-processor.js b/ai-bridge/services/claude/stream-event-processor.js index 6f397065..93bd0d9c 100644 --- a/ai-bridge/services/claude/stream-event-processor.js +++ b/ai-bridge/services/claude/stream-event-processor.js @@ -1,6 +1,11 @@ import { emitAccumulatedUsage, mergeUsage } from '../../utils/usage-utils.js'; import { truncateErrorContent, truncateToolResultBlock } from './message-output-filter.js'; -import { normalizeStreamDelta, resolveSnapshotDelta, resetTurnBlockState } from './stream-delta-normalizer.js'; +import { + normalizeStreamDelta, + resolveSnapshotDelta, + resetTurnBlockState, + prepareAssistantSnapshotBlock, +} from './stream-delta-normalizer.js'; export function emitUsageTag(msg) { if (msg.type === 'assistant' && msg.message?.usage) { @@ -62,6 +67,17 @@ export function processStreamEvent(msg, turnState) { } } + if (event.type === 'content_block_start' && event.content_block) { + const hasPreviousBlock = + (turnState.textBlockContentByIndex instanceof Map && turnState.textBlockContentByIndex.size > 0) + || (turnState.thinkingBlockContentByIndex instanceof Map + && turnState.thinkingBlockContentByIndex.size > 0); + if (hasPreviousBlock) { + resetTurnBlockState(turnState); + process.stdout.write('[BLOCK_RESET]\n'); + } + } + if (event.type === 'message_delta' && event.usage) { turnState.accumulatedUsage = mergeUsage(turnState.accumulatedUsage, event.usage); emitAccumulatedUsage(turnState.accumulatedUsage); @@ -92,12 +108,23 @@ export function processMessageContent(msg, turnState) { for (let i = 0; i < content.length; i += 1) { const block = content[i]; if (block.type === 'text') { - emitSnapshotText(block.text || '', turnState, i); + const text = block.text || ''; + if (prepareAssistantSnapshotBlock(turnState, 'text', i, text, msg)) { + process.stdout.write('[BLOCK_RESET]\n'); + } + emitSnapshotText(text, turnState, i); } else if (block.type === 'thinking') { - emitSnapshotThinking(block.thinking || block.text || '', turnState, i); + const thinking = block.thinking || block.text || ''; + if (prepareAssistantSnapshotBlock(turnState, 'thinking', i, thinking, msg)) { + process.stdout.write('[BLOCK_RESET]\n'); + } + emitSnapshotThinking(thinking, turnState, i); } } } else if (typeof content === 'string') { + if (prepareAssistantSnapshotBlock(turnState, 'text', 0, content, msg)) { + process.stdout.write('[BLOCK_RESET]\n'); + } emitSnapshotText(content, turnState, 0); } } diff --git a/ai-bridge/services/claude/stream-event-processor.test.js b/ai-bridge/services/claude/stream-event-processor.test.js index aa7a5d24..c82a5fd9 100644 --- a/ai-bridge/services/claude/stream-event-processor.test.js +++ b/ai-bridge/services/claude/stream-event-processor.test.js @@ -138,6 +138,82 @@ test('shouldOutputMessage: streaming assistant with multiple tool_use blocks ret assert.equal(shouldOutputMessage(msg, state), true); }); +test('REGRESSION: non-cumulative assistant blocks keep thinking boundaries within one response', () => { + const state = makeTurnState(true); + const chunks = [ + '**Planning document updates for fixes**\n**Analyzing architecture for Nacos config**', + '**Defining MFA role configuration semantics**\n**Designing role-based MFA configuration map**', + '**Defining explicit role-based MFA configuration**\n**Mapping front-end and backend MFA configs**', + ]; + + const captured = captureStdout(() => { + chunks.forEach((thinking, index) => { + processMessageContent( + { + type: 'assistant', + uuid: `assistant-block-${index}`, + message: { + id: 'response-1', + content: [{ type: 'thinking', thinking }], + }, + }, + state, + ); + }); + }); + + const resetLines = tagLines(captured, '[BLOCK_RESET]'); + const thinkingLines = tagLines(captured, '[THINKING_DELTA]'); + const emitted = thinkingLines + .map((line) => JSON.parse(line.replace(/^\[THINKING_DELTA\]\s+/, '').trim())) + .join(''); + + assert.equal(resetLines.length, chunks.length - 1, 'each later normalized block needs one boundary'); + assert.equal(thinkingLines.length, chunks.length); + assert.equal(emitted, chunks.join('')); +}); + +test('REGRESSION: assistant snapshots replaying stream thinking deltas are absorbed', () => { + const state = makeTurnState(true); + const chunks = ['A', 'B', 'C']; + + const captured = captureStdout(() => { + chunks.forEach((thinking, index) => { + processStreamEvent( + { + type: 'stream_event', + event: { + type: 'content_block_delta', + index: 0, + delta: { type: 'thinking_delta', thinking }, + }, + }, + state, + ); + processMessageContent( + { + type: 'assistant', + uuid: `assistant-block-${index}`, + message: { + id: 'response-1', + content: [{ type: 'thinking', thinking }], + }, + }, + state, + ); + }); + }); + + const thinkingLines = tagLines(captured, '[THINKING_DELTA]'); + const emitted = thinkingLines + .map((line) => JSON.parse(line.replace(/^\[THINKING_DELTA\]\s+/, '').trim())) + .join(''); + + assert.equal(thinkingLines.length, chunks.length, 'snapshot replays must not emit duplicate deltas'); + assert.equal(emitted, 'ABC'); + assert.equal(tagLines(captured, '[BLOCK_RESET]').length, 0); +}); + test('end-to-end: streaming pure-text response emits no [MESSAGE], no duplicate [CONTENT_DELTA]', () => { const state = makeTurnState(true); @@ -1162,3 +1238,36 @@ test('REGRESSION (#1371) companion: snapshot path absorbs incoming === previous assert.equal(deltaLines.length, 2, `snapshot replay must not emit; got ${JSON.stringify(deltaLines)}`); assert.equal(emitted, 'Hello world', `accumulated content must remain "Hello world"; got "${emitted}"`); }); + +test('BLOCK_RESET: later content_block_start separates streamed thinking blocks', () => { + const state = makeTurnState(true); + const captured = captureStdout(() => { + for (const thinking of ['first', 'second']) { + processStreamEvent( + { + type: 'stream_event', + event: { type: 'content_block_start', index: 0, content_block: { type: 'thinking' } }, + }, + state, + ); + processStreamEvent( + { + type: 'stream_event', + event: { + type: 'content_block_delta', + index: 0, + delta: { type: 'thinking_delta', thinking }, + }, + }, + state, + ); + } + }); + + assert.equal(tagLines(captured, '[BLOCK_RESET]').length, 1); + assert.deepEqual( + tagLines(captured, '[THINKING_DELTA]').map((line) => + JSON.parse(line.replace(/^\[THINKING_DELTA\]\s+/, '').trim())), + ['first', 'second'], + ); +}); diff --git a/src/main/java/com/github/claudecodegui/provider/claude/ClaudeProcessInvoker.java b/src/main/java/com/github/claudecodegui/provider/claude/ClaudeProcessInvoker.java index aac5e2d5..bd581262 100644 --- a/src/main/java/com/github/claudecodegui/provider/claude/ClaudeProcessInvoker.java +++ b/src/main/java/com/github/claudecodegui/provider/claude/ClaudeProcessInvoker.java @@ -272,6 +272,7 @@ class ClaudeProcessInvoker { || line.startsWith("[TOOL_RESULT]") || line.startsWith("[USAGE]") || line.startsWith("[MESSAGE_START]") + || line.startsWith("[BLOCK_RESET]") || line.startsWith("[MESSAGE_END]"); } diff --git a/src/main/java/com/github/claudecodegui/session/MessageMerger.java b/src/main/java/com/github/claudecodegui/session/MessageMerger.java index 6ef4c093..825e7d50 100644 --- a/src/main/java/com/github/claudecodegui/session/MessageMerger.java +++ b/src/main/java/com/github/claudecodegui/session/MessageMerger.java @@ -230,7 +230,7 @@ public class MessageMerger { // prefix relation, whereas the lenient suffix-prefix overlap would fire // on incidental shared boundaries (code fences, Markdown markers) and // wrongly merge a new segment into the previous one. - return textLooksRelatedStrict(getTextContent(existingBlock), getTextContent(incomingBlock)); + return contentLooksRelatedStrict(getTextContent(existingBlock), getTextContent(incomingBlock)); } if ("thinking".equals(type)) { @@ -241,7 +241,11 @@ public class MessageMerger { if (existingThinking.isEmpty() || incomingThinking.isEmpty()) { return true; } - return textLooksRelated(existingThinking, incomingThinking); + // Thinking blocks can cross the same segment boundaries as text blocks. + // A suffix-prefix overlap is especially easy to trigger with Markdown + // markers (for example, adjacent "**...**" summaries), so only a + // prefix-related snapshot may update the existing block. + return contentLooksRelatedStrict(existingThinking, incomingThinking); } return existingBlock.equals(incomingBlock); @@ -296,56 +300,28 @@ public class MessageMerger { return getTextContent(block); } - // Whether two non-empty texts are equal or one is a prefix of the other: - // the "same segment, possibly grown" relation shared by the strict text - // check and the lenient thinking check's prefix stage. + // Whether two non-empty block contents are equal or one is a prefix of the other: + // both relations describe one segment that is being filled by a fuller snapshot. private boolean isPrefixRelated(String existing, String incoming) { return existing.equals(incoming) || existing.startsWith(incoming) || incoming.startsWith(existing); } - // Strict prefix-only relatedness for text blocks across segments. Omits the - // suffix-prefix overlap that textLooksRelated keeps for fragmented thinking: - // for text, overlap frequently fires on incidental shared boundaries (code - // fences, Markdown markers) between two segments separated by a tool_use, - // wrongly merging a new segment into the previous one. - private boolean textLooksRelatedStrict(String existingText, String incomingText) { + // Strict prefix-only relatedness for unkeyed text/thinking blocks. Omitting + // suffix-prefix overlap prevents incidental shared boundaries (for example, + // Markdown markers) from joining two independent streaming segments. + private boolean contentLooksRelatedStrict(String existingText, String incomingText) { String existing = existingText != null ? existingText : ""; String incoming = incomingText != null ? incomingText : ""; // isPrefixRelated already treats an empty string as a prefix of any string, - // so an empty text block and a non-empty one are the same segment (the empty + // so an empty block and a non-empty one are the same segment (the empty // one is the segment's leading edge before content arrives). This lets a // later, fuller snapshot fill an empty placeholder instead of duplicating - // it, mirroring the thinking branch's empty-is-same-segment rule. + // it. return isPrefixRelated(existing, incoming); } - private boolean textLooksRelated(String existingText, String incomingText) { - String existing = existingText != null ? existingText : ""; - String incoming = incomingText != null ? incomingText : ""; - - if (existing.isEmpty() || incoming.isEmpty()) { - return existing.isEmpty() && incoming.isEmpty(); - } - - if (isPrefixRelated(existing, incoming)) { - return true; - } - - // Check suffix-prefix overlap (streaming may produce partial overlaps) - int maxOverlap = Math.min(existing.length(), incoming.length()); - maxOverlap = Math.min(maxOverlap, 200); - int eLen = existing.length(); - for (int overlap = maxOverlap; overlap > 0; overlap--) { - if (existing.regionMatches(eLen - overlap, incoming, 0, overlap)) { - return true; - } - } - - return false; - } - private String preferMoreCompleteContent(String existingText, String incomingText) { String existing = existingText != null ? existingText : ""; String incoming = incomingText != null ? incomingText : ""; diff --git a/src/main/java/com/github/claudecodegui/session/SessionCallbackAdapter.java b/src/main/java/com/github/claudecodegui/session/SessionCallbackAdapter.java index 3af159f3..672d8b73 100644 --- a/src/main/java/com/github/claudecodegui/session/SessionCallbackAdapter.java +++ b/src/main/java/com/github/claudecodegui/session/SessionCallbackAdapter.java @@ -343,7 +343,12 @@ public class SessionCallbackAdapter implements ClaudeSession.SessionCallback { if (isInactive()) { return; } - // Reset throttlers for the new turn's deltas + // Flush BEFORE resetting: block boundaries now fire mid-response (one per + // content-block edge, not just per tool-loop turn), so deltas buffered in + // the throttlers belong to the ending block. reset() alone would silently + // drop them and force the frontend to fall back to updateMessages snapshots. + contentDeltaThrottler.flushNow(); + thinkingDeltaThrottler.flushNow(); contentDeltaThrottler.reset(); thinkingDeltaThrottler.reset(); ApplicationManager.getApplication().invokeLater(() -> { diff --git a/src/test/java/com/github/claudecodegui/provider/claude/ClaudeSDKBridgeRefactorTest.java b/src/test/java/com/github/claudecodegui/provider/claude/ClaudeSDKBridgeRefactorTest.java index e0e8a9d7..8b66ab30 100644 --- a/src/test/java/com/github/claudecodegui/provider/claude/ClaudeSDKBridgeRefactorTest.java +++ b/src/test/java/com/github/claudecodegui/provider/claude/ClaudeSDKBridgeRefactorTest.java @@ -8,6 +8,7 @@ import com.google.gson.Gson; import com.google.gson.JsonObject; import org.junit.Test; +import java.lang.reflect.Method; import java.util.ArrayList; import java.util.List; import java.util.concurrent.CompletableFuture; @@ -128,6 +129,42 @@ public class ClaudeSDKBridgeRefactorTest { assertEquals(null, lastNodeError.get()); } + @Test + public void streamAdapterRoutesBlockResetLines() { + ClaudeStreamAdapter adapter = new ClaudeStreamAdapter(new Gson()); + RecordingCallback callback = new RecordingCallback(); + SDKResult result = new SDKResult(); + AtomicBoolean hadSendError = new AtomicBoolean(false); + AtomicReference lastNodeError = new AtomicReference<>(null); + AtomicBoolean wasAborted = new AtomicBoolean(false); + + adapter.processOutputLine("[BLOCK_RESET]", callback, result, new StringBuilder(), + hadSendError, lastNodeError, wasAborted); + + assertEquals(1, callback.events.size()); + assertEquals("block_reset", callback.events.get(0).type); + assertEquals("", callback.events.get(0).payload); + } + + @Test + public void processInvokerRecognizesBlockResetAsBridgeLine() throws Exception { + ClaudeProcessInvoker invoker = new ClaudeProcessInvoker( + null, + new Gson(), + null, + null, + null, + null, + null, + null, + new ClaudeStreamAdapter(new Gson()) + ); + Method method = ClaudeProcessInvoker.class.getDeclaredMethod("isRecognizedBridgeLine", String.class); + method.setAccessible(true); + + assertTrue((Boolean) method.invoke(invoker, "[BLOCK_RESET]")); + } + @Test public void streamAdapterMarksSendErrorsAndPreservesParsedMessage() { ClaudeStreamAdapter adapter = new ClaudeStreamAdapter(new Gson()); diff --git a/src/test/java/com/github/claudecodegui/session/ClaudeMessageHandlerDedupTest.java b/src/test/java/com/github/claudecodegui/session/ClaudeMessageHandlerDedupTest.java index 6c3f55ca..cbea56b0 100644 --- a/src/test/java/com/github/claudecodegui/session/ClaudeMessageHandlerDedupTest.java +++ b/src/test/java/com/github/claudecodegui/session/ClaudeMessageHandlerDedupTest.java @@ -2,6 +2,7 @@ package com.github.claudecodegui.session; import com.google.gson.Gson; import com.google.gson.GsonBuilder; +import com.google.gson.JsonArray; import org.junit.Before; import org.junit.Test; @@ -38,6 +39,23 @@ public class ClaudeMessageHandlerDedupTest { ); } + @Test + public void blockReset_startsNewThinkingBlockForIndependentAssistantTurn() { + handler.onMessage("stream_start", ""); + handler.onMessage("thinking_delta", "first thought"); + handler.onMessage("block_reset", ""); + handler.onMessage("thinking_delta", "second thought"); + + List messages = callbackHandler.messageUpdates.get( + callbackHandler.messageUpdates.size() - 1 + ); + JsonArray content = messages.get(0).raw.getAsJsonObject("message").getAsJsonArray("content"); + + assertEquals("Independent thinking turns must remain separate blocks", 2, content.size()); + assertEquals("first thought", content.get(0).getAsJsonObject().get("thinking").getAsString()); + assertEquals("second thought", content.get(1).getAsJsonObject().get("thinking").getAsString()); + } + /** * Test that content delta is skipped when it duplicates existing content after conservative sync. */ @@ -300,6 +318,7 @@ public class ClaudeMessageHandlerDedupTest { private static class RecordingCallbackHandler extends CallbackHandler { final List contentDeltas = new ArrayList<>(); final List thinkingDeltas = new ArrayList<>(); + final List> messageUpdates = new ArrayList<>(); int streamStartCount = 0; int streamEndCount = 0; @@ -308,6 +327,11 @@ public class ClaudeMessageHandlerDedupTest { thinkingDeltas.clear(); } + @Override + public void notifyMessageUpdate(List messages) { + messageUpdates.add(messages); + } + @Override public void notifyContentDelta(String delta) { contentDeltas.add(delta); diff --git a/src/test/java/com/github/claudecodegui/session/MessageMergerTest.java b/src/test/java/com/github/claudecodegui/session/MessageMergerTest.java index 84ed3093..52a655f3 100644 --- a/src/test/java/com/github/claudecodegui/session/MessageMergerTest.java +++ b/src/test/java/com/github/claudecodegui/session/MessageMergerTest.java @@ -227,6 +227,32 @@ public class MessageMergerTest { assertEquals("bash-1", mergedContent.get(2).getAsJsonObject().get("id").getAsString()); } + @Test + public void mergeAssistantMessageKeepsMarkdownThinkingBlocksSeparatedAcrossSegments() { + MessageMerger merger = new MessageMerger(); + + JsonObject existingRaw = assistantMessage( + thinkingBlock("**Inspecting multi-factor login controller toggles**") + ); + JsonObject newRaw = assistantMessage( + thinkingBlock("**Investigating login filter and authentication flow**") + ); + + JsonArray mergedContent = merger.mergeAssistantMessage(existingRaw, newRaw) + .getAsJsonObject("message") + .getAsJsonArray("content"); + + assertEquals("Independent thinking snapshots must remain separate blocks", 2, mergedContent.size()); + assertEquals( + "**Inspecting multi-factor login controller toggles**", + mergedContent.get(0).getAsJsonObject().get("thinking").getAsString() + ); + assertEquals( + "**Investigating login filter and authentication flow**", + mergedContent.get(1).getAsJsonObject().get("thinking").getAsString() + ); + } + @Test public void mergeAssistantMessageKeepsMoreCompleteThinkingBlockFromSnapshot() { MessageMerger merger = new MessageMerger(); diff --git a/src/test/java/com/github/claudecodegui/session/SessionCallbackAdapterStreamEndTest.java b/src/test/java/com/github/claudecodegui/session/SessionCallbackAdapterStreamEndTest.java index a28659bf..cfaa74da 100644 --- a/src/test/java/com/github/claudecodegui/session/SessionCallbackAdapterStreamEndTest.java +++ b/src/test/java/com/github/claudecodegui/session/SessionCallbackAdapterStreamEndTest.java @@ -1,7 +1,11 @@ package com.github.claudecodegui.session; +import com.intellij.openapi.application.Application; +import com.intellij.openapi.application.ApplicationManager; +import org.jetbrains.annotations.NotNull; import org.junit.Test; +import java.lang.reflect.Proxy; import java.util.ArrayList; import java.util.List; import java.util.concurrent.atomic.AtomicLong; @@ -176,5 +180,73 @@ public class SessionCallbackAdapterStreamEndTest { assertEquals(1, jsTarget.calls.size()); assertEquals("onStreamEnd:77", jsTarget.calls.get(0)); } + + /** + * Regression: block boundaries now fire mid-response (one per content-block + * edge), so deltas still buffered in the delta throttlers when onBlockReset + * runs belong to the ending block. The adapter must flush them to the + * frontend BEFORE resetting; a bare reset() silently drops the buffered tail + * and forces the frontend onto updateMessages snapshot rendering (visible as + * thinking text jumping in chunks instead of streaming). + * + *

Exercised through the real SessionCallbackAdapter with test-friendly + * collaborators; the throttlers' default constructor flushes synchronously + * via flushNow(), so no IntelliJ Application is needed — the headless + * invokeLater in onBlockReset happens after the ordering under test. + */ + @Test + public void blockResetFlushesBufferedDeltasBeforeClearing() throws Exception { + RecordingJsTarget jsTarget = new RecordingJsTarget(); + SessionCallbackAdapter adapter = new SessionCallbackAdapter( + null, + jsTarget, + null, + () -> true, + null + ); + + // Deltas arrive and sit in the throttlers' 33ms window... + adapter.onContentDelta("text-tail"); + adapter.onThinkingDelta("thinking-tail"); + + // onBlockReset dispatches its JS notification via invokeLater, which has + // no Application in headless tests. The flush-before-reset ordering under + // test completes before that call, so a benign proxy stub suffices. + // Not restored afterwards: setApplication(null, ...) is rejected by the + // platform's @NotNull contract, and each Gradle test fork owns its JVM. + ApplicationManager.setApplication(invokeLaterInlineApplication()); + adapter.onBlockReset(); + + assertTrue(jsTarget.calls.contains("onContentDelta:text-tail")); + assertTrue(jsTarget.calls.contains("onThinkingDelta:thinking-tail")); + adapter.deactivate(); + } + + /** + * Headless Application stub whose invokeLater runs the runnable inline — the + * adapter only needs a non-null application for its EDT dispatch. + */ + private static @NotNull Application invokeLaterInlineApplication() { + return (Application) Proxy.newProxyInstance( + Application.class.getClassLoader(), + new Class[] { Application.class }, + (proxy, method, args) -> { + if ("invokeLater".equals(method.getName()) && args != null && args.length >= 1) { + ((Runnable) args[0]).run(); + return null; + } + if (method.getName().equals("isDispatchThread")) { + return Boolean.TRUE; + } + // Unimplemented platform calls are irrelevant to this test; + // returning defaults keeps the stub minimal. + Class type = method.getReturnType(); + if (type == boolean.class) { + return Boolean.FALSE; + } + return null; + } + ); + } }