Branching conversations

// Each message is a node in the tree
type ChatNode = {
  id: string;
  chatId: string;
  parentId: string | null; // null for root
  role: "user" | "assistant";
  message: UIMessage; // the full AI SDK message
  createdAt: Date;
};

root
├── user: "Hello"
│   └── assistant: "Hi there!"
│       ├── user: "What's the weather?" ← branch A
│       │   └── assistant: "It's sunny!"
│       └── user: "Tell me a joke" ← branch B (active)
│           └── assistant: "Why did the..."

// Resolve the active path: walk from leaf to root, reverse
async function getActiveBranch(chatId: string): Promise<UIMessage[]> {
  const nodes = await db.chatNode.findMany({ where: { chatId } });
  const byId = new Map(nodes.map((n) => [n.id, n]));

  // Find active leaf (most recently created leaf node)
  const childIds = new Set(nodes.map((n) => n.parentId).filter(Boolean));
  const leaves = nodes.filter((n) => !childIds.has(n.id));
  const activeLeaf = leaves.sort((a, b) => b.createdAt - a.createdAt)[0];
  if (!activeLeaf) return [];

  // Walk to root
  const path: UIMessage[] = [];
  let current: ChatNode | undefined = activeLeaf;
  while (current) {
    path.unshift(current.message);
    current = current.parentId ? byId.get(current.parentId) : undefined;
  }
  return path;
}

// Append a message as a child of the current leaf
async function appendMessage(chatId: string, message: UIMessage): Promise<void> {
  const branch = await getActiveBranch(chatId);
  const parentId = branch.length > 0 ? branch[branch.length - 1]!.id : null;

  await db.chatNode.create({
    data: { id: message.id, chatId, parentId, role: message.role, message, createdAt: new Date() },
  });
}

import { chat } from "@trigger.dev/sdk/ai";
import { streamText, stepCountIs } from "ai";
import { anthropic } from "@ai-sdk/anthropic";
import { z } from "zod";

export const myChat = chat.agent({
  id: "branching-chat",

  // Load the active branch from the DB on every turn.
  // The frontend's message array is ignored — the tree is the source of truth.
  hydrateMessages: async ({ chatId, trigger, incomingMessages }) => {
    if (trigger === "submit-message" && incomingMessages.length > 0) {
      await appendMessage(chatId, incomingMessages[incomingMessages.length - 1]!);
    }
    return getActiveBranch(chatId);
  },

  actionSchema: z.discriminatedUnion("type", [
    // Edit a previous user message — creates a sibling node in the tree
    z.object({ type: z.literal("edit"), messageId: z.string(), text: z.string() }),
    // Switch to a different branch by selecting a leaf node
    z.object({ type: z.literal("switch-branch"), leafId: z.string() }),
    // Undo the last user + assistant exchange
    z.object({ type: z.literal("undo") }),
  ]),

  onAction: async ({ action, chatId }) => {
    switch (action.type) {
      case "edit": {
        // Find the original message's parent, create a sibling with new content
        const original = await db.chatNode.findUnique({ where: { id: action.messageId } });
        if (!original) break;

        const newId = generateId();
        await db.chatNode.create({
          data: {
            id: newId,
            chatId,
            parentId: original.parentId, // same parent = sibling
            role: "user",
            message: { id: newId, role: "user", parts: [{ type: "text", text: action.text }] },
            createdAt: new Date(),
          },
        });
        // Active branch now resolves through the new sibling (most recent leaf)
        break;
      }

      case "switch-branch": {
        // Mark this leaf as the most recently accessed so getActiveBranch picks it
        await db.chatNode.update({
          where: { id: action.leafId },
          data: { createdAt: new Date() },
        });
        break;
      }

      case "undo": {
        // Remove the last two nodes (user + assistant) from the active branch
        const branch = await getActiveBranch(chatId);
        if (branch.length >= 2) {
          const lastTwo = branch.slice(-2);
          await db.chatNode.deleteMany({
            where: { id: { in: lastTwo.map((m) => m.id) } },
          });
        }
        break;
      }
    }

    // Reload the (now modified) active branch into the accumulator
    const updated = await getActiveBranch(chatId);
    chat.history.set(updated);
  },

  onTurnComplete: async ({ chatId, responseMessage }) => {
    // Persist the assistant's response as a new node
    if (responseMessage) {
      await appendMessage(chatId, responseMessage);
    }
  },

  run: async ({ messages, signal }) => {
    return streamText({
      model: anthropic("claude-sonnet-4-5"),
      messages,
      abortSignal: signal,
      stopWhen: stepCountIs(15),
    });
  },
});

function MessageActions({ message, chatId }: { message: UIMessage; chatId: string }) {
  const transport = useTransport();
  const [editing, setEditing] = useState(false);
  const [editText, setEditText] = useState("");

  if (message.role !== "user") return null;

  return (
    <div>
      {editing ? (
        <form onSubmit={() => {
          transport.sendAction(chatId, { type: "edit", messageId: message.id, text: editText });
          setEditing(false);
        }}>
          <input value={editText} onChange={(e) => setEditText(e.target.value)} />
          <button type="submit">Save</button>
        </form>
      ) : (
        <button onClick={() => { setEditText(getMessageText(message)); setEditing(true); }}>
          Edit
        </button>
      )}
    </div>
  );
}

function BranchSwitcher({ message, chatId, siblings }: {
  message: UIMessage;
  chatId: string;
  siblings: { id: string; createdAt: string }[];
}) {
  const transport = useTransport();
  if (siblings.length <= 1) return null;

  const currentIndex = siblings.findIndex((s) => s.id === message.id);

  return (
    <div>
      <button
        disabled={currentIndex === 0}
        onClick={() => {
          // Find the leaf of the previous sibling's subtree
          transport.sendAction(chatId, {
            type: "switch-branch",
            leafId: siblings[currentIndex - 1]!.id,
          });
        }}
      >
        &lt;
      </button>
      <span>{currentIndex + 1}/{siblings.length}</span>
      <button
        disabled={currentIndex === siblings.length - 1}
        onClick={() => {
          transport.sendAction(chatId, {
            type: "switch-branch",
            leafId: siblings[currentIndex + 1]!.id,
          });
        }}
      >
        &gt;
      </button>
    </div>
  );
}