refactor(chat): store canonical thread snapshots

Author: FranciscoMoretti

apps/chat/lib/stores/with-threads.test.ts

@@ -13,6 +13,7 @@ function createMessage({
   parallelGroupId = null,
   parallelIndex = null,
   activeStreamId = null,
+  text = "",
 }: {
   id: string;
   role: ChatMessage["role"];
@@ -21,11 +22,12 @@ function createMessage({
   parallelGroupId?: string | null;
   parallelIndex?: number | null;
   activeStreamId?: string | null;
+  text?: string;
 }): ChatMessage {
   return {
     id,
     role,
-    parts: [],
+    parts: text ? [{ type: "text", text }] : [],
     metadata: {
       createdAt: new Date(createdAt),
       parentMessageId,
@@ -44,14 +46,82 @@ function createThreadStore(initialMessages: ChatMessage[]) {
     withThreads<ChatMessage, BaseChatStoreState<ChatMessage>>(
       (set) =>
         ({
+          _messageIndex: { update: () => undefined },
+          _memoizedSelectors: new Map(),
+          _throttledMessages: initialMessages,
           messages: initialMessages,
           setMessages: (messages: ChatMessage[]) => set({ messages }),
-        }) as BaseChatStoreState<ChatMessage>
+        }) as unknown as BaseChatStoreState<ChatMessage>
     )
   );
 }
 
 describe("withThreads", () => {
+  it("preserves hidden branches when ThreadChat publishes an active-path snapshot", () => {
+    const userA = createMessage({
+      id: "user-a",
+      role: "user",
+      createdAt: "2024-01-01T00:00:00.000Z",
+    });
+    const assistantA = createMessage({
+      id: "assistant-a",
+      role: "assistant",
+      createdAt: "2024-01-01T00:00:01.000Z",
+      parentMessageId: userA.id,
+    });
+    const userB = createMessage({
+      id: "user-b",
+      role: "user",
+      createdAt: "2024-01-01T00:00:02.000Z",
+    });
+    const assistantB = createMessage({
+      id: "assistant-b",
+      role: "assistant",
+      createdAt: "2024-01-01T00:00:03.000Z",
+      parentMessageId: userB.id,
+    });
+    const store = createThreadStore([userB, assistantB]);
+
+    store.getState().setAllMessages([userA, assistantA, userB, assistantB]);
+    store.getState().setTreeSnapshot({
+      childrenByParentId: {
+        __root__: [userB.id],
+        [userB.id]: [assistantB.id],
+      },
+      cursorId: assistantB.id,
+      messagesById: {
+        [userB.id]: userB,
+        [assistantB.id]: assistantB,
+      },
+      parentById: {
+        [userB.id]: null,
+        [assistantB.id]: userB.id,
+      },
+      rootIds: [userB.id],
+      version: 1,
+    });
+
+    const siblingInfo = store.getState().getMessageSiblingInfo(userB.id);
+    assert.deepEqual(
+      siblingInfo?.siblings.map((message) => message.id),
+      [userA.id, userB.id]
+    );
+    const previousThread = store.getState().switchToSibling(userB.id, "prev");
+    const previousThreadIds = [userA.id, assistantA.id];
+    assert.deepEqual(
+      previousThread?.map((message) => message.id),
+      previousThreadIds
+    );
+    assert.deepEqual(
+      store.getState().messages.map((message) => message.id),
+      previousThreadIds
+    );
+    assert.deepEqual(
+      store.getState()._throttledMessages?.map((message) => message.id),
+      previousThreadIds
+    );
+  });
+
   it("preserves local-only optimistic branch nodes across server syncs", () => {
     const rootUser = createMessage({
       id: "user-root",
@@ -118,9 +188,9 @@ describe("withThreads", () => {
     assert.deepEqual(allMessageIds, [
       "user-root",
       "assistant-a",
-      "assistant-b",
       "user-nested",
       "assistant-nested-a",
+      "assistant-b",
       "assistant-nested-b",
     ]);
 
@@ -134,4 +204,182 @@ describe("withThreads", () => {
       "pending:assistant-nested-b"
     );
   });
+
+  it("preserves a pending stream marker when a run inserts its assistant shell", () => {
+    const rootUser = createMessage({
+      id: "user-root",
+      role: "user",
+      createdAt: "2024-01-01T00:00:00.000Z",
+    });
+    const assistant = createMessage({
+      id: "assistant-a",
+      role: "assistant",
+      createdAt: "2024-01-01T00:00:01.000Z",
+      parentMessageId: rootUser.id,
+      activeStreamId: "pending:assistant-a",
+    });
+    const assistantWithoutMetadata = {
+      id: assistant.id,
+      parts: [],
+      role: "assistant",
+    } as unknown as ChatMessage;
+
+    const store = createThreadStore([rootUser, assistant]);
+    store.getState().addMessageToTree(assistantWithoutMetadata);
+
+    const updatedAssistant = store.getState().allMessages.at(-1);
+    assert.equal(
+      updatedAssistant?.metadata.activeStreamId,
+      "pending:assistant-a"
+    );
+    assert.equal(
+      updatedAssistant?.metadata.selectedModel,
+      assistant.metadata.selectedModel
+    );
+  });
+
+  it("replaces a selected placeholder with completed server content when ready", () => {
+    const rootUser = createMessage({
+      id: "user-root",
+      role: "user",
+      createdAt: "2024-01-01T00:00:00.000Z",
+    });
+    const placeholder = createMessage({
+      id: "assistant-b",
+      role: "assistant",
+      createdAt: "2024-01-01T00:00:01.000Z",
+      parentMessageId: rootUser.id,
+      parallelGroupId: "group-root",
+      parallelIndex: 1,
+      activeStreamId: "pending:assistant-b",
+    });
+    const completed = createMessage({
+      id: placeholder.id,
+      role: "assistant",
+      createdAt: "2024-01-01T00:00:01.000Z",
+      parentMessageId: rootUser.id,
+      parallelGroupId: "group-root",
+      parallelIndex: 1,
+      text: "Completed secondary response",
+    });
+    const store = createThreadStore([rootUser, placeholder]);
+
+    store.getState().setAllMessages([rootUser, completed]);
+
+    const completedPart = store.getState().messages.at(-1)?.parts.at(0);
+    assert.equal(completedPart?.type, "text");
+    assert.equal(
+      completedPart?.type === "text" ? completedPart.text : null,
+      "Completed secondary response"
+    );
+    assert.equal(
+      store.getState().messages.at(-1)?.metadata.activeStreamId,
+      null
+    );
+    const throttledPart = store
+      .getState()
+      ._throttledMessages?.at(-1)
+      ?.parts.at(0);
+    assert.equal(
+      throttledPart?.type === "text" ? throttledPart.text : null,
+      "Completed secondary response"
+    );
+  });
+
+  it("fills metadata for ThreadChat assistant shells in tree snapshots", () => {
+    const rootUser = createMessage({
+      id: "user-root",
+      role: "user",
+      createdAt: "2024-01-01T00:00:00.000Z",
+    });
+    const assistantWithoutMetadata = {
+      id: "assistant-a",
+      parts: [],
+      role: "assistant",
+    } as unknown as ChatMessage;
+
+    const store = createThreadStore([rootUser]);
+
+    store.getState().setTreeSnapshot({
+      childrenByParentId: {
+        __root__: [rootUser.id],
+        [rootUser.id]: [assistantWithoutMetadata.id],
+      },
+      cursorId: assistantWithoutMetadata.id,
+      messagesById: {
+        [rootUser.id]: rootUser,
+        [assistantWithoutMetadata.id]: assistantWithoutMetadata,
+      },
+      parentById: {
+        [rootUser.id]: null,
+        [assistantWithoutMetadata.id]: rootUser.id,
+      },
+      rootIds: [rootUser.id],
+      version: 1,
+    });
+
+    const assistant = store.getState().messages.at(-1);
+    assert.equal(assistant?.metadata.parentMessageId, rootUser.id);
+    assert.equal(
+      assistant?.metadata.selectedModel,
+      rootUser.metadata.selectedModel
+    );
+  });
+
+  it("adds exactly one user sibling when editing after branch navigation", () => {
+    const userA = createMessage({
+      id: "user-a",
+      role: "user",
+      createdAt: "2024-01-01T00:00:00.000Z",
+    });
+    const assistantA = createMessage({
+      id: "assistant-a",
+      role: "assistant",
+      createdAt: "2024-01-01T00:00:01.000Z",
+      parentMessageId: userA.id,
+    });
+    const userB = createMessage({
+      id: "user-b",
+      role: "user",
+      createdAt: "2024-01-01T00:00:02.000Z",
+    });
+    const assistantB = createMessage({
+      id: "assistant-b",
+      role: "assistant",
+      createdAt: "2024-01-01T00:00:03.000Z",
+      parentMessageId: userB.id,
+    });
+    const userC = createMessage({
+      id: "user-c",
+      role: "user",
+      createdAt: "2024-01-01T00:00:04.000Z",
+    });
+    const assistantC = createMessage({
+      id: "assistant-c",
+      role: "assistant",
+      createdAt: "2024-01-01T00:00:05.000Z",
+      parentMessageId: userC.id,
+    });
+
+    const store = createThreadStore([userA, assistantA]);
+    store.getState().addMessageToTree(userB);
+    store.getState().addMessageToTree(assistantB);
+
+    assert.equal(
+      store.getState().getMessageSiblingInfo(userB.id)?.siblings.length,
+      2
+    );
+
+    store.getState().switchToSibling(userB.id, "prev");
+    store.getState().setMessages([]);
+    store.getState().addMessageToTree(userC);
+    store.getState().addMessageToTree(assistantC);
+
+    const rootSiblingIds = store
+      .getState()
+      .getMessageSiblingInfo(userC.id)
+      ?.siblings.map((message: ChatMessage) => message.id);
+
+    assert.deepEqual(rootSiblingIds, [userA.id, userB.id, userC.id]);
+  });
 });