Reuse stream indexes across delta updates

src/features/chat/application/state/thread-stream.ts

@@ -441,7 +441,17 @@ function shouldUseActiveSegment(
 }
 
 function appendActiveSegmentItem(segment: ChatThreadStreamActiveSegment, item: ThreadStreamItem): ChatThreadStreamActiveSegment {
-  return activeSegmentFromItems(segment.turnId, [...segment.items, item]);
+  const index = segment.items.length;
+  const indexById = new Map(segment.indexById);
+  indexById.set(item.id, index);
+  const indexBySourceItemId = new Map(segment.indexBySourceItemId);
+  if (item.sourceItemId) indexBySourceItemId.set(item.sourceItemId, index);
+  return {
+    turnId: segment.turnId,
+    items: [...segment.items, item],
+    indexById,
+    indexBySourceItemId,
+  };
 }
 
 function upsertActiveSegmentItem(segment: ChatThreadStreamActiveSegment, item: ThreadStreamItem): ChatThreadStreamActiveSegment {
@@ -461,6 +471,9 @@ function replaceActiveSegmentItem(
   if (next === previous) return segment;
   const items = [...segment.items];
   items[index] = next;
+  if (next.id === previous.id && next.sourceItemId === previous.sourceItemId) {
+    return { ...segment, items };
+  }
   return activeSegmentFromItems(segment.turnId, items);
 }
 

tests/features/chat/application/state/thread-stream.test.ts

@@ -1,13 +1,43 @@
 import { describe, expect, it } from "vitest";
 import {
   initialChatThreadStreamState,
+  reduceThreadStreamSlice,
   threadStreamRollbackCandidate,
+  threadStreamStartActiveSegment,
   threadStreamTurnsAfterTurnId,
   threadStreamWithActiveTurnItems,
 } from "../../../../../src/features/chat/application/state/thread-stream";
 import type { ThreadStreamItem } from "../../../../../src/features/chat/domain/thread-stream/items";
 
 describe("thread stream selectors", () => {
+  it("reuses active-segment indexes while appending deltas to an existing item", () => {
+    const initial = threadStreamStartActiveSegment(initialChatThreadStreamState(), "turn-1", [
+      {
+        id: "assistant-1",
+        sourceItemId: "source-1",
+        kind: "dialogue",
+        dialogueKind: "assistantResponse",
+        dialogueState: "streaming",
+        role: "assistant",
+        text: "hello",
+        turnId: "turn-1",
+      },
+    ]);
+    const previousSegment = initial.activeSegment;
+    if (!previousSegment) throw new Error("Expected active segment");
+
+    const next = reduceThreadStreamSlice(initial, {
+      type: "thread-stream/assistant-delta-appended",
+      itemId: "source-1",
+      turnId: "turn-1",
+      delta: " world",
+    });
+
+    expect(next.activeSegment?.items[0]).toMatchObject({ text: "hello world" });
+    expect(next.activeSegment?.indexById).toBe(previousSegment.indexById);
+    expect(next.activeSegment?.indexBySourceItemId).toBe(previousSegment.indexBySourceItemId);
+  });
+
   it("counts turns after a turn id from thread stream state", () => {
     const state = initialChatThreadStreamState(items());