feat: root-as-head stack rendering with tree expansion

server/lib/threadChain.test.ts

@@ -0,0 +1,175 @@
+import { describe, it, expect, vi, beforeEach } from 'vitest';
+import type { Thread, ThreadChainNode } from '../../shared/types.js';
+
+vi.mock('./threadCrud.js', () => ({
+  getThreads: vi.fn(),
+}));
+
+import { getThreadChain } from './threadChain.js';
+import { getThreads } from './threadCrud.js';
+
+const mockedGetThreads = vi.mocked(getThreads);
+
+function makeThread(overrides: Partial<Thread> & { id: string }): Thread {
+  return {
+    title: `Thread ${overrides.id}`,
+    lastUpdated: '2 hours ago',
+    visibility: 'Private' as const,
+    messages: 5,
+    ...overrides,
+  };
+}
+
+function collectNodeIds(nodes: ThreadChainNode[]): string[] {
+  const ids: string[] = [];
+  for (const node of nodes) {
+    ids.push(node.thread.id);
+    ids.push(...collectNodeIds(node.children));
+  }
+  return ids;
+}
+
+beforeEach(() => {
+  vi.clearAllMocks();
+});
+
+describe('getThreadChain', () => {
+  it('returns empty chain for a thread with no relationships', async () => {
+    mockedGetThreads.mockResolvedValue({
+      threads: [makeThread({ id: 'solo' })],
+      nextCursor: null,
+      hasMore: false,
+    });
+
+    const chain = await getThreadChain('solo');
+    expect(chain.ancestors).toEqual([]);
+    expect(chain.current?.id).toBe('solo');
+    expect(chain.descendantsTree).toEqual([]);
+  });
+
+  it('returns linear ancestors for a simple chain', async () => {
+    mockedGetThreads.mockResolvedValue({
+      threads: [
+        makeThread({ id: 'root' }),
+        makeThread({ id: 'middle', handoffParentId: 'root' }),
+        makeThread({ id: 'leaf', handoffParentId: 'middle' }),
+      ],
+      nextCursor: null,
+      hasMore: false,
+    });
+
+    const chain = await getThreadChain('leaf');
+    expect(chain.ancestors.map((a) => a.id)).toEqual(['root', 'middle']);
+    expect(chain.current?.id).toBe('leaf');
+    expect(chain.descendantsTree).toEqual([]);
+  });
+
+  it('returns tree with fork for multiple children', async () => {
+    // root → mid → [child-a, child-b]
+    mockedGetThreads.mockResolvedValue({
+      threads: [
+        makeThread({ id: 'root' }),
+        makeThread({ id: 'mid', handoffParentId: 'root' }),
+        makeThread({ id: 'child-a', handoffParentId: 'mid' }),
+        makeThread({ id: 'child-b', handoffParentId: 'mid' }),
+      ],
+      nextCursor: null,
+      hasMore: false,
+    });
+
+    const chain = await getThreadChain('root');
+    expect(chain.ancestors).toEqual([]);
+    expect(chain.current?.id).toBe('root');
+
+    // Tree: root's direct child is mid, mid has two children
+    expect(chain.descendantsTree).toHaveLength(1);
+    const midNode = chain.descendantsTree[0];
+    expect(midNode?.thread.id).toBe('mid');
+    expect(midNode?.children).toHaveLength(2);
+    const childIds = midNode?.children.map((c) => c.thread.id) ?? [];
+    expect(childIds).toContain('child-a');
+    expect(childIds).toContain('child-b');
+
+    // All 3 descendants are in the tree
+    const allIds = collectNodeIds(chain.descendantsTree);
+    expect(allIds).toContain('mid');
+    expect(allIds).toContain('child-a');
+    expect(allIds).toContain('child-b');
+  });
+
+  it('builds tree from middle of a chain', async () => {
+    // root → mid → [child-a → grandchild, child-b]
+    mockedGetThreads.mockResolvedValue({
+      threads: [
+        makeThread({ id: 'root' }),
+        makeThread({ id: 'mid', handoffParentId: 'root' }),
+        makeThread({ id: 'child-a', handoffParentId: 'mid' }),
+        makeThread({ id: 'child-b', handoffParentId: 'mid' }),
+        makeThread({ id: 'grandchild', handoffParentId: 'child-a' }),
+      ],
+      nextCursor: null,
+      hasMore: false,
+    });
+
+    const chain = await getThreadChain('mid');
+
+    // Ancestors: just root
+    expect(chain.ancestors.map((a) => a.id)).toEqual(['root']);
+    expect(chain.current?.id).toBe('mid');
+
+    // Descendants tree: two children of mid
+    expect(chain.descendantsTree).toHaveLength(2);
+    const allDescIds = collectNodeIds(chain.descendantsTree);
+    expect(allDescIds).toContain('child-a');
+    expect(allDescIds).toContain('child-b');
+    expect(allDescIds).toContain('grandchild');
+
+    // child-a should have grandchild as its child in the tree
+    const childANode = chain.descendantsTree.find((n) => n.thread.id === 'child-a');
+    expect(childANode?.children).toHaveLength(1);
+    expect(childANode?.children[0]?.thread.id).toBe('grandchild');
+
+    // child-b is a leaf
+    const childBNode = chain.descendantsTree.find((n) => n.thread.id === 'child-b');
+    expect(childBNode?.children).toEqual([]);
+  });
+
+  it('returns null current for unknown thread id', async () => {
+    mockedGetThreads.mockResolvedValue({
+      threads: [makeThread({ id: 'other' })],
+      nextCursor: null,
+      hasMore: false,
+    });
+
+    const chain = await getThreadChain('nonexistent');
+    expect(chain.current).toBeNull();
+    expect(chain.ancestors).toEqual([]);
+    expect(chain.descendantsTree).toEqual([]);
+  });
+
+  it('tree contains all descendants in a complex fork', async () => {
+    // root → [a → [a1, a2], b]
+    mockedGetThreads.mockResolvedValue({
+      threads: [
+        makeThread({ id: 'root' }),
+        makeThread({ id: 'a', handoffParentId: 'root' }),
+        makeThread({ id: 'b', handoffParentId: 'root' }),
+        makeThread({ id: 'a1', handoffParentId: 'a' }),
+        makeThread({ id: 'a2', handoffParentId: 'a' }),
+      ],
+      nextCursor: null,
+      hasMore: false,
+    });
+
+    const chain = await getThreadChain('root');
+    const treeIds = collectNodeIds(chain.descendantsTree).sort();
+    expect(treeIds).toEqual(['a', 'a1', 'a2', 'b']);
+
+    // Verify tree structure: root has 2 direct children
+    expect(chain.descendantsTree).toHaveLength(2);
+    const nodeA = chain.descendantsTree.find((n) => n.thread.id === 'a');
+    expect(nodeA?.children).toHaveLength(2);
+    const nodeB = chain.descendantsTree.find((n) => n.thread.id === 'b');
+    expect(nodeB?.children).toEqual([]);
+  });
+});

