Merge pull request #1733 from gadfly3173/fix/claude-thinking-block-boundaries

fix(claude): preserve thinking block boundaries in streaming output

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') {

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;
 }

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);
   }
 }

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'],
+  );
+});

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]");
     }
 

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 : "";

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(() -> {

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<String> 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());

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<ClaudeSession.Message> 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<String> contentDeltas = new ArrayList<>();
         final List<String> thinkingDeltas = new ArrayList<>();
+        final List<List<ClaudeSession.Message>> messageUpdates = new ArrayList<>();
         int streamStartCount = 0;
         int streamEndCount = 0;
 
@@ -308,6 +327,11 @@ public class ClaudeMessageHandlerDedupTest {
             thinkingDeltas.clear();
         }
 
+        @Override
+        public void notifyMessageUpdate(List<ClaudeSession.Message> messages) {
+            messageUpdates.add(messages);
+        }
+
         @Override
         public void notifyContentDelta(String delta) {
             contentDeltas.add(delta);

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();

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).
+     *
+     * <p>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;
+                }
+        );
+    }
 }