feat: label session branches by first message

app/api/sessions/[id]/route.ts

@@ -12,108 +12,9 @@ import {
 } from "@/lib/session-reader";
 import { sessionPathKey } from "@/lib/session-path";
 import { getRpcSession } from "@/lib/rpc-manager";
+import { projectTreeForResponse } from "@/lib/project-tree";
 import { computeSessionTotalActiveMs } from "@/lib/session-timing";
 
-// BranchNavigator still traverses recursively, so keep the response tree shallow.
-const MAX_PROJECTED_TREE_DEPTH = 200;
-
-/**
- * Project the session tree into the shallow navigation tree sent to the client.
- * Keeps roots, branch points, and leaves while contracting single-child chains
- * without recursive traversal. Contracted entry IDs are attached to the next
- * visible node so the UI can still recognize an active leaf inside the chain.
- */
-function projectTreeForResponse<T extends { entry: { id: string }; children: T[]; compressedEntryIds?: string[] }>(
-  nodes: T[]
-): T[] {
-  const keep = new Set<T>();
-  const roots = new Set(nodes);
-  const seen = new Set<T>();
-  const stack = [...nodes];
-
-  while (stack.length > 0) {
-    const node = stack.pop()!;
-    if (seen.has(node)) continue;
-    seen.add(node);
-
-    if (
-      roots.has(node) ||
-      node.children.length !== 1
-    ) {
-      keep.add(node);
-    }
-
-    for (const child of node.children) {
-      stack.push(child);
-    }
-  }
-
-  const cloneNode = (node: T, compressedEntryIds?: string[]): T => ({
-    ...node,
-    children: [],
-    ...(compressedEntryIds?.length ? { compressedEntryIds } : {}),
-  });
-  const projectedRoots = nodes.map((node) => cloneNode(node));
-  const tasks = nodes.map((source, index) => ({
-    source,
-    projected: projectedRoots[index],
-    depth: 1,
-  }));
-
-  const appendFlattenedKeptDescendants = (source: T, projectedParent: T) => {
-    const pending = [{ node: source, compressedEntryIds: [] as string[] }];
-    const flattenedSeen = new Set<T>();
-
-    while (pending.length > 0) {
-      const { node, compressedEntryIds } = pending.pop()!;
-      if (flattenedSeen.has(node)) continue;
-      flattenedSeen.add(node);
-
-      if (keep.has(node)) {
-        projectedParent.children.push(cloneNode(node, compressedEntryIds));
-      }
-
-      for (let i = node.children.length - 1; i >= 0; i--) {
-        pending.push({
-          node: node.children[i],
-          compressedEntryIds: keep.has(node)
-            ? []
-            : [...compressedEntryIds, node.entry.id],
-        });
-      }
-    }
-  };
-
-  while (tasks.length > 0) {
-    const { source, projected, depth } = tasks.pop()!;
-
-    for (const sourceChild of source.children) {
-      let child = sourceChild;
-
-      if (depth >= MAX_PROJECTED_TREE_DEPTH) {
-        appendFlattenedKeptDescendants(child, projected);
-        continue;
-      }
-
-      const compressedEntryIds: string[] = [];
-      while (!keep.has(child) && child.children.length === 1) {
-        compressedEntryIds.push(child.entry.id);
-        child = child.children[0];
-      }
-
-      if (!keep.has(child)) {
-        continue;
-      }
-
-      const projectedChild = cloneNode(child, compressedEntryIds);
-      projected.children.push(projectedChild);
-      tasks.push({ source: child, projected: projectedChild, depth: depth + 1 });
-    }
-  }
-
-  return projectedRoots;
-}
-
 export async function GET(
   req: Request,
   { params }: { params: Promise<{ id: string }> }

components/BranchNavigator.test.mjs

@@ -0,0 +1,106 @@
+import assert from "node:assert/strict";
+import test from "node:test";
+import { createJiti } from "jiti";
+
+const jiti = createJiti(import.meta.url, {
+  jsx: { runtime: "automatic" },
+  tsconfigPaths: true,
+});
+const { compressChain, selectTopLevelBranches } = await jiti.import("./BranchNavigator.tsx");
+
+const msg = (id, role, text) => ({ type: "message", id, parentId: null, timestamp: "t", message: { role, content: text } });
+const info = (id) => ({ type: "session_info", id, parentId: null, timestamp: "t", name: "x" });
+const model = (id) => ({ type: "model_change", id, parentId: null, timestamp: "t", provider: "test", modelId: "test" });
+const node = (entry, children = []) => ({ entry, children });
+
+test("compressChain labels a chain by its first message entry", () => {
+  const chain = node(msg("u1", "user", "原问题"), [node(msg("a1", "assistant", "回答"))]);
+  const { labelEntry, node: rep } = compressChain(chain);
+  assert.equal(labelEntry.id, "u1");
+  assert.equal(rep.entry.id, "a1");
+});
+
+test("compressChain skips non-message entries such as session_info", () => {
+  const chain = node(info("s1"), [node(msg("u1", "user", "原始问题"), [node(msg("a1", "assistant", "答"))])]);
+  const { labelEntry, node: rep, skipped } = compressChain(chain);
+  assert.equal(labelEntry.id, "u1");
+  assert.equal(rep.entry.id, "a1");
+  assert.equal(skipped, 2);
+});
+
+test("compressChain labels a projected chain by its preview but selects its representative", () => {
+  const representative = {
+    entry: msg("a1", "assistant", "回答"),
+    children: [],
+    compressedEntryIds: ["u1"],
+    branchPreview: { role: "user", text: "原始问题" },
+  };
+  const chain = node(info("s1"), [representative]);
+  const { branchPreview, node: rep, skipped } = compressChain(chain);
+  assert.deepEqual(branchPreview, { role: "user", text: "原始问题" });
+  assert.equal(rep.entry.id, "a1");
+  assert.equal(skipped, 2);
+});
+
+test("compressChain falls back to the chain end when no message entry exists", () => {
+  const chain = node(info("s1"), [node(info("s2"))]);
+  const { labelEntry } = compressChain(chain);
+  assert.equal(labelEntry.id, "s2");
+});
+
+test("selectTopLevelBranches returns all roots for multi-root trees", () => {
+  const r1 = node(msg("u1", "user", "第一问"));
+  const r2 = node(msg("u1b", "user", "第一问改"));
+  assert.deepEqual(selectTopLevelBranches([r1, r2]).map((n) => n.entry.id), ["u1", "u1b"]);
+});
+
+test("selectTopLevelBranches returns children of the first branching node", () => {
+  const b1 = node(msg("u2", "user", "分支一"));
+  const b2 = node(msg("u2b", "user", "分支二"));
+  const root = node(msg("u1", "user", "第一问"), [node(msg("a1", "assistant", "答"), [b1, b2])]);
+  assert.deepEqual(selectTopLevelBranches([root]).map((n) => n.entry.id), ["u2", "u2b"]);
+});
+
+test("selectTopLevelBranches returns empty for a linear session", () => {
+  const root = node(msg("u1", "user", "第一问"), [node(msg("a1", "assistant", "答"))]);
+  assert.deepEqual(selectTopLevelBranches([root]), []);
+});
+
+test("selectTopLevelBranches works on preview-only server projections", () => {
+  const arm1 = {
+    entry: msg("a2", "assistant", "答一"),
+    children: [],
+    compressedEntryIds: ["s1", "u2"],
+    branchPreview: { role: "user", text: "分支一" },
+  };
+  const arm2 = {
+    entry: msg("a2b", "assistant", "答二"),
+    children: [],
+    compressedEntryIds: ["u2b"],
+    branchPreview: { role: "user", text: "分支二" },
+  };
+  const branchPoint = { entry: msg("a1", "assistant", "答"), children: [arm1, arm2] };
+  const root = { entry: msg("u1", "user", "第一问"), children: [branchPoint] };
+  const topLevel = selectTopLevelBranches([root]);
+  assert.deepEqual(topLevel.map((n) => n.entry.id), ["a2", "a2b"]);
+  assert.deepEqual(compressChain(topLevel[0]).branchPreview, { role: "user", text: "分支一" });
+  assert.equal(compressChain(topLevel[0]).node.entry.id, "a2");
+});
+
+test("multi-root metadata chains use their user previews and assistant representatives", () => {
+  const r1 = node(model("m1"), [{
+    entry: msg("a1", "assistant", "回答一"),
+    children: [],
+    compressedEntryIds: ["u1"],
+    branchPreview: { role: "user", text: "第一问" },
+  }]);
+  const r2 = node(info("s2"), [{
+    entry: msg("a2", "assistant", "回答二"),
+    children: [],
+    compressedEntryIds: ["u2"],
+    branchPreview: { role: "user", text: "第二问" },
+  }]);
+  const topLevel = selectTopLevelBranches([r1, r2]);
+  assert.deepEqual(topLevel.map((n) => compressChain(n).branchPreview.text), ["第一问", "第二问"]);
+  assert.deepEqual(topLevel.map((n) => compressChain(n).node.entry.id), ["a1", "a2"]);
+});

components/BranchNavigator.tsx

@@ -1,7 +1,7 @@
 "use client";
 
 import { useState, useCallback, useMemo, useRef, useEffect } from "react";
-import type { SessionEntry, SessionTreeNode } from "@/lib/types";
+import type { BranchPreview, SessionEntry, SessionTreeNode } from "@/lib/types";
 import { useI18n } from "@/hooks/useI18n";
 
 interface Props {
@@ -40,16 +40,41 @@ function buildActivePath(nodes: SessionTreeNode[], targetId: string | null): Set
   return new Set(search(nodes, []) ?? []);
 }
 
+function isMessageEntry(entry: SessionEntry): boolean {
+  return entry.type === "message" && "message" in entry;
+}
+
 // Compress a visible linear chain into the first branching/leaf node.
 // Server-side compressed IDs also count as skipped nodes.
-function compress(node: SessionTreeNode): { node: SessionTreeNode; skipped: number } {
+// branchPreview is the bounded preview of the first message on the source
+// chain. labelEntry keeps unprojected/test shapes working as a fallback.
+export function compressChain(node: SessionTreeNode): {
+  node: SessionTreeNode;
+  skipped: number;
+  branchPreview?: BranchPreview;
+  labelEntry: SessionEntry;
+} {
   let current = node;
+  let branchPreview = current.branchPreview;
+  let labelEntry: SessionEntry | null = isMessageEntry(current.entry) ? current.entry : null;
   let skipped = current.compressedEntryIds?.length ?? 0;
   while (current.children.length === 1) {
     current = current.children[0];
+    branchPreview ??= current.branchPreview;
+    if (!labelEntry && isMessageEntry(current.entry)) labelEntry = current.entry;
     skipped += 1 + (current.compressedEntryIds?.length ?? 0);
   }
-  return { node: current, skipped };
+  return { node: current, skipped, branchPreview, labelEntry: labelEntry ?? current.entry };
+}
+
+// Top-level rows of the panel: with multiple roots (a branch was started from
+// the very first message) the roots themselves are the branches; otherwise the
+// children of the first branching node.
+export function selectTopLevelBranches(tree: SessionTreeNode[]): SessionTreeNode[] {
+  if (tree.length > 1) return tree;
+  if (tree.length === 0) return [];
+  const first = compressChain(tree[0]).node;
+  return first.children.length > 1 ? first.children : [];
 }
 
 function getLabel(entry: SessionEntry): string {
@@ -74,6 +99,7 @@ function getLabel(entry: SessionEntry): string {
 
 // Does the tree have any branching at all?
 function hasBranch(nodes: SessionTreeNode[]): boolean {
+  if (nodes.length > 1) return true;
   for (const node of nodes) {
     if (node.children.length > 1) return true;
     if (hasBranch(node.children)) return true;
@@ -91,13 +117,15 @@ interface TreeNodeProps {
 }
 
 function TreeNodeView({ node, activePathIds, depth, isLast, parentLines, onSelect }: TreeNodeProps) {
-  const { node: rep, skipped } = compress(node);
+  const { node: rep, skipped, branchPreview, labelEntry } = compressChain(node);
   const isActive = activePathIds.has(rep.entry.id);
   const isOnPath = activePathIds.has(node.entry.id) || activePathIds.has(rep.entry.id);
-  const label = getLabel(rep.entry);
-  const role = rep.entry.type === "message" && "message" in rep.entry
-    ? (rep.entry.message as { role: string }).role
-    : null;
+  const label = branchPreview?.text ?? getLabel(labelEntry);
+  const role = branchPreview
+    ? branchPreview.role ?? null
+    : isMessageEntry(labelEntry)
+      ? (labelEntry as { message: { role: string } }).message.role
+      : null;
 
   return (
     <div>
@@ -253,10 +281,8 @@ export function BranchNavigator({ tree, activeLeafId, onLeafChange, inline, cont
       ? t("i18n.noBranches")
       : null;
 
-  // Find first meaningful node (skip pure linear prefix)
-  const compressed = tree.length > 0 ? compress(tree[0]) : null;
-  const firstNode = compressed?.node ?? null;
-  const hasContent = !noBranchReason && firstNode && firstNode.children.length > 1;
+  const topLevel = selectTopLevelBranches(tree);
+  const hasContent = !noBranchReason && topLevel.length > 0;
 
   const branchIcon = (
     <svg width="12" height="12" viewBox="0 0 24 24" fill="none" stroke="currentColor" strokeWidth="2" strokeLinecap="round" strokeLinejoin="round" style={{ color: hasContent ? "var(--accent)" : "var(--text-dim)", flexShrink: 0 }}>
@@ -315,15 +341,15 @@ export function BranchNavigator({ tree, activeLeafId, onLeafChange, inline, cont
             borderBottom: "1px solid var(--border)",
             zIndex: 500,
           }}>
-            {hasContent && firstNode ? (
+            {hasContent ? (
               <div style={{ padding: "4px 12px 8px 12px", maxHeight: 260, overflowY: "auto" }}>
-                {firstNode.children.map((child, idx) => (
+                {topLevel.map((child, idx) => (
                   <TreeNodeView
                     key={child.entry.id}
                     node={child}
                     activePathIds={activePathIds}
                     depth={0}
-                    isLast={idx === firstNode.children.length - 1}
+                    isLast={idx === topLevel.length - 1}
                     parentLines={[]}
                     onSelect={handleSelect}
                   />
@@ -376,15 +402,15 @@ export function BranchNavigator({ tree, activeLeafId, onLeafChange, inline, cont
           boxShadow: "0 4px 12px rgba(0,0,0,0.1)",
           zIndex: 100,
         }}>
-          {hasContent && firstNode ? (
+          {hasContent ? (
             <div style={{ padding: "4px 12px 8px 12px", maxHeight: 260, overflowY: "auto" }}>
-              {firstNode.children.map((child, idx) => (
+              {topLevel.map((child, idx) => (
                 <TreeNodeView
                   key={child.entry.id}
                   node={child}
                   activePathIds={activePathIds}
                   depth={0}
-                  isLast={idx === firstNode.children.length - 1}
+                  isLast={idx === topLevel.length - 1}
                   parentLines={[]}
                   onSelect={handleSelect}
                 />