server/lib/threadChain.ts

@@ -1,6 +1,7 @@
 import { readFile } from 'fs/promises';
 import { join } from 'path';
-import type { Thread, ThreadChain, ChainThread } from '../../shared/types.js';
+import type { Thread, ThreadChain, ChainThread, ThreadChainNode } from '../../shared/types.js';
+import { buildHandoffGraph } from '../../shared/utils.js';
 import { runAmp, stripAnsi } from './utils.js';
 import { getThreads } from './threadCrud.js';
 import { THREADS_DIR, isHandoffRelationship, type ThreadFile } from './threadTypes.js';
@@ -9,53 +10,54 @@ import { getThreadMetadata, updateLinkedIssue } from './database.js';
 
 export async function getThreadChain(threadId: string): Promise<ThreadChain> {
   const { threads } = await getThreads({ limit: 1000 });
-  const threadMap = new Map(threads.map((t) => [t.id, t]));
-
-  // Build parent → children map using handoffParentId on each thread.
-  const childrenMap = new Map<string, Thread[]>();
-  for (const thread of threads) {
-    if (thread.handoffParentId) {
-      let children = childrenMap.get(thread.handoffParentId);
-      if (!children) {
-        children = [];
-        childrenMap.set(thread.handoffParentId, children);
-      }
-      children.push(thread);
-    }
-  }
+  const { threadMap, parentToChildren } = buildHandoffGraph(threads);
 
-  // 1. Identify chain members by traversal
+  // Walk up to collect ancestor IDs (linear path to root)
   const ancestorIds: string[] = [];
   const visited = new Set<string>([threadId]);
-
-  let currentId: string | null | undefined = threadMap.get(threadId)?.handoffParentId;
-  while (currentId && !visited.has(currentId)) {
-    visited.add(currentId);
-    if (!threadMap.has(currentId)) break;
-    ancestorIds.unshift(currentId);
-    currentId = threadMap.get(currentId)?.handoffParentId;
+  let walkId = threadId;
+  // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition -- runtime guard
+  while (true) {
+    const thread = threadMap.get(walkId);
+    const parentId = thread?.handoffParentId;
+    if (!parentId || visited.has(parentId)) break;
+    if (!threadMap.has(parentId)) break;
+    visited.add(parentId);
+    ancestorIds.unshift(parentId);
+    walkId = parentId;
   }
 
-  const descendantIds: string[] = [];
-  function collectDescendants(id: string): void {
-    const children = childrenMap.get(id) || [];
-    for (const child of children) {
-      if (visited.has(child.id)) continue;
-      visited.add(child.id);
-      descendantIds.push(child.id);
-      collectDescendants(child.id);
+  // Build descendants tree (recursive, supports forks)
+  function buildDescendantNode(id: string): ThreadChainNode | null {
+    const t = threadMap.get(id);
+    if (!t) return null;
+    const childIds = parentToChildren.get(id) || [];
+    const children: ThreadChainNode[] = [];
+    for (const childId of childIds) {
+      if (visited.has(childId)) continue;
+      visited.add(childId);
+      const node = buildDescendantNode(childId);
+      if (node) children.push(node);
     }
+    return { thread: toChainThread(t), children };
   }
-  collectDescendants(threadId);
 
-  // 2. Read handoff comments from thread files for all chain members.
+  const descendantsTree: ThreadChainNode[] = [];
+  const directChildIds = parentToChildren.get(threadId) || [];
+  for (const childId of directChildIds) {
+    if (visited.has(childId)) continue;
+    visited.add(childId);
+    const node = buildDescendantNode(childId);
+    if (node) descendantsTree.push(node);
+  }
+
+  // Read handoff comments from thread files for all chain members.
   // Each thread's `role: "parent"` relationship has the comment describing
   // what was handed off to the next thread in the chain.
-  const allChainIds = [...ancestorIds, threadId, ...descendantIds];
   const commentMap = new Map<string, string>();
 
   await Promise.all(
-    allChainIds.map(async (id) => {
+    [...visited].map(async (id) => {
       try {
         const content = await readFile(join(THREADS_DIR, `${id}.json`), 'utf-8');
         const data = JSON.parse(content) as ThreadFile;
@@ -73,27 +75,37 @@ export async function getThreadChain(threadId: string): Promise<ThreadChain> {
     }),
   );
 
-  // 3. Build ChainThread objects with comments
-  function toChainThread(id: string): ChainThread {
-    const thread = threadMap.get(id);
-    if (!thread) throw new Error(`Thread ${id} not found in map`);
-
-    return {
-      id: thread.id,
-      title: thread.title,
-      lastUpdated: thread.lastUpdatedDate || thread.lastUpdated,
-      workspace: thread.workspace,
-      comment: commentMap.get(id),
-    };
+  // Apply comments to descendants tree
+  function applyComments(node: ThreadChainNode): void {
+    const comment = commentMap.get(node.thread.id);
+    if (comment) node.thread.comment = comment;
+    for (const child of node.children) applyComments(child);
   }
+  for (const node of descendantsTree) applyComments(node);
 
-  const ancestors = ancestorIds.map(toChainThread);
-  const descendants = descendantIds.map(toChainThread);
+  const ancestors = ancestorIds
+    .map((id) => {
+      const t = threadMap.get(id);
+      return t ? toChainThread(t, commentMap.get(id)) : null;
+    })
+    .filter((t): t is ChainThread => t != null);
 
   const currentThread = threadMap.get(threadId);
-  const current: ChainThread | null = currentThread ? toChainThread(threadId) : null;
+  const current: ChainThread | null = currentThread
+    ? toChainThread(currentThread, commentMap.get(threadId))
+    : null;
+
+  return { ancestors, current, descendantsTree };
+}
 
-  return { ancestors, current, descendants };
+function toChainThread(t: Thread, comment?: string): ChainThread {
+  return {
+    id: t.id,
+    title: t.title,
+    lastUpdated: t.lastUpdated,
+    workspace: t.workspace,
+    ...(comment != null ? { comment } : {}),
+  };
 }
 
 interface HandoffResult {

server/lib/threadCrud.ts

@@ -127,18 +127,19 @@ export async function getThreads({
           // Extract handoff relationships
           const relationships = data.relationships || [];
           let handoffParentId: string | null = null;
-          let handoffChildId: string | null = null;
+          const handoffChildIds: string[] = [];
           for (const rel of relationships) {
             if (isHandoffRelationship(rel)) {
               if (rel.role === 'child') {
                 // "I am the child" → threadID is my parent
                 handoffParentId = rel.threadID;
               } else {
-                // "I am the parent" → threadID is my child (use last one seen)
-                handoffChildId = rel.threadID;
+                // "I am the parent" → threadID is my child
+                handoffChildIds.push(rel.threadID);
               }
             }
           }
+          const uniqueChildIds = [...new Set(handoffChildIds)];
 
           // Extract touched files from tool uses
           const touchedFiles = new Set<string>();
@@ -167,7 +168,7 @@ export async function getThreads({
             repo,
             touchedFiles: [...touchedFiles],
             handoffParentId,
-            handoffChildId,
+            handoffChildIds: uniqueChildIds,
           };
         } catch (e) {
           const error = e as Error;