revert(vscode-ide-companion): undo #3450 split-stream timestamp sharing

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<typeof useMessageHandling>;
@@ -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,24 +98,89 @@ 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(
       (message): message is TextMessage => message.role === 'assistant',
     );
 
     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);
+    },
+  );
 });

packages/vscode-ide-companion/src/webview/hooks/message/useMessageHandling.ts

@@ -35,8 +35,6 @@ export const useMessageHandling = () => {
   const streamingMessageIndexRef = useRef<number | null>(null);
   // Track the index of the current aggregated thinking message
   const thinkingMessageIndexRef = useRef<number | null>(null);
-  // Preserve one stable timestamp for all message segments in the same turn.
-  const currentStreamTimestampRef = useRef<number | null>(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: '',