diff --git a/packages/vscode-ide-companion/src/webview/hooks/message/useMessageHandling.test.tsx b/packages/vscode-ide-companion/src/webview/hooks/message/useMessageHandling.test.tsx index 3c805b4fb..cf1921407 100644 --- a/packages/vscode-ide-companion/src/webview/hooks/message/useMessageHandling.test.tsx +++ b/packages/vscode-ide-companion/src/webview/hooks/message/useMessageHandling.test.tsx @@ -8,7 +8,7 @@ import { act } from 'react'; import { createRoot, type Root } from 'react-dom/client'; -import { afterEach, beforeEach, describe, expect, it } from 'vitest'; +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { useMessageHandling, type TextMessage } from './useMessageHandling.js'; type MessageHandlingApi = ReturnType; @@ -41,6 +41,26 @@ function renderHookHarness() { }; } +/** + * The webview merges text messages and tool calls and sorts them by + * `timestamp` for rendering. Two known bugs push that sort in opposite + * directions: + * + * - Tool-call interleave: a tool call that arrives between two assistant + * segments of the same turn must sort strictly between them. This + * requires seg1.ts < toolCall.ts < seg2.ts. + * + * - #3273 (user question appears above the previous assistant answer): + * a user message belonging to a later turn must sort after every + * segment / tool call of the previous turn, even if the later segment + * was created after the user message was added. This requires all + * segments of turn N to be strictly less than any message/tool call + * from turn N+1. + * + * A single timestamp strategy cannot satisfy both simultaneously without a + * monotonic-sequence layer. These tests pin the current behaviour of the + * hook so we can wire a proper fix in next. + */ describe('useMessageHandling', () => { let root: Root | null = null; let container: HTMLDivElement | null = null; @@ -52,6 +72,7 @@ describe('useMessageHandling', () => { }); afterEach(() => { + vi.useRealTimers(); if (root) { act(() => { root?.unmount(); @@ -64,7 +85,10 @@ describe('useMessageHandling', () => { } }); - it('keeps the original stream timestamp when a tool call splits one assistant reply into multiple segments', () => { + it('assigns the second assistant segment a newer timestamp so a tool call can sort between the two segments', () => { + vi.useFakeTimers(); + vi.setSystemTime(1_000); + const rendered = renderHookHarness(); root = rendered.root; container = rendered.container; @@ -74,15 +98,20 @@ describe('useMessageHandling', () => { }); act(() => { - rendered.api.appendStreamChunk('before tool call'); + rendered.api.appendStreamChunk('seg1'); }); + // Tool call arrives and triggers a segment break. + vi.setSystemTime(2_000); + const toolCallTimestamp = Date.now(); act(() => { rendered.api.breakAssistantSegment(); }); + // Next chunk lands after the tool call. + vi.setSystemTime(3_000); act(() => { - rendered.api.appendStreamChunk('after tool call'); + rendered.api.appendStreamChunk('seg2'); }); const assistantMessages = rendered.api.messages.filter( @@ -90,8 +119,68 @@ describe('useMessageHandling', () => { ); expect(assistantMessages).toHaveLength(2); - expect(assistantMessages.map((message) => message.timestamp)).toEqual([ - 1_000, 1_000, - ]); + expect(assistantMessages[0].timestamp).toBeLessThan(toolCallTimestamp); + expect(assistantMessages[1].timestamp).toBeGreaterThan(toolCallTimestamp); }); + + it.fails( + 'keeps every assistant segment of a turn before a user message that was sent between segments (#3273)', + () => { + // Reproduces the race that #3273 describes: React batching (or any + // other delay) causes the second segment's placeholder to materialize + // AFTER the next user message has already been pushed into the list. + // Because the placeholder uses Date.now() at materialization time, it + // receives a timestamp greater than the new user message, and the + // user bubble ends up sandwiched between the two assistant segments + // once the list is sorted. + vi.useFakeTimers(); + vi.setSystemTime(1_000); + + const rendered = renderHookHarness(); + root = rendered.root; + container = rendered.container; + + act(() => { + rendered.api.startStreaming(1_000); + }); + + act(() => { + rendered.api.appendStreamChunk('seg1'); + }); + + vi.setSystemTime(2_000); + act(() => { + rendered.api.breakAssistantSegment(); + }); + + // The user types and sends their next question before the delayed + // second-segment chunk is flushed. + vi.setSystemTime(3_000); + act(() => { + rendered.api.addMessage({ + role: 'user', + content: 'next question', + timestamp: Date.now(), + }); + }); + + // Second-segment chunk finally runs. + vi.setSystemTime(4_000); + act(() => { + rendered.api.appendStreamChunk('seg2'); + }); + + const sorted = [...rendered.api.messages].sort( + (a, b) => a.timestamp - b.timestamp, + ); + const roles = sorted.map((m) => m.role); + const userIdx = roles.indexOf('user'); + const lastAssistantIdx = roles.lastIndexOf('assistant'); + + // Expected (and currently violated) invariant: the user message marks + // the start of a new turn, so every assistant segment of the previous + // turn must come before it. + expect(lastAssistantIdx).toBeLessThan(userIdx); + }, + ); }); diff --git a/packages/vscode-ide-companion/src/webview/hooks/message/useMessageHandling.ts b/packages/vscode-ide-companion/src/webview/hooks/message/useMessageHandling.ts index 0ac2726ba..ff6af0125 100644 --- a/packages/vscode-ide-companion/src/webview/hooks/message/useMessageHandling.ts +++ b/packages/vscode-ide-companion/src/webview/hooks/message/useMessageHandling.ts @@ -35,8 +35,6 @@ export const useMessageHandling = () => { const streamingMessageIndexRef = useRef(null); // Track the index of the current aggregated thinking message const thinkingMessageIndexRef = useRef(null); - // Preserve one stable timestamp for all message segments in the same turn. - const currentStreamTimestampRef = useRef(null); /** * Add message @@ -56,9 +54,6 @@ export const useMessageHandling = () => { * Start streaming response */ const startStreaming = useCallback((timestamp?: number) => { - const resolvedTimestamp = - typeof timestamp === 'number' ? timestamp : Date.now(); - currentStreamTimestampRef.current = resolvedTimestamp; // Create an assistant placeholder message immediately so tool calls won't jump before it setMessages((prev) => { // Record index of the placeholder to update on chunks @@ -68,8 +63,8 @@ export const useMessageHandling = () => { { role: 'assistant', content: '', - // Use one stable turn timestamp so later split segments sort correctly. - timestamp: resolvedTimestamp, + // Use provided timestamp (from extension) to keep ordering stable + timestamp: typeof timestamp === 'number' ? timestamp : Date.now(), }, ]; }); @@ -94,11 +89,7 @@ export const useMessageHandling = () => { if (idx === null) { idx = next.length; streamingMessageIndexRef.current = idx; - next.push({ - role: 'assistant', - content: '', - timestamp: currentStreamTimestampRef.current ?? Date.now(), - }); + next.push({ role: 'assistant', content: '', timestamp: Date.now() }); } if (idx < 0 || idx >= next.length) { @@ -131,7 +122,6 @@ export const useMessageHandling = () => { setIsStreaming(false); streamingMessageIndexRef.current = null; thinkingMessageIndexRef.current = null; - currentStreamTimestampRef.current = null; }, []); /** @@ -183,7 +173,7 @@ export const useMessageHandling = () => { assistantIdx >= 0 && assistantIdx < next.length ? next[assistantIdx].timestamp - : (currentStreamTimestampRef.current ?? Date.now()); + : Date.now(); next.push({ role: 'thinking', content: '',