@aleph-alpha/chat-kit 5.5.0 → 5.6.0

package/src/helpers/conversationTree.spec.ts

@@ -518,6 +518,67 @@ describe('cutBranchAfter', () => {
 
     expect(Object.keys(tree.nodes).sort()).toEqual(before);
   });
+
+  it('clears comparison state on the target — no children left to compare', () => {
+    const tree = makeBranchingTree();
+    tree.nodes.u2.comparing = true;
+    tree.nodes.u2.comparisonWinnerId = 'a2';
+
+    cutBranchAfter(tree, 'u2');
+
+    expect(tree.nodes.u2.comparing).toBeUndefined();
+    expect(tree.nodes.u2.comparisonWinnerId).toBeUndefined();
+  });
+
+  it('spares the exempted child and makes it the sole active child', () => {
+    const tree = makeBranchingTree();
+    // u2 has children [a2, a3]; keep a3.
+    cutBranchAfter(tree, 'u2', 'a3');
+
+    expect(tree.nodes.a2).toBeUndefined();
+    expect(tree.nodes.a3).toBeDefined();
+    expect(tree.nodes.u2.childrenIds).toEqual(['a3']);
+    expect(tree.nodes.u2.activeChildIndex).toBe(0);
+    expect(getLinearPath(tree).map((m) => m.id)).toEqual([
+      'u1',
+      'a1',
+      'u2',
+      'a3',
+    ]);
+  });
+
+  it('keeps the exempted child subtree intact', () => {
+    const tree = makeBranchingTree();
+    // Give a3 a follow-up so we can confirm the whole branch survives.
+    addNode(tree, makeTreeMessage({ id: 'u4', parentId: 'a3', role: 'user' }));
+
+    cutBranchAfter(tree, 'u2', 'a3');
+
+    expect(tree.nodes.a3).toBeDefined();
+    expect(tree.nodes.u4).toBeDefined();
+  });
+
+  it('clears comparison state even when a child is spared (single lane left)', () => {
+    const tree = makeBranchingTree();
+    tree.nodes.u2.comparing = true;
+    tree.nodes.u2.comparisonWinnerId = 'a3';
+
+    cutBranchAfter(tree, 'u2', 'a3');
+
+    expect(tree.nodes.u2.childrenIds).toEqual(['a3']);
+    expect(tree.nodes.u2.comparing).toBeUndefined();
+    expect(tree.nodes.u2.comparisonWinnerId).toBeUndefined();
+  });
+
+  it('ignores an unknown or non-child keepChildId and removes every branch', () => {
+    const tree = makeBranchingTree();
+
+    cutBranchAfter(tree, 'u2', 'nonexistent');
+
+    expect(tree.nodes.a2).toBeUndefined();
+    expect(tree.nodes.a3).toBeUndefined();
+    expect(tree.nodes.u2.childrenIds).toEqual([]);
+  });
 });
 
 describe('cutBranchAt', () => {
@@ -637,6 +698,43 @@ describe('cutBranchAt', () => {
 
     expect(Object.keys(tree.nodes).sort()).toEqual(before);
   });
+
+  it('clears the parent comparison when fewer than two lanes remain', () => {
+    const tree = makeThreeSiblingTree();
+    tree.nodes.p.comparing = true;
+    // Down to [cA] after removing cB and cC: nothing left to compare.
+    cutBranchAt(tree, 'cB');
+    cutBranchAt(tree, 'cC');
+
+    expect(tree.nodes.p.childrenIds).toEqual(['cA']);
+    expect(tree.nodes.p.comparing).toBeUndefined();
+    expect(tree.nodes.p.comparisonWinnerId).toBeUndefined();
+  });
+
+  it('drops the dangling winner when the winning lane is removed but others remain', () => {
+    const tree = makeThreeSiblingTree();
+    tree.nodes.p.comparing = true;
+    tree.nodes.p.comparisonWinnerId = 'cA';
+
+    cutBranchAt(tree, 'cA');
+
+    // Still a comparison (cB, cC remain) but the recorded vote is gone.
+    expect(tree.nodes.p.childrenIds).toEqual(['cB', 'cC']);
+    expect(tree.nodes.p.comparing).toBe(true);
+    expect(tree.nodes.p.comparisonWinnerId).toBeUndefined();
+  });
+
+  it('keeps the comparison and winner when a losing lane is removed', () => {
+    const tree = makeThreeSiblingTree();
+    tree.nodes.p.comparing = true;
+    tree.nodes.p.comparisonWinnerId = 'cA';
+
+    cutBranchAt(tree, 'cB');
+
+    expect(tree.nodes.p.childrenIds).toEqual(['cA', 'cC']);
+    expect(tree.nodes.p.comparing).toBe(true);
+    expect(tree.nodes.p.comparisonWinnerId).toBe('cA');
+  });
 });
 
 describe('integration: full branching scenario', () => {

package/src/helpers/conversationTree.ts

@@ -1,12 +1,87 @@
-import type { BaseMessage, ConversationTree, TreeMessage } from '../types';
+import type {
+  ActiveComparison,
+  BaseMessage,
+  ComparisonLane,
+  ConversationTree,
+  TreeMessage,
+} from '../types';
 
 export function createEmptyTree(): ConversationTree {
   return { nodes: {}, rootId: null };
 }
 
+/** Project a tree node to a `BaseMessage`, enriched with sibling metadata. */
+function toBaseMessage(
+  nodes: Record<string, TreeMessage>,
+  node: TreeMessage,
+): BaseMessage {
+  const base: BaseMessage = {
+    id: node.id,
+    role: node.role,
+    content: node.content,
+    createdAt: node.createdAt,
+    status: node.status,
+    ...(node.metadata !== undefined ? { metadata: node.metadata } : {}),
+  };
+
+  if (node.parentId !== null) {
+    const parent = nodes[node.parentId];
+    if (parent && parent.childrenIds.length > 1) {
+      base.siblingIds = [...parent.childrenIds];
+      base.siblingIndex = parent.childrenIds.indexOf(node.id);
+    }
+  }
+
+  return base;
+}
+
+/**
+ * Whether a node's children should render as an **active** comparison — i.e.
+ * `comparing` is set, it has at least two lanes, and it hasn't yet collapsed to
+ * history. Unvoted forks are always active (blind); a voted fork is active only
+ * until the winner branch gains a follow-up turn (after which the comparison is
+ * history and the walk follows the winner linearly).
+ *
+ * A single lane is never a comparison: mid-stream a fork can momentarily hold
+ * one lane (the other still arriving) or end up with one (a lane produced no
+ * assistant node). Treating that as active would stop `projectActivePath` early
+ * and leave no linear leaf to send from, so we require two or more children.
+ */
+function isActiveComparison(
+  nodes: Record<string, TreeMessage>,
+  node: TreeMessage,
+): boolean {
+  if (!node.comparing || node.childrenIds.length < 2) return false;
+  const winnerId = node.comparisonWinnerId;
+  if (winnerId == null) return true;
+  const winner = nodes[winnerId];
+  return !winner || winner.childrenIds.length === 0;
+}
+
+/** Walk one lane's active sub-path from a fork-child head downward. */
+function laneWalk(
+  nodes: Record<string, TreeMessage>,
+  headId: string,
+): BaseMessage[] {
+  const path: BaseMessage[] = [];
+  let current: string | null = headId;
+  while (current !== null) {
+    const node: TreeMessage | undefined = nodes[current];
+    if (!node) break;
+    path.push(toBaseMessage(nodes, node));
+    if (node.childrenIds.length === 0) break;
+    current = node.childrenIds[node.activeChildIndex];
+  }
+  return path;
+}
+
 /**
  * Walks from rootId following activeChildIndex at each node.
  * Enriches each node with sibling metadata for the UI.
+ *
+ * Comparison-unaware: it always descends the active child. Consumers that
+ * render comparisons should use {@link projectActivePath}, which stops the
+ * linear trunk at an active comparison fork and returns its lanes separately.
  */
 export function getLinearPath(tree: ConversationTree): BaseMessage[] {
   const { nodes, rootId } = tree;
@@ -19,24 +94,7 @@ export function getLinearPath(tree: ConversationTree): BaseMessage[] {
     const node: TreeMessage | undefined = nodes[current];
     if (!node) break;
 
-    const base: BaseMessage = {
-      id: node.id,
-      role: node.role,
-      content: node.content,
-      createdAt: node.createdAt,
-      status: node.status,
-      ...(node.metadata !== undefined ? { metadata: node.metadata } : {}),
-    };
-
-    if (node.parentId !== null) {
-      const parent = nodes[node.parentId];
-      if (parent && parent.childrenIds.length > 1) {
-        base.siblingIds = [...parent.childrenIds];
-        base.siblingIndex = parent.childrenIds.indexOf(node.id);
-      }
-    }
-
-    path.push(base);
+    path.push(toBaseMessage(nodes, node));
 
     if (node.childrenIds.length === 0) break;
     current = node.childrenIds[node.activeChildIndex];
@@ -45,20 +103,99 @@ export function getLinearPath(tree: ConversationTree): BaseMessage[] {
   return path;
 }
 
-export function addNode(tree: ConversationTree, node: TreeMessage): void {
+/**
+ * Single-pass projection of the active view into a linear `trunk` plus (at
+ * most) one `activeComparison`. The trunk follows `activeChildIndex` from the
+ * root and **stops at an active comparison fork** (the fork parent is the last
+ * trunk item); that fork's children become the comparison lanes. History
+ * comparisons (voted + continued) are transparent — the trunk flows through
+ * the winner as ordinary linear messages.
+ */
+export function projectActivePath(tree: ConversationTree): {
+  trunk: BaseMessage[];
+  activeComparison: ActiveComparison | null;
+} {
+  const { nodes, rootId } = tree;
+  const trunk: BaseMessage[] = [];
+  if (!rootId || !nodes[rootId]) return { trunk, activeComparison: null };
+
+  let current: string | null = rootId;
+  while (current !== null) {
+    const node: TreeMessage | undefined = nodes[current];
+    if (!node) break;
+
+    trunk.push(toBaseMessage(nodes, node));
+
+    if (isActiveComparison(nodes, node)) {
+      const lanes: ComparisonLane[] = node.childrenIds.map((branchId) => ({
+        branchId,
+        path: laneWalk(nodes, branchId),
+      }));
+      const winnerBranchId = node.comparisonWinnerId ?? null;
+      return {
+        trunk,
+        activeComparison: {
+          forkParentId: node.id,
+          resolved: winnerBranchId !== null,
+          winnerBranchId,
+          lanes,
+        },
+      };
+    }
+
+    if (node.childrenIds.length === 0) break;
+    current = node.childrenIds[node.activeChildIndex];
+  }
+
+  return { trunk, activeComparison: null };
+}
+
+export function addNode(
+  tree: ConversationTree,
+  node: TreeMessage,
+  opts: { activate?: boolean } = { activate: true },
+): void {
   tree.nodes[node.id] = node;
 
   if (node.parentId !== null) {
     const parent = tree.nodes[node.parentId];
     if (parent && !parent.childrenIds.includes(node.id)) {
       parent.childrenIds.push(node.id);
-      parent.activeChildIndex = parent.childrenIds.length - 1;
+      // Opening a comparison adds N siblings without picking a winner, so the
+      // caller suppresses activation; ordinary appends activate the new leaf.
+      if (opts.activate)
+        parent.activeChildIndex = parent.childrenIds.length - 1;
     }
   } else if (tree.rootId === null) {
     tree.rootId = node.id;
   }
 }
 
+/** Mark a node's children as an (unresolved) comparison. */
+export function markComparison(tree: ConversationTree, parentId: string): void {
+  const parent = tree.nodes[parentId];
+  if (parent) parent.comparing = true;
+}
+
+/**
+ * Record a comparison vote: set the winner (mirrored to `activeChildIndex`) and
+ * reveal. Leaves `comparing` set — the comparison stays a labelled, revealed
+ * view until the user continues past the winner (which collapses it to
+ * history). No-op if the parent/winner is unknown.
+ */
+export function voteComparison(
+  tree: ConversationTree,
+  parentId: string,
+  winnerBranchId: string,
+): void {
+  const parent = tree.nodes[parentId];
+  if (!parent) return;
+  const idx = parent.childrenIds.indexOf(winnerBranchId);
+  if (idx === -1) return;
+  parent.comparisonWinnerId = winnerBranchId;
+  parent.activeChildIndex = idx;
+}
+
 export function updateNode(
   tree: ConversationTree,
   id: string,
@@ -125,17 +262,48 @@ function deleteSubtrees(tree: ConversationTree, ids: readonly string[]): void {
  * itself intact. The node's `childrenIds` is cleared and `activeChildIndex`
  * is reset to 0, which makes the node the new leaf of its branch.
  *
+ * Pass `keepChildId` (a *direct* child of the node) to exempt one branch: that
+ * child and its subtree survive and become the node's sole, active child while
+ * every sibling branch is dropped. Useful when picking a comparison winner and
+ * discarding the alternatives. An unknown or non-child `keepChildId` is ignored
+ * (every branch is removed, as with no exception).
+ *
+ * Either way the node is left with at most one child, so any comparison state
+ * on it (`comparing` / `comparisonWinnerId`) — impossible with a single lane —
+ * is cleared.
+ *
  * Useful when editing a message and discarding the conversation that was
  * generated after it.
  */
-export function cutBranchAfter(tree: ConversationTree, id: string): void {
+export function cutBranchAfter(
+  tree: ConversationTree,
+  id: string,
+  keepChildId?: string,
+): void {
   const node = tree.nodes[id];
-  if (!node) return;
+  // Guard against a crafted id (`__proto__`, `constructor`) resolving to a
+  // built-in prototype object rather than a real node: the mutations below
+  // would otherwise corrupt `Object.prototype`.
+  if (!node || !Object.prototype.hasOwnProperty.call(tree.nodes, id)) return;
 
-  deleteSubtrees(tree, node.childrenIds);
+  const kept =
+    keepChildId !== undefined && node.childrenIds.includes(keepChildId)
+      ? keepChildId
+      : null;
 
-  node.childrenIds = [];
+  deleteSubtrees(
+    tree,
+    kept === null
+      ? node.childrenIds
+      : node.childrenIds.filter((childId) => childId !== kept),
+  );
+
+  node.childrenIds = kept === null ? [] : [kept];
   node.activeChildIndex = 0;
+  // With at most one child left there is nothing to compare, so any comparison
+  // state on the node is stale — drop it so the node reads as an ordinary leaf.
+  delete node.comparing;
+  delete node.comparisonWinnerId;
 }
 
 /**
@@ -146,7 +314,10 @@ export function cutBranchAfter(tree: ConversationTree, id: string): void {
  * (e.g. an interrupted regenerate) without cropping the others.
  *
  * The parent's `activeChildIndex` is clamped to a remaining sibling so the
- * active path stays valid. Removing the root clears the tree.
+ * active path stays valid, and its comparison state is reconciled: the whole
+ * comparison is dropped once fewer than two lanes survive, and a dangling
+ * `comparisonWinnerId` is cleared when the removed branch was the winner.
+ * Removing the root clears the tree.
  *
  * Prefer {@link cutBranchAfter} when it fits the case: unlinking here scans
  * and splices the parent's `childrenIds` (O(siblings)), whereas
@@ -177,6 +348,20 @@ export function cutBranchAt(tree: ConversationTree, id: string): void {
   if (parent.activeChildIndex >= parent.childrenIds.length) {
     parent.activeChildIndex = Math.max(0, parent.childrenIds.length - 1);
   }
+
+  if (parent.comparing) {
+    // A comparison needs at least two lanes; once we're down to a single (or
+    // no) remaining child there is nothing left to compare, so clear the
+    // comparison entirely and let the survivor render as ordinary history.
+    if (parent.childrenIds.length < 2) {
+      delete parent.comparing;
+      delete parent.comparisonWinnerId;
+    } else if (parent.comparisonWinnerId === id) {
+      // The removed branch was the recorded winner; drop the (now dangling)
+      // vote so the comparison falls back to the blind, awaiting-vote state.
+      delete parent.comparisonWinnerId;
+    }
+  }
 }
 
 /**

package/src/index.ts

@@ -138,7 +138,9 @@ export { formatRelativeTime } from './helpers';
 
 // Shared Types
 export type {
+  ActiveComparison,
   BaseMessage,
+  ComparisonLane,
   Conversation,
   ConversationTree,
   MessageContent,

package/src/services/types.ts

@@ -28,6 +28,13 @@ export interface ConversationSummary {
   title: string;
   createdAt: Date;
   updatedAt: Date;
+  /**
+   * The raw conversation metadata, passed through untouched — including the
+   * `title` (also surfaced above as a convenience). Lets consumers read
+   * app-specific keys — e.g. arena winners — from a `fetchById` result without
+   * a second request. `null` when the server returned no metadata.
+   */
+  metadata?: Record<string, unknown> | null;
 }
 
 /**

package/src/types.ts

@@ -102,6 +102,21 @@ export interface TreeMessage {
   status: MessageStatus;
   createdAt: Date;
   metadata?: MessageMetadata;
+  /**
+   * When set, this node's `childrenIds` are a **comparison** (a preference
+   * pair / draft set rendered side-by-side) rather than ordinary regenerate
+   * alternatives. Set when the comparison is opened and left set once voted —
+   * the comparison stays a distinct, labelled view until the user continues
+   * past the winner, at which point it collapses to linear history.
+   */
+  comparing?: boolean;
+  /**
+   * The child id chosen when a `comparing` node is voted on (mirrored to
+   * `activeChildIndex`). Absent = not yet voted. Its presence flips the
+   * comparison from the blind "awaiting vote" state to the revealed "resolved"
+   * state; once the winner branch gains children the comparison is history.
+   */
+  comparisonWinnerId?: string;
 }
 
 export interface ConversationTree {
@@ -109,6 +124,31 @@ export interface ConversationTree {
   rootId: string | null;
 }
 
+/**
+ * One side of an active comparison: the head child that defines the lane plus
+ * that lane's full active sub-path (length 1 for a fresh pair, longer once the
+ * conversation has continued inside the lane).
+ */
+export interface ComparisonLane {
+  branchId: string;
+  path: BaseMessage[];
+}
+
+/**
+ * The single active comparison surfaced alongside the linear `messages` path.
+ * Non-null only while the active path's tail fork is unresolved (blind) or
+ * voted-but-not-continued (revealed). Its `forkParentId` is the last item of
+ * `messages`; everything below the fork lives in `lanes`.
+ */
+export interface ActiveComparison {
+  forkParentId: string;
+  /** false = awaiting vote (blind); true = voted, not yet continued (revealed). */
+  resolved: boolean;
+  /** The chosen lane's head id once resolved. */
+  winnerBranchId: string | null;
+  lanes: ComparisonLane[];
+}
+
 /**
  * High-level conversation record consumed by `CkHistory`.
  *