fix: flush pending message_end on new turn, finalize orphaned current

dimakis · claude · dimakis · commit 7ef2843545b9 · 2026-04-03T02:06:22.000+01:00
In multi-turn streaming-input sessions, message_start for turn N+1
could arrive before the deferred message_end for turn N flushed.
The message_start handler reset pendingMessageEnd to null, silently
dropping turn N's message_end. On the frontend, MESSAGE_START
overwrote current without finalizing it into messages[].

Server: extract forceFlushPendingMessage() helper that force-closes
open blocks and emits the pending message_end. Called from both the
message_start handler (before resetting turn state) and the result
handler (before session_end). DRYs the duplicated flush logic.

Frontend: MESSAGE_START now finalizes any orphaned current into
messages[] before creating the new streaming message. Belt-and-
suspenders against any remaining server ordering edge cases.

Co-Authored-By: Claude Opus 4.6 &lt;noreply@anthropic.com&gt;
Made-with: Cursor
diff --git a/frontend/src/hooks/__tests__/chatMessagesReducer.test.ts b/frontend/src/hooks/__tests__/chatMessagesReducer.test.ts
@@ -23,14 +23,35 @@ describe('MESSAGE_START', () => {
     expect(state.current!.blocks.size).toBe(0);
   });
 
-  it('replaces an existing current when a new message starts', () => {
-    const withCurrent = chatMessagesReducer(INITIAL, { type: 'MESSAGE_START', messageId: 'msg-1' });
-    const replaced = chatMessagesReducer(withCurrent, {
-      type: 'MESSAGE_START',
-      messageId: 'msg-2',
-    });
-    expect(replaced.current!.messageId).toBe('msg-2');
-    expect(replaced.current!.blockOrder).toHaveLength(0);
+  it('finalizes orphaned current into messages when a new message starts', () => {
+    let state = chatMessagesReducer(INITIAL, { type: 'MESSAGE_START', messageId: 'msg-1' });
+    state = chatMessagesReducer(state, {
+      type: 'BLOCK_START',
+      messageId: 'msg-1',
+      blockId: 'b1',
+      blockType: 'text',
+    });
+    state = chatMessagesReducer(state, {
+      type: 'BLOCK_DELTA',
+      messageId: 'msg-1',
+      blockId: 'b1',
+      blockType: 'text',
+      delta: 'orphaned content',
+    });
+    // New message starts before old one got MESSAGE_END
+    state = chatMessagesReducer(state, { type: 'MESSAGE_START', messageId: 'msg-2' });
+    expect(state.current!.messageId).toBe('msg-2');
+    expect(state.current!.blockOrder).toHaveLength(0);
+    expect(state.messages).toHaveLength(1);
+    expect(state.messages[0].messageId).toBe('msg-1');
+    expect(state.messages[0].blocks[0].content).toBe('orphaned content');
+  });
+
+  it('creates new current cleanly when no prior current exists', () => {
+    const state = chatMessagesReducer(INITIAL, { type: 'MESSAGE_START', messageId: 'msg-2' });
+    expect(state.current!.messageId).toBe('msg-2');
+    expect(state.current!.blockOrder).toHaveLength(0);
+    expect(state.messages).toHaveLength(0);
   });
 });
 
diff --git a/frontend/src/hooks/useChatMessages.ts b/frontend/src/hooks/useChatMessages.ts
@@ -124,10 +124,16 @@ export function chatMessagesReducer(
 ): ChatMessagesState {
   switch (action.type) {
     case 'MESSAGE_START': {
-      const block = new Map<string, StreamingBlock>();
+      const base = state.current
+        ? { ...state, messages: [...state.messages, finishCurrent(state.current)] }
+        : state;
       return {
-        ...state,
-        current: { messageId: action.messageId, blocks: block, blockOrder: [] },
+        ...base,
+        current: {
+          messageId: action.messageId,
+          blocks: new Map<string, StreamingBlock>(),
+          blockOrder: [],
+        },
       };
     }
 
diff --git a/server/__tests__/query-loop.test.ts b/server/__tests__/query-loop.test.ts
@@ -292,6 +292,75 @@ describe('runQueryLoop', () => {
     expect(blockEndIdx).toBeLessThan(messageEndIdx);
   });
 
+  it('flushes pending message_end when new message_start arrives (multi-turn race)', async () => {
+    const events: Record<string, unknown>[] = [
+      // Turn 1: thinking block, assistant fires before block_stop
+      { type: 'stream_event', event: { type: 'message_start', message: { id: 'msg-t1' } } },
+      {
+        type: 'stream_event',
+        event: { type: 'content_block_start', index: 0, content_block: { type: 'thinking' } },
+      },
+      {
+        type: 'stream_event',
+        event: {
+          type: 'content_block_delta',
+          index: 0,
+          delta: { type: 'thinking_delta', thinking: 'hmm' },
+        },
+      },
+      { type: 'stream_event', event: { type: 'content_block_stop', index: 0 } },
+      {
+        type: 'stream_event',
+        event: { type: 'content_block_start', index: 1, content_block: { type: 'text' } },
+      },
+      {
+        type: 'stream_event',
+        event: {
+          type: 'content_block_delta',
+          index: 1,
+          delta: { type: 'text_delta', text: 'Turn 1 response' },
+        },
+      },
+      // assistant fires BEFORE content_block_stop for the text block
+      { type: 'assistant', message: { content: [] }, session_id: 'sess-mt' },
+      // Turn 2 starts before Turn 1's text block_stop arrives
+      { type: 'stream_event', event: { type: 'message_start', message: { id: 'msg-t2' } } },
+      // Turn 1's block_stop is now orphaned — server should have force-flushed
+      {
+        type: 'stream_event',
+        event: { type: 'content_block_start', index: 0, content_block: { type: 'text' } },
+      },
+      {
+        type: 'stream_event',
+        event: {
+          type: 'content_block_delta',
+          index: 0,
+          delta: { type: 'text_delta', text: 'Turn 2 response' },
+        },
+      },
+      { type: 'stream_event', event: { type: 'content_block_stop', index: 0 } },
+      { type: 'assistant', message: { content: [] }, session_id: 'sess-mt' },
+      { type: 'result', session_id: 'sess-mt' },
+    ];
+
+    await runQueryLoop(eventStream(events), clientId, registry, abortController, ws);
+
+    const sent = ws.sent;
+
+    // Turn 1 message_end must exist and come before Turn 2 message_start
+    const t1MsgEnd = sent.findIndex((m) => m.type === 'message_end' && m.messageId === 'msg-t1');
+    const t2MsgStart = sent.findIndex(
+      (m) => m.type === 'message_start' && m.messageId === 'msg-t2',
+    );
+    expect(t1MsgEnd).toBeGreaterThan(-1);
+    expect(t2MsgStart).toBeGreaterThan(-1);
+    expect(t1MsgEnd).toBeLessThan(t2MsgStart);
+
+    // Turn 2 message_end must also exist
+    const t2MsgEnd = sent.findIndex((m) => m.type === 'message_end' && m.messageId === 'msg-t2');
+    expect(t2MsgEnd).toBeGreaterThan(t2MsgStart);
+  });
+
   it('does not emit old-style text or text_delta events', async () => {
     const events: Record<string, unknown>[] = [
       { type: 'stream_event', event: { type: 'message_start', message: { id: 'msg-nodupe' } } },
diff --git a/server/query-loop.ts b/server/query-loop.ts
@@ -73,6 +73,55 @@ export async function runQueryLoop(
     }
   }
 
+  function forceFlushPendingMessage(
+    ws: WebSocket,
+    session: { currentSnapshot: { blocks: SnapshotBlock[] } | null },
+  ) {
+    if (!pendingMessageEnd) return;
+    for (const [index, bid] of blockIdByIndex) {
+      if (openBlockCount <= 0) break;
+      const snap = session.currentSnapshot?.blocks.find((b) => b.blockId === bid);
+      if (snap && !snap.done) {
+        snap.done = true;
+        const toolEntry = toolInputBuffers.get(index);
+        if (toolEntry) {
+          toolInputBuffers.delete(index);
+          let toolInput: Record<string, unknown> = {};
+          try {
+            toolInput = JSON.parse(toolEntry.inputBuf || '{}');
+          } catch {
+            /* empty */
+          }
+          send(
+            ws,
+            v2('block_end', {
+              messageId: currentMessageId,
+              blockId: bid,
+              blockType: 'tool_use',
+              toolName: toolEntry.name,
+              toolId: toolEntry.id,
+              input: summarizeToolInput(toolEntry.name, toolInput),
+            }),
+          );
+        } else {
+          send(
+            ws,
+            v2('block_end', {
+              messageId: currentMessageId,
+              blockId: bid,
+              blockType: snap.blockType,
+            }),
+          );
+        }
+        openBlockCount = Math.max(0, openBlockCount - 1);
+      }
+    }
+    send(ws, pendingMessageEnd);
+    pendingMessageEnd = null;
+    currentMessageId = null;
+    session.currentSnapshot = null;
+  }
+
   log.info('query loop started', { clientId });
 
   try {
@@ -100,63 +149,19 @@ export async function runQueryLoop(
       } else if (msg.type === 'result') {
         log.info('result received', { clientId, sessionId: msg.session_id });
         if (msg.session_id) send(currentWs, { type: 'session_id', sessionId: msg.session_id });
-        // Force-close any open blocks and flush pending message_end before session_end.
-        if (pendingMessageEnd) {
-          for (const [index, bid] of blockIdByIndex) {
-            if (openBlockCount <= 0) break;
-            const snap = currentSession.currentSnapshot?.blocks.find((b) => b.blockId === bid);
-            if (snap && !snap.done) {
-              snap.done = true;
-              const toolEntry = toolInputBuffers.get(index);
-              if (toolEntry) {
-                toolInputBuffers.delete(index);
-                let toolInput: Record<string, unknown> = {};
-                try {
-                  toolInput = JSON.parse(toolEntry.inputBuf || '{}');
-                } catch {
-                  /* empty */
-                }
-                send(
-                  currentWs,
-                  v2('block_end', {
-                    messageId: currentMessageId,
-                    blockId: bid,
-                    blockType: 'tool_use',
-                    toolName: toolEntry.name,
-                    toolId: toolEntry.id,
-                    input: summarizeToolInput(toolEntry.name, toolInput),
-                  }),
-                );
-              } else {
-                send(
-                  currentWs,
-                  v2('block_end', {
-                    messageId: currentMessageId,
-                    blockId: bid,
-                    blockType: snap.blockType,
-                  }),
-                );
-              }
-              openBlockCount = Math.max(0, openBlockCount - 1);
-            }
-          }
-          send(currentWs, pendingMessageEnd);
-          pendingMessageEnd = null;
-          currentMessageId = null;
-          currentSession.currentSnapshot = null;
-        }
+        forceFlushPendingMessage(currentWs, currentSession);
         doneSent = true;
         send(currentWs, v2('session_end', { sessionId: msg.session_id }));
       } else if (msg.type === 'stream_event') {
         const evt = msg.event as Record<string, unknown> | undefined;
         log.debug('stream event', { clientId, evtType: evt?.type });
 
         if (evt?.type === 'message_start') {
+          forceFlushPendingMessage(currentWs, currentSession);
           toolInputBuffers.clear();
           blockIdByIndex.clear();
           blockCounter = 0;
           openBlockCount = 0;
-          pendingMessageEnd = null;
           // Use API message ID if available, otherwise generate one.
           const apiMsg = evt.message as Record<string, unknown> | undefined;
           currentMessageId = (apiMsg?.id as string | undefined) ?? `msg-${Date.now()}`;
