@chat-js/cli 0.3.0 → 0.4.0

package/templates/chat-app/lib/db/migrations/meta/_journal.json

@@ -309,6 +309,13 @@
       "when": 1769940346934,
       "tag": "0043_lyrical_pride",
       "breakpoints": true
+    },
+    {
+      "idx": 44,
+      "version": "7",
+      "when": 1773499312081,
+      "tag": "0044_gray_red_shift",
+      "breakpoints": true
     }
   ]
-}
+}

package/templates/chat-app/lib/db/queries.ts

@@ -17,6 +17,7 @@ import type {
   ToolName,
   ToolOutput,
 } from "@/lib/ai/types";
+import { isSelectedModelValue } from "@/lib/ai/types";
 import { createModuleLogger } from "@/lib/logger";
 import { chatMessageToDbMessage } from "@/lib/message-conversion";
 
@@ -79,6 +80,35 @@ export async function saveChat({
   }
 }
 
+export async function saveChatIfNotExists({
+  id,
+  userId,
+  title,
+  projectId,
+}: {
+  id: string;
+  userId: string;
+  title: string;
+  projectId?: string;
+}) {
+  try {
+    return await db
+      .insert(chat)
+      .values({
+        id,
+        createdAt: new Date(),
+        updatedAt: new Date(),
+        userId,
+        title,
+        projectId: projectId ?? null,
+      })
+      .onConflictDoNothing();
+  } catch (error) {
+    console.error("Failed to save chat in database");
+    throw error;
+  }
+}
+
 export async function deleteChatById({ id }: { id: string }) {
   try {
     // Get all messages for this chat to clean up their attachments
@@ -341,6 +371,43 @@ export async function saveMessage({
   }
 }
 
+export async function saveMessageIfNotExists({
+  id,
+  chatId,
+  message: chatMessage,
+}: {
+  id: string;
+  chatId: string;
+  message: ChatMessage;
+}) {
+  try {
+    return await db.transaction(async (tx) => {
+      const dbMessage = chatMessageToDbMessage(chatMessage, chatId);
+      dbMessage.id = id;
+
+      const insertedMessages = await tx
+        .insert(message)
+        .values(dbMessage)
+        .onConflictDoNothing()
+        .returning({ id: message.id });
+
+      if (insertedMessages.length === 0) {
+        return;
+      }
+
+      const mappedDBParts = mapUIMessagePartsToDBParts(chatMessage.parts, id);
+      if (mappedDBParts.length > 0) {
+        await tx.insert(part).values(mappedDBParts);
+      }
+
+      await updateChatUpdatedAt({ chatId });
+    });
+  } catch (error) {
+    logger.error({ error, chatId, id }, "saveMessageIfNotExists failed");
+    throw error;
+  }
+}
+
 export async function saveChatMessages({
   messages,
 }: {
@@ -413,6 +480,11 @@ export async function updateMessage({
           attachments: dbMessage.attachments,
           createdAt: dbMessage.createdAt,
           parentMessageId: dbMessage.parentMessageId,
+          selectedModel: dbMessage.selectedModel,
+          selectedTool: dbMessage.selectedTool,
+          parallelGroupId: dbMessage.parallelGroupId,
+          parallelIndex: dbMessage.parallelIndex,
+          isPrimaryParallel: dbMessage.isPrimaryParallel,
           lastContext: dbMessage.lastContext,
           activeStreamId: dbMessage.activeStreamId,
         })
@@ -492,8 +564,12 @@ export async function getAllMessagesByChatId({
           createdAt: msg.createdAt,
           activeStreamId: msg.activeStreamId,
           parentMessageId: msg.parentMessageId,
-          selectedModel: (msg.selectedModel ||
-            "") as ChatMessage["metadata"]["selectedModel"],
+          parallelGroupId: msg.parallelGroupId,
+          parallelIndex: msg.parallelIndex,
+          isPrimaryParallel: msg.isPrimaryParallel,
+          selectedModel: isSelectedModelValue(msg.selectedModel)
+            ? msg.selectedModel
+            : ("" as ChatMessage["metadata"]["selectedModel"]),
           selectedTool: (msg.selectedTool ||
             undefined) as ChatMessage["metadata"]["selectedTool"],
           usage: msg.lastContext as ChatMessage["metadata"]["usage"],
@@ -797,8 +873,12 @@ export async function getChatMessageWithPartsById({
           createdAt: dbMessage.createdAt,
           activeStreamId: dbMessage.activeStreamId,
           parentMessageId: dbMessage.parentMessageId,
-          selectedModel: (dbMessage.selectedModel ||
-            "") as ChatMessage["metadata"]["selectedModel"],
+          parallelGroupId: dbMessage.parallelGroupId,
+          parallelIndex: dbMessage.parallelIndex,
+          isPrimaryParallel: dbMessage.isPrimaryParallel,
+          selectedModel: isSelectedModelValue(dbMessage.selectedModel)
+            ? dbMessage.selectedModel
+            : ("" as ChatMessage["metadata"]["selectedModel"]),
           selectedTool: (dbMessage.selectedTool ||
             undefined) as ChatMessage["metadata"]["selectedTool"],
           usage: dbMessage.lastContext as ChatMessage["metadata"]["usage"],

package/templates/chat-app/lib/db/schema.ts

@@ -113,8 +113,11 @@ export const message = pgTable("Message", {
   attachments: json("attachments").notNull(),
   createdAt: timestamp("createdAt").notNull(),
   annotations: json("annotations"),
-  selectedModel: varchar("selectedModel", { length: 256 }).default(""),
+  selectedModel: json("selectedModel"),
   selectedTool: varchar("selectedTool", { length: 256 }).default(""),
+  parallelGroupId: uuid("parallelGroupId"),
+  parallelIndex: integer("parallelIndex"),
+  isPrimaryParallel: boolean("isPrimaryParallel"),
   lastContext: json("lastContext"),
   activeStreamId: varchar("activeStreamId", { length: 64 }),
   /** Timestamp when this message's stream was canceled by the user. Null means not canceled. */

package/templates/chat-app/lib/message-conversion.ts

@@ -1,7 +1,11 @@
 import type { ModelId } from "@/lib/ai/app-models";
 import type { Chat, DBMessage } from "@/lib/db/schema";
 import type { UIChat } from "@/lib/types/ui-chat";
-import type { ChatMessage, UiToolName } from "./ai/types";
+import {
+  isSelectedModelValue,
+  type ChatMessage,
+  type UiToolName,
+} from "./ai/types";
 
 // Helper functions for type conversion
 export function dbChatToUIChat(chat: Chat): UIChat {
@@ -29,7 +33,12 @@ function _dbMessageToChatMessage(message: DBMessage): ChatMessage {
       createdAt: message.createdAt,
       activeStreamId: message.activeStreamId,
       parentMessageId: message.parentMessageId,
-      selectedModel: (message.selectedModel as ModelId) || ("" as ModelId),
+      parallelGroupId: message.parallelGroupId,
+      parallelIndex: message.parallelIndex,
+      isPrimaryParallel: message.isPrimaryParallel,
+      selectedModel: isSelectedModelValue(message.selectedModel)
+        ? message.selectedModel
+        : ("" as ModelId),
       selectedTool: (message.selectedTool as UiToolName | null) || undefined,
       usage: message.lastContext as ChatMessage["metadata"]["usage"],
     },
@@ -66,6 +75,9 @@ export function chatMessageToDbMessage(
     parentMessageId,
     selectedModel,
     selectedTool: message.metadata?.selectedTool || null,
+    parallelGroupId: message.metadata?.parallelGroupId || null,
+    parallelIndex: message.metadata?.parallelIndex ?? null,
+    isPrimaryParallel: message.metadata?.isPrimaryParallel ?? null,
     activeStreamId: message.metadata?.activeStreamId || null,
     canceledAt: null,
   };

package/templates/chat-app/lib/stores/hooks-threads.ts

@@ -7,7 +7,7 @@ import {
   type CustomChatStoreState,
   useCustomChatStoreApi,
 } from "./custom-store-provider";
-import type { MessageSiblingInfo } from "./with-threads";
+import type { MessageSiblingInfo, ParallelGroupInfo } from "./with-threads";
 
 function useThreadStore<T>(
   selector: (store: CustomChatStoreState<ChatMessage>) => T,
@@ -95,3 +95,39 @@ export const useSwitchToSibling = () => {
     [store]
   );
 };
+
+export function useParallelGroupInfo(
+  messageId: string
+): ParallelGroupInfo<ChatMessage> | null {
+  return useThreadStore(
+    (state) => state.getParallelGroupInfo(messageId),
+    (a, b) => {
+      if (a === null && b === null) {
+        return true;
+      }
+
+      if (a === null || b === null) {
+        return false;
+      }
+
+      return (
+        a.parallelGroupId === b.parallelGroupId &&
+        a.selectedMessageId === b.selectedMessageId &&
+        a.messages.length === b.messages.length &&
+        a.messages.every(
+          (msg, i) =>
+            msg.id === b.messages[i]?.id &&
+            msg.metadata?.activeStreamId === b.messages[i]?.metadata?.activeStreamId
+        )
+      );
+    }
+  );
+}
+
+export const useSwitchToMessage = () => {
+  const store = useCustomChatStoreApi<ChatMessage>();
+  return useCallback(
+    (messageId: string) => store.getState().switchToMessage(messageId),
+    [store]
+  );
+};

package/templates/chat-app/lib/stores/with-threads.test.ts

@@ -0,0 +1,137 @@
+import assert from "node:assert/strict";
+import type { StoreState as BaseChatStoreState } from "@ai-sdk-tools/store";
+import { describe, it } from "vitest";
+import { createStore } from "zustand/vanilla";
+import type { ChatMessage } from "../ai/types";
+import { withThreads, type ThreadAugmentedState } from "./with-threads";
+
+function createMessage({
+  id,
+  role,
+  createdAt,
+  parentMessageId = null,
+  parallelGroupId = null,
+  parallelIndex = null,
+  activeStreamId = null,
+}: {
+  id: string;
+  role: ChatMessage["role"];
+  createdAt: string;
+  parentMessageId?: string | null;
+  parallelGroupId?: string | null;
+  parallelIndex?: number | null;
+  activeStreamId?: string | null;
+}): ChatMessage {
+  return {
+    id,
+    role,
+    parts: [],
+    metadata: {
+      createdAt: new Date(createdAt),
+      parentMessageId,
+      parallelGroupId,
+      parallelIndex,
+      isPrimaryParallel: parallelIndex === null ? null : parallelIndex === 0,
+      selectedModel: "openai/gpt-4o-mini",
+      activeStreamId,
+      selectedTool: undefined,
+    },
+  };
+}
+
+function createThreadStore(initialMessages: ChatMessage[]) {
+  return createStore<ThreadAugmentedState<ChatMessage>>()(
+    withThreads<ChatMessage, BaseChatStoreState<ChatMessage>>(
+      (set) =>
+        ({
+          messages: initialMessages,
+          setMessages: (messages: ChatMessage[]) => set({ messages }),
+        }) as BaseChatStoreState<ChatMessage>
+    )
+  );
+}
+
+describe("withThreads", () => {
+  it("preserves local-only optimistic branch nodes across server syncs", () => {
+    const rootUser = createMessage({
+      id: "user-root",
+      role: "user",
+      createdAt: "2024-01-01T00:00:00.000Z",
+    });
+    const branchA = createMessage({
+      id: "assistant-a",
+      role: "assistant",
+      createdAt: "2024-01-01T00:00:01.000Z",
+      parentMessageId: rootUser.id,
+      parallelGroupId: "group-root",
+      parallelIndex: 0,
+    });
+    const branchB = createMessage({
+      id: "assistant-b",
+      role: "assistant",
+      createdAt: "2024-01-01T00:00:02.000Z",
+      parentMessageId: rootUser.id,
+      parallelGroupId: "group-root",
+      parallelIndex: 1,
+    });
+    const nestedUser = createMessage({
+      id: "user-nested",
+      role: "user",
+      createdAt: "2024-01-01T00:00:03.000Z",
+      parentMessageId: branchA.id,
+    });
+    const nestedBranchA = createMessage({
+      id: "assistant-nested-a",
+      role: "assistant",
+      createdAt: "2024-01-01T00:00:04.000Z",
+      parentMessageId: nestedUser.id,
+      parallelGroupId: "group-nested",
+      parallelIndex: 0,
+      activeStreamId: "pending:assistant-nested-a",
+    });
+    const nestedBranchB = createMessage({
+      id: "assistant-nested-b",
+      role: "assistant",
+      createdAt: "2024-01-01T00:00:05.000Z",
+      parentMessageId: nestedUser.id,
+      parallelGroupId: "group-nested",
+      parallelIndex: 1,
+      activeStreamId: "pending:assistant-nested-b",
+    });
+
+    const store = createThreadStore([
+      rootUser,
+      branchA,
+      nestedUser,
+      nestedBranchA,
+    ]);
+
+    store.getState().addMessageToTree(branchB);
+    store.getState().addMessageToTree(nestedBranchB);
+
+    store.getState().setMessagesWithEpoch([rootUser, branchB]);
+    store.getState().setAllMessages([rootUser, branchA, branchB]);
+
+    const allMessageIds = store
+      .getState()
+      .allMessages.map((message: ChatMessage) => message.id);
+    assert.deepEqual(allMessageIds, [
+      "user-root",
+      "assistant-a",
+      "assistant-b",
+      "user-nested",
+      "assistant-nested-a",
+      "assistant-nested-b",
+    ]);
+
+    const restoredThread = store.getState().switchToMessage(branchA.id);
+    assert.deepEqual(
+      restoredThread?.map((message: ChatMessage) => message.id),
+      ["user-root", "assistant-a", "user-nested", "assistant-nested-b"]
+    );
+    assert.equal(
+      restoredThread?.at(-1)?.metadata.activeStreamId,
+      "pending:assistant-nested-b"
+    );
+  });
+});

package/templates/chat-app/lib/stores/with-threads.ts

@@ -19,6 +19,12 @@ export interface MessageSiblingInfo<UM> {
   siblings: UM[];
 }
 
+export interface ParallelGroupInfo<UM> {
+  messages: UM[];
+  parallelGroupId: string;
+  selectedMessageId: string | null;
+}
+
 export type ThreadAugmentedState<UM extends UIMessage> =
   BaseChatStoreState<UM> & {
     threadEpoch: number;
@@ -41,6 +47,7 @@ export type ThreadAugmentedState<UM extends UIMessage> =
     addMessageToTree: (message: UM) => void;
     /** Look up sibling info for a message. */
     getMessageSiblingInfo: (messageId: string) => MessageSiblingInfo<UM> | null;
+    getParallelGroupInfo: (messageId: string) => ParallelGroupInfo<UM> | null;
     /**
      * Switch to a sibling thread. Returns the new thread array,
      * or null if no switch was possible.
@@ -49,6 +56,7 @@ export type ThreadAugmentedState<UM extends UIMessage> =
       messageId: string,
       direction: "prev" | "next"
     ) => UM[] | null;
+    switchToMessage: (messageId: string) => UM[] | null;
   };
 
 export const withThreads =
@@ -64,6 +72,33 @@ export const withThreads =
     const rebuildMap = (msgs: UI_MESSAGE[]) =>
       buildChildrenMap(msgs as (UI_MESSAGE & MessageNode)[]);
 
+    const mergeTreeMessages = (
+      serverMessages: UI_MESSAGE[],
+      existingTreeMessages: UI_MESSAGE[],
+      currentVisibleMessages: UI_MESSAGE[]
+    ): UI_MESSAGE[] => {
+      const merged = new Map<string, UI_MESSAGE>();
+
+      for (const message of serverMessages) {
+        merged.set(message.id, message);
+      }
+
+      // Preserve every local-only tree node until the server returns a message with
+      // the same id. Restricting this to pending assistant shells orphaned optimistic
+      // user messages when switching away from an in-flight branch mid-stream.
+      for (const message of existingTreeMessages) {
+        if (!merged.has(message.id)) {
+          merged.set(message.id, message);
+        }
+      }
+
+      for (const message of currentVisibleMessages) {
+        merged.set(message.id, message);
+      }
+
+      return Array.from(merged.values());
+    };
+
     return {
       ...base,
       threadEpoch: 0,
@@ -96,10 +131,45 @@ export const withThreads =
       },
 
       setAllMessages: (messages: UI_MESSAGE[]) => {
-        set((state) => ({
-          ...state,
-          allMessages: messages,
-          childrenMap: rebuildMap(messages),
+        const state = get();
+        const currentVisibleMessages = state.messages;
+        const existingTreeMessages = state.allMessages;
+        const mergedMessages = mergeTreeMessages(
+          messages,
+          existingTreeMessages,
+          currentVisibleMessages
+        );
+
+        // While the SDK is actively streaming, updating the visible thread with
+        // server data would mix the SDK's client-generated message ID with the
+        // server's assistantMessageId. The mismatch causes the SDK to push a
+        // second assistant message on the next chunk, bumping the epoch and
+        // remounting ChatSync mid-stream. Only update the tree index here and
+        // let the normal post-stream invalidation apply the full visible update.
+        if (state.status === "streaming" || state.status === "submitted") {
+          set((prev) => ({
+            ...prev,
+            allMessages: mergedMessages,
+            childrenMap: rebuildMap(mergedMessages),
+          }));
+          return;
+        }
+
+        const currentLeafId = currentVisibleMessages.at(-1)?.id;
+        const nextVisibleThread = currentLeafId
+          ? (buildThreadFromLeaf(
+              mergedMessages as (UI_MESSAGE & MessageNode)[],
+              currentLeafId
+            ) as UI_MESSAGE[])
+          : currentVisibleMessages;
+
+        originalSetMessages(nextVisibleThread);
+        set((prev) => ({
+          ...prev,
+          messages: nextVisibleThread,
+          threadInitialMessages: nextVisibleThread,
+          allMessages: mergedMessages,
+          childrenMap: rebuildMap(mergedMessages),
         }));
       },
 
@@ -135,6 +205,66 @@ export const withThreads =
         return { siblings, siblingIndex };
       },
 
+      getParallelGroupInfo: (
+        messageId: string
+      ): ParallelGroupInfo<UI_MESSAGE> | null => {
+        const state = get();
+        const message = state.allMessages.find((item) => item.id === messageId);
+        if (!message) {
+          return null;
+        }
+
+        const metadata = (message as UI_MESSAGE & MessageNode).metadata;
+        const parallelGroupId = metadata?.parallelGroupId || null;
+        const parentId =
+          message.role === "user"
+            ? message.id
+            : metadata?.parentMessageId || null;
+
+        if (!(parentId && parallelGroupId)) {
+          return null;
+        }
+
+        const groupMessages = (
+          (state.childrenMap.get(parentId) ?? []) as UI_MESSAGE[]
+        )
+          .filter(
+            (candidate) =>
+              (candidate as UI_MESSAGE & MessageNode).metadata
+                ?.parallelGroupId === parallelGroupId
+          )
+          .sort((a, b) => {
+            const aIndex =
+              (a as UI_MESSAGE & MessageNode).metadata?.parallelIndex ??
+              Number.MAX_SAFE_INTEGER;
+            const bIndex =
+              (b as UI_MESSAGE & MessageNode).metadata?.parallelIndex ??
+              Number.MAX_SAFE_INTEGER;
+
+            if (aIndex !== bIndex) {
+              return aIndex - bIndex;
+            }
+
+            return 0;
+          });
+
+        if (groupMessages.length <= 1) {
+          return null;
+        }
+
+        const visibleMessageIds = new Set(state.messages.map((m) => m.id));
+        const selectedMessageId =
+          groupMessages.find((candidate) =>
+            visibleMessageIds.has(candidate.id)
+          )?.id ?? null;
+
+        return {
+          messages: groupMessages,
+          parallelGroupId,
+          selectedMessageId,
+        };
+      },
+
       switchToSibling: (
         messageId: string,
         direction: "prev" | "next"
@@ -170,6 +300,29 @@ export const withThreads =
         return newThread;
       },
 
+      switchToMessage: (messageId: string): UI_MESSAGE[] | null => {
+        const state = get();
+        const { allMessages, childrenMap } = state;
+        const message = allMessages.find(
+          (candidate) => candidate.id === messageId
+        );
+        if (!message) {
+          return null;
+        }
+
+        const leaf = findLeafDfsToRightFromMessageId(
+          childrenMap as Map<string | null, (UI_MESSAGE & MessageNode)[]>,
+          messageId
+        );
+        const newThread = buildThreadFromLeaf(
+          allMessages as (UI_MESSAGE & MessageNode)[],
+          leaf ? leaf.id : messageId
+        ) as UI_MESSAGE[];
+
+        state.setMessagesWithEpoch(newThread);
+        return newThread;
+      },
+
       // Override setMessages to auto-bump epoch when thread changes
       setMessages: (messages: UI_MESSAGE[]) => {
         const currentMessages = get().messages;

package/templates/chat-app/lib/thread-utils.ts

@@ -3,6 +3,9 @@ export interface MessageNode {
   id: string;
   metadata?: {
     parentMessageId: string | null;
+    parallelGroupId?: string | null;
+    parallelIndex?: number | null;
+    activeStreamId?: string | null;
     createdAt: Date;
   };
 }
@@ -102,8 +105,26 @@ export function buildChildrenMap<T extends MessageNode>(
   }
   for (const siblings of map.values()) {
     siblings.sort(
-      (a, b) =>
-        toTimestamp(a.metadata?.createdAt) - toTimestamp(b.metadata?.createdAt)
+      (a, b) => {
+        const aParallelIndex = a.metadata?.parallelIndex;
+        const bParallelIndex = b.metadata?.parallelIndex;
+        const sameParallelGroup =
+          a.metadata?.parallelGroupId &&
+          a.metadata?.parallelGroupId === b.metadata?.parallelGroupId;
+
+        if (
+          sameParallelGroup &&
+          typeof aParallelIndex === "number" &&
+          typeof bParallelIndex === "number" &&
+          aParallelIndex !== bParallelIndex
+        ) {
+          return aParallelIndex - bParallelIndex;
+        }
+
+        return (
+          toTimestamp(a.metadata?.createdAt) - toTimestamp(b.metadata?.createdAt)
+        );
+      }
     );
   }
   return map;