From 8a6141cbaa179b45304e42cad39b7edaf2a47a33 Mon Sep 17 00:00:00 2001 From: 111-test-111 Date: Sun, 9 Aug 2026 23:20:48 +0800 Subject: [PATCH 1/2] feat: label session branches by first message Extract the session-tree projection into lib/project-tree.ts and keep the first message entry of each diverging arm, so BranchNavigator can label each branch with the message it diverged with. BranchNavigator renders those first-message labels and supports multi-root trees. --- app/api/sessions/[id]/route.ts | 101 +------------------------ components/BranchNavigator.test.mjs | 62 ++++++++++++++++ components/BranchNavigator.tsx | 50 +++++++++---- lib/project-tree.test.mjs | 37 ++++++++++ lib/project-tree.ts | 111 ++++++++++++++++++++++++++++ 5 files changed, 245 insertions(+), 116 deletions(-) create mode 100644 components/BranchNavigator.test.mjs create mode 100644 lib/project-tree.test.mjs create mode 100644 lib/project-tree.ts diff --git a/app/api/sessions/[id]/route.ts b/app/api/sessions/[id]/route.ts index 899d01371..cd3a7fee0 100644 --- a/app/api/sessions/[id]/route.ts +++ b/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( - nodes: T[] -): T[] { - const keep = new Set(); - const roots = new Set(nodes); - const seen = new Set(); - 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(); - - 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 }> } diff --git a/components/BranchNavigator.test.mjs b/components/BranchNavigator.test.mjs new file mode 100644 index 000000000..481c201c7 --- /dev/null +++ b/components/BranchNavigator.test.mjs @@ -0,0 +1,62 @@ +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 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 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 server-projected tree shapes", () => { + const arm1 = { entry: msg("u2", "user", "分支一"), children: [{ entry: msg("a2", "assistant", "答一"), children: [], compressedEntryIds: ["s1"] }] }; + const arm2 = { entry: msg("u2b", "user", "分支二"), children: [{ entry: msg("a2b", "assistant", "答二"), children: [] }] }; + 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), ["u2", "u2b"]); + assert.equal(compressChain(topLevel[0]).labelEntry.id, "u2"); +}); diff --git a/components/BranchNavigator.tsx b/components/BranchNavigator.tsx index 570020602..d17b7ce23 100644 --- a/components/BranchNavigator.tsx +++ b/components/BranchNavigator.tsx @@ -40,16 +40,35 @@ 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 } { +// labelEntry is the first message entry on the chain — the message the branch +// diverged with — so labels read as the question, not the tail of the chain. +// Falls back to the chain end when the chain has no message entry at all. +export function compressChain(node: SessionTreeNode): { node: SessionTreeNode; skipped: number; labelEntry: SessionEntry } { let current = node; + 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]; + if (!labelEntry && isMessageEntry(current.entry)) labelEntry = current.entry; skipped += 1 + (current.compressedEntryIds?.length ?? 0); } - return { node: current, skipped }; + return { node: current, skipped, 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 +93,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,12 +111,12 @@ interface TreeNodeProps { } function TreeNodeView({ node, activePathIds, depth, isLast, parentLines, onSelect }: TreeNodeProps) { - const { node: rep, skipped } = compress(node); + const { node: rep, skipped, 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 + const label = getLabel(labelEntry); + const role = isMessageEntry(labelEntry) + ? (labelEntry as { message: { role: string } }).message.role : null; return ( @@ -253,10 +273,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 = ( @@ -315,15 +333,15 @@ export function BranchNavigator({ tree, activeLeafId, onLeafChange, inline, cont borderBottom: "1px solid var(--border)", zIndex: 500, }}> - {hasContent && firstNode ? ( + {hasContent ? (
- {firstNode.children.map((child, idx) => ( + {topLevel.map((child, idx) => ( @@ -376,15 +394,15 @@ export function BranchNavigator({ tree, activeLeafId, onLeafChange, inline, cont boxShadow: "0 4px 12px rgba(0,0,0,0.1)", zIndex: 100, }}> - {hasContent && firstNode ? ( + {hasContent ? (
- {firstNode.children.map((child, idx) => ( + {topLevel.map((child, idx) => ( diff --git a/lib/project-tree.test.mjs b/lib/project-tree.test.mjs new file mode 100644 index 000000000..08fed857a --- /dev/null +++ b/lib/project-tree.test.mjs @@ -0,0 +1,37 @@ +import assert from "node:assert/strict"; +import test from "node:test"; + +const { projectTreeForResponse } = await import("./project-tree.ts"); + +const msg = (id, text) => ({ type: "message", id, parentId: null, timestamp: "t", message: { role: "user", content: text } }); +const info = (id) => ({ type: "session_info", id, parentId: null, timestamp: "t", name: "x" }); +const node = (entry, children = []) => ({ entry, children }); + +test("keeps the first message entry of each diverging arm", () => { + // A1 branches: arm 1 starts with session_info then U2; arm 2 starts with U2b + const arm1Leaf = node(msg("a2", "答一")); + const arm1 = node(info("s1"), [node(msg("u2", "分支一的问题"), [arm1Leaf])]); + const arm2 = node(msg("u2b", "分支二的问题"), [node(msg("a2b", "答二"))]); + const a1 = node(msg("a1", "答"), [arm1, arm2]); + const root = node(msg("u1", "第一问"), [a1]); + + const [projectedRoot] = projectTreeForResponse([root]); + const projectedA1 = projectedRoot.children[0]; + assert.equal(projectedA1.entry.id, "a1"); + // arm 1: s1 compressed away, u2 visible with the question text + assert.equal(projectedA1.children[0].entry.id, "u2"); + assert.deepEqual(projectedA1.children[0].compressedEntryIds, ["s1"]); + // arm 2: u2b visible directly + assert.equal(projectedA1.children[1].entry.id, "u2b"); + // leaves still kept below the arm heads + assert.equal(projectedA1.children[0].children[0].entry.id, "a2"); + assert.equal(projectedA1.children[1].children[0].entry.id, "a2b"); +}); + +test("linear sessions still project to root + leaf only", () => { + const root = node(msg("u1", "第一问"), [node(msg("a1", "答"))]); + const [projected] = projectTreeForResponse([root]); + assert.equal(projected.entry.id, "u1"); + assert.equal(projected.children.length, 1); + assert.equal(projected.children[0].entry.id, "a1"); +}); diff --git a/lib/project-tree.ts b/lib/project-tree.ts new file mode 100644 index 000000000..0143d0fe3 --- /dev/null +++ b/lib/project-tree.ts @@ -0,0 +1,111 @@ +// BranchNavigator still traverses recursively, so keep the response tree shallow. +export 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. + */ +export function projectTreeForResponse( + nodes: T[] +): T[] { + const keep = new Set(); + const roots = new Set(nodes); + const seen = new Set(); + 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); + } + + // Keep the first message entry of each diverging arm so the client can + // label the branch with the message it diverged with + if (node.children.length > 1) { + for (const child of node.children) { + let cursor = child; + while (!keep.has(cursor) && cursor.entry.type !== "message" && cursor.children.length === 1) { + cursor = cursor.children[0]; + } + if (cursor.entry.type === "message") keep.add(cursor); + } + } + + 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(); + + 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; +} From 582f7684f9c66f2e2f5636f4774e7e58b6882aee Mon Sep 17 00:00:00 2001 From: Alex Yang Date: Tue, 11 Aug 2026 21:41:29 +0800 Subject: [PATCH 2/2] fix: keep branch previews lightweight --- components/BranchNavigator.test.mjs | 54 +++++++++++-- components/BranchNavigator.tsx | 30 +++++--- lib/project-tree.test.mjs | 114 +++++++++++++++++++++++----- lib/project-tree.ts | 103 ++++++++++++++++++++----- lib/types.ts | 6 ++ 5 files changed, 254 insertions(+), 53 deletions(-) diff --git a/components/BranchNavigator.test.mjs b/components/BranchNavigator.test.mjs index 481c201c7..324aced4d 100644 --- a/components/BranchNavigator.test.mjs +++ b/components/BranchNavigator.test.mjs @@ -10,6 +10,7 @@ const { compressChain, selectTopLevelBranches } = await jiti.import("./BranchNav 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", () => { @@ -27,6 +28,20 @@ test("compressChain skips non-message entries such as session_info", () => { 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); @@ -51,12 +66,41 @@ test("selectTopLevelBranches returns empty for a linear session", () => { assert.deepEqual(selectTopLevelBranches([root]), []); }); -test("selectTopLevelBranches works on server-projected tree shapes", () => { - const arm1 = { entry: msg("u2", "user", "分支一"), children: [{ entry: msg("a2", "assistant", "答一"), children: [], compressedEntryIds: ["s1"] }] }; - const arm2 = { entry: msg("u2b", "user", "分支二"), children: [{ entry: msg("a2b", "assistant", "答二"), children: [] }] }; +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), ["u2", "u2b"]); - assert.equal(compressChain(topLevel[0]).labelEntry.id, "u2"); + 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"]); }); diff --git a/components/BranchNavigator.tsx b/components/BranchNavigator.tsx index d17b7ce23..a1ebbf3fa 100644 --- a/components/BranchNavigator.tsx +++ b/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 { @@ -46,19 +46,25 @@ function isMessageEntry(entry: SessionEntry): boolean { // Compress a visible linear chain into the first branching/leaf node. // Server-side compressed IDs also count as skipped nodes. -// labelEntry is the first message entry on the chain — the message the branch -// diverged with — so labels read as the question, not the tail of the chain. -// Falls back to the chain end when the chain has no message entry at all. -export function compressChain(node: SessionTreeNode): { node: SessionTreeNode; skipped: number; labelEntry: SessionEntry } { +// 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, labelEntry: labelEntry ?? current.entry }; + return { node: current, skipped, branchPreview, labelEntry: labelEntry ?? current.entry }; } // Top-level rows of the panel: with multiple roots (a branch was started from @@ -111,13 +117,15 @@ interface TreeNodeProps { } function TreeNodeView({ node, activePathIds, depth, isLast, parentLines, onSelect }: TreeNodeProps) { - const { node: rep, skipped, labelEntry } = compressChain(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(labelEntry); - const role = isMessageEntry(labelEntry) - ? (labelEntry as { message: { role: string } }).message.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 (
diff --git a/lib/project-tree.test.mjs b/lib/project-tree.test.mjs index 08fed857a..ac9fde469 100644 --- a/lib/project-tree.test.mjs +++ b/lib/project-tree.test.mjs @@ -1,35 +1,113 @@ import assert from "node:assert/strict"; import test from "node:test"; -const { projectTreeForResponse } = await import("./project-tree.ts"); +const { MAX_PROJECTED_TREE_DEPTH, projectTreeForResponse } = await import("./project-tree.ts"); -const msg = (id, text) => ({ type: "message", id, parentId: null, timestamp: "t", message: { role: "user", content: text } }); +const msg = (id, role, content) => ({ type: "message", id, parentId: null, timestamp: "t", message: { role, content } }); 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("keeps the first message entry of each diverging arm", () => { - // A1 branches: arm 1 starts with session_info then U2; arm 2 starts with U2b - const arm1Leaf = node(msg("a2", "答一")); - const arm1 = node(info("s1"), [node(msg("u2", "分支一的问题"), [arm1Leaf])]); - const arm2 = node(msg("u2b", "分支二的问题"), [node(msg("a2b", "答二"))]); - const a1 = node(msg("a1", "答"), [arm1, arm2]); - const root = node(msg("u1", "第一问"), [a1]); +function findProjectedNode(nodes, id) { + const pending = [...nodes]; + while (pending.length > 0) { + const current = pending.pop(); + if (current.entry.id === id) return current; + pending.push(...current.children); + } + return undefined; +} + +test("attaches the first branch message preview to the contracted representative", () => { + const arm1Leaf = node(msg("a2", "assistant", "答一")); + const arm1 = node(info("s1"), [node(msg("u2", "user", [ + { type: "text", text: "分支一的问题" }, + { type: "image", data: "base64-secret", mimeType: "image/png" }, + ]), [arm1Leaf])]); + const arm2 = node(msg("u2b", "user", "分支二的问题"), [node(msg("a2b", "assistant", "答二"))]); + const a1 = node(msg("a1", "assistant", "答"), [arm1, arm2]); + const root = node(msg("u1", "user", "第一问"), [a1]); const [projectedRoot] = projectTreeForResponse([root]); const projectedA1 = projectedRoot.children[0]; assert.equal(projectedA1.entry.id, "a1"); - // arm 1: s1 compressed away, u2 visible with the question text - assert.equal(projectedA1.children[0].entry.id, "u2"); - assert.deepEqual(projectedA1.children[0].compressedEntryIds, ["s1"]); - // arm 2: u2b visible directly - assert.equal(projectedA1.children[1].entry.id, "u2b"); - // leaves still kept below the arm heads - assert.equal(projectedA1.children[0].children[0].entry.id, "a2"); - assert.equal(projectedA1.children[1].children[0].entry.id, "a2b"); + assert.equal(projectedA1.children[0].entry.id, "a2"); + assert.deepEqual(projectedA1.children[0].compressedEntryIds, ["s1", "u2"]); + assert.deepEqual(projectedA1.children[0].branchPreview, { role: "user", text: "分支一的问题" }); + assert.equal(projectedA1.children[1].entry.id, "a2b"); + assert.deepEqual(projectedA1.children[1].compressedEntryIds, ["u2b"]); + assert.deepEqual(projectedA1.children[1].branchPreview, { role: "user", text: "分支二的问题" }); + assert.equal(JSON.stringify(projectedRoot).includes("base64-secret"), false); +}); + +test("does not copy thinking or tool payloads from a compressed assistant label source", () => { + const assistant = node(msg("a1", "assistant", [ + { type: "thinking", thinking: "thinking-secret" }, + { type: "text", text: "可见回答" }, + { type: "toolCall", id: "tc1", name: "read", arguments: { value: "tool-secret" } }, + ]), [node(info("leaf1"))]); + const sibling = node(msg("u2", "user", "另一个分支"), [node(info("leaf2"))]); + const projected = projectTreeForResponse([node(info("root"), [assistant, sibling])]); + const leaf = findProjectedNode(projected, "leaf1"); + + assert.deepEqual(leaf.branchPreview, { role: "assistant", text: "可见回答" }); + const serialized = JSON.stringify(projected); + assert.equal(serialized.includes("thinking-secret"), false); + assert.equal(serialized.includes("tool-secret"), false); +}); + +test("carries previews through non-message roots in multi-root trees", () => { + const root1 = node(model("m1"), [node(msg("u1", "user", "第一个问题"), [node(msg("a1", "assistant", "回答一"))])]); + const root2 = node(info("s2"), [node(msg("u2", "user", "第二个问题"), [node(msg("a2", "assistant", "回答二"))])]); + const projected = projectTreeForResponse([root1, root2]); + + assert.deepEqual(projected[0].children[0].branchPreview, { role: "user", text: "第一个问题" }); + assert.deepEqual(projected[1].children[0].branchPreview, { role: "user", text: "第二个问题" }); + assert.deepEqual(projected[0].children[0].compressedEntryIds, ["u1"]); + assert.deepEqual(projected[1].children[0].compressedEntryIds, ["u2"]); +}); + +test("normalizes and bounds preview text and labels image-only messages", () => { + const longText = ` 第一行\n\n第二行 ${"x".repeat(80)}`; + const textArm = node(msg("u1", "user", longText), [node(info("leaf1"))]); + const imageArm = node(msg("u2", "user", [{ type: "image", data: "secret-image", mimeType: "image/png" }]), [node(info("leaf2"))]); + const projected = projectTreeForResponse([node(info("root"), [textArm, imageArm])]); + const textPreview = findProjectedNode(projected, "leaf1").branchPreview; + const imagePreview = findProjectedNode(projected, "leaf2").branchPreview; + + assert.equal(textPreview.text.startsWith("第一行 第二行 "), true); + assert.equal(textPreview.text.length, 41); + assert.equal(textPreview.text.endsWith("…"), true); + assert.deepEqual(imagePreview, { role: "user", text: "[image]" }); + assert.equal(JSON.stringify(projected).includes("secret-image"), false); +}); + +test("does not copy unknown message roles into previews", () => { + const unknownRole = `unknown-role-${"x".repeat(80)}`; + const arm = node(msg("m1", unknownRole, "可见内容"), [node(info("leaf1"))]); + const sibling = node(msg("u2", "user", "另一个分支"), [node(info("leaf2"))]); + const projected = projectTreeForResponse([node(info("root"), [arm, sibling])]); + const preview = findProjectedNode(projected, "leaf1").branchPreview; + + assert.deepEqual(preview, { text: "可见内容" }); + assert.equal(JSON.stringify(projected).includes("unknown-role-"), false); +}); + +test("carries previews through the depth-limit flattening path", () => { + let deepArm = node(info("prefix"), [node(msg("deep-user", "user", "深层问题"), [node(info("deep-leaf"))])]); + for (let i = 0; i < MAX_PROJECTED_TREE_DEPTH + 2; i++) { + deepArm = node(info(`branch-${i}`), [deepArm, node(info(`side-${i}`))]); + } + + const projected = projectTreeForResponse([deepArm]); + const leaf = findProjectedNode(projected, "deep-leaf"); + assert.deepEqual(leaf.branchPreview, { role: "user", text: "深层问题" }); + assert.deepEqual(leaf.compressedEntryIds, ["prefix", "deep-user"]); + assert.equal(findProjectedNode(projected, "deep-user"), undefined); }); test("linear sessions still project to root + leaf only", () => { - const root = node(msg("u1", "第一问"), [node(msg("a1", "答"))]); + const root = node(msg("u1", "user", "第一问"), [node(msg("a1", "assistant", "答"))]); const [projected] = projectTreeForResponse([root]); assert.equal(projected.entry.id, "u1"); assert.equal(projected.children.length, 1); diff --git a/lib/project-tree.ts b/lib/project-tree.ts index 0143d0fe3..5e65db10c 100644 --- a/lib/project-tree.ts +++ b/lib/project-tree.ts @@ -1,5 +1,73 @@ +import type { BranchPreview } from "@/lib/types"; + // BranchNavigator still traverses recursively, so keep the response tree shallow. export const MAX_PROJECTED_TREE_DEPTH = 200; +const MAX_BRANCH_PREVIEW_LENGTH = 40; + +type ProjectableEntry = { + id: string; + type: string; + message?: unknown; +}; + +type ProjectableTreeNode = { + entry: ProjectableEntry; + children: T[]; + compressedEntryIds?: string[]; + branchPreview?: BranchPreview; +}; + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value); +} + +function appendPreviewText(current: string, value: unknown): string { + if (typeof value !== "string" || current.length > MAX_BRANCH_PREVIEW_LENGTH) return current; + const normalized = value.replace(/\s+/g, " ").trim(); + if (!normalized) return current; + const separator = current ? " " : ""; + const prefix = current + separator; + if (prefix.length >= MAX_BRANCH_PREVIEW_LENGTH + 1) { + return prefix.slice(0, MAX_BRANCH_PREVIEW_LENGTH + 1); + } + const remaining = MAX_BRANCH_PREVIEW_LENGTH + 1 - prefix.length; + return prefix + normalized.slice(0, remaining); +} + +function previewForEntry(entry: ProjectableEntry): BranchPreview | undefined { + if (entry.type !== "message" || !isRecord(entry.message) || typeof entry.message.role !== "string") { + return undefined; + } + + const content = entry.message.content; + let text = ""; + let hasImage = false; + if (typeof content === "string") { + text = appendPreviewText(text, content); + } else if (Array.isArray(content)) { + for (const block of content) { + if (!isRecord(block)) continue; + if (block.type === "image") hasImage = true; + if (block.type === "text") text = appendPreviewText(text, block.text); + if (text.length > MAX_BRANCH_PREVIEW_LENGTH) break; + } + } + + if (text.length > MAX_BRANCH_PREVIEW_LENGTH) { + text = text.slice(0, MAX_BRANCH_PREVIEW_LENGTH) + "…"; + } else if (!text) { + text = hasImage + ? "[image]" + : entry.message.role === "assistant" + ? "[assistant]" + : "message"; + } + + const role = entry.message.role === "user" || entry.message.role === "assistant" + ? entry.message.role + : undefined; + return { ...(role ? { role } : {}), text }; +} /** * Project the session tree into the shallow navigation tree sent to the client. @@ -7,7 +75,7 @@ export const MAX_PROJECTED_TREE_DEPTH = 200; * 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. */ -export function projectTreeForResponse( +export function projectTreeForResponse>( nodes: T[] ): T[] { const keep = new Set(); @@ -27,29 +95,18 @@ export function projectTreeForResponse 1) { - for (const child of node.children) { - let cursor = child; - while (!keep.has(cursor) && cursor.entry.type !== "message" && cursor.children.length === 1) { - cursor = cursor.children[0]; - } - if (cursor.entry.type === "message") keep.add(cursor); - } - } - for (const child of node.children) { stack.push(child); } } - const cloneNode = (node: T, compressedEntryIds?: string[]): T => ({ + const cloneNode = (node: T, compressedEntryIds?: string[], branchPreview?: BranchPreview): T => ({ ...node, children: [], ...(compressedEntryIds?.length ? { compressedEntryIds } : {}), + ...(branchPreview ? { branchPreview } : {}), }); - const projectedRoots = nodes.map((node) => cloneNode(node)); + const projectedRoots = nodes.map((node) => cloneNode(node, undefined, previewForEntry(node.entry))); const tasks = nodes.map((source, index) => ({ source, projected: projectedRoots[index], @@ -57,16 +114,21 @@ export function projectTreeForResponse { - const pending = [{ node: source, compressedEntryIds: [] as string[] }]; + const pending = [{ + node: source, + compressedEntryIds: [] as string[], + branchPreview: undefined as BranchPreview | undefined, + }]; const flattenedSeen = new Set(); while (pending.length > 0) { - const { node, compressedEntryIds } = pending.pop()!; + const { node, compressedEntryIds, branchPreview } = pending.pop()!; if (flattenedSeen.has(node)) continue; flattenedSeen.add(node); + const nextPreview = branchPreview ?? previewForEntry(node.entry); if (keep.has(node)) { - projectedParent.children.push(cloneNode(node, compressedEntryIds)); + projectedParent.children.push(cloneNode(node, compressedEntryIds, nextPreview)); } for (let i = node.children.length - 1; i >= 0; i--) { @@ -75,6 +137,7 @@ export function projectTreeForResponse