@aleph-alpha/chat-kit 1.11.1 → 1.12.0

package/src/helpers/conversationTree.spec.ts

@@ -7,6 +7,7 @@ import type {
 import {
   addNode,
   createEmptyTree,
+  cutBranchAfter,
   getLinearPath,
   navigate,
   updateNode,
@@ -193,6 +194,35 @@ describe('getLinearPath', () => {
     expect(u1Msg.status).toBe('completed');
     expect(u1Msg.createdAt).toEqual(new Date('2025-01-01'));
   });
+
+  it('forwards node metadata onto the BaseMessage projection', () => {
+    const tree = createEmptyTree();
+    addNode(
+      tree,
+      makeTreeMessage({
+        id: 'u1',
+        content: [text('Hello')],
+        metadata: {
+          attachments: [
+            { name: 'report.pdf', fileId: 'f1', searchStoreId: 's1' },
+          ],
+        },
+      }),
+    );
+
+    const path = getLinearPath(tree);
+    expect(path[0].metadata).toEqual({
+      attachments: [{ name: 'report.pdf', fileId: 'f1', searchStoreId: 's1' }],
+    });
+  });
+
+  it('omits metadata when the node has none', () => {
+    const tree = createEmptyTree();
+    addNode(tree, makeTreeMessage({ id: 'u1', content: [text('Hello')] }));
+
+    const path = getLinearPath(tree);
+    expect('metadata' in path[0]).toBe(false);
+  });
 });
 
 describe('addNode', () => {
@@ -287,6 +317,38 @@ describe('updateNode', () => {
     updateNode(tree, 'nonexistent', { content: [text('x')] });
     expect(Object.keys(tree.nodes).length).toBe(nodeCount);
   });
+
+  it('sets metadata when supplied', () => {
+    const tree = makeBranchingTree();
+    updateNode(tree, 'u1', {
+      metadata: {
+        attachments: [{ name: 'report.pdf', fileId: 'f1' }],
+      },
+    });
+    expect(tree.nodes.u1.metadata).toEqual({
+      attachments: [{ name: 'report.pdf', fileId: 'f1' }],
+    });
+  });
+
+  it('clears metadata when explicitly set to undefined', () => {
+    const tree = makeBranchingTree();
+    tree.nodes.u1.metadata = {
+      attachments: [{ name: 'old.pdf' }],
+    };
+    updateNode(tree, 'u1', { metadata: undefined });
+    expect(tree.nodes.u1.metadata).toBeUndefined();
+  });
+
+  it('leaves metadata untouched when the key is absent from updates', () => {
+    const tree = makeBranchingTree();
+    tree.nodes.u1.metadata = {
+      attachments: [{ name: 'keep.pdf' }],
+    };
+    updateNode(tree, 'u1', { content: [text('Updated')] });
+    expect(tree.nodes.u1.metadata).toEqual({
+      attachments: [{ name: 'keep.pdf' }],
+    });
+  });
 });
 
 describe('updateNodeId', () => {
@@ -385,6 +447,78 @@ describe('navigate', () => {
   });
 });
 
+describe('cutBranchAfter', () => {
+  it('removes every descendant of the target node', () => {
+    const tree = makeBranchingTree();
+
+    cutBranchAfter(tree, 'u2');
+
+    expect(tree.nodes.a2).toBeUndefined();
+    expect(tree.nodes.a3).toBeUndefined();
+  });
+
+  it('keeps the target node and clears its children', () => {
+    const tree = makeBranchingTree();
+
+    cutBranchAfter(tree, 'u2');
+
+    expect(tree.nodes.u2).toBeDefined();
+    expect(tree.nodes.u2.childrenIds).toEqual([]);
+    expect(tree.nodes.u2.activeChildIndex).toBe(0);
+  });
+
+  it('makes the target node the leaf of the active path', () => {
+    const tree = makeBranchingTree();
+
+    cutBranchAfter(tree, 'u2');
+
+    const path = getLinearPath(tree);
+    expect(path.map((m) => m.id)).toEqual(['u1', 'a1', 'u2']);
+  });
+
+  it('removes the entire subtree, including grandchildren', () => {
+    const tree = makeBranchingTree();
+
+    cutBranchAfter(tree, 'u1');
+
+    expect(Object.keys(tree.nodes)).toEqual(['u1']);
+    expect(tree.nodes.u1.childrenIds).toEqual([]);
+  });
+
+  it('only cuts descendants — sibling branches are untouched', () => {
+    const tree = makeBranchingTree();
+
+    cutBranchAfter(tree, 'a1');
+
+    expect(tree.nodes.u2).toBeUndefined();
+    expect(tree.nodes.a2).toBeUndefined();
+    expect(tree.nodes.a3).toBeUndefined();
+
+    expect(tree.nodes.a1b).toBeDefined();
+    expect(tree.nodes.u3).toBeDefined();
+    expect(tree.nodes.u1.childrenIds).toEqual(['a1', 'a1b']);
+  });
+
+  it('is a no-op when the node has no children', () => {
+    const tree = makeBranchingTree();
+    const nodeCount = Object.keys(tree.nodes).length;
+
+    cutBranchAfter(tree, 'a2');
+
+    expect(Object.keys(tree.nodes).length).toBe(nodeCount);
+    expect(tree.nodes.a2.childrenIds).toEqual([]);
+  });
+
+  it('is a no-op for a nonexistent id', () => {
+    const tree = makeBranchingTree();
+    const before = Object.keys(tree.nodes).sort();
+
+    cutBranchAfter(tree, 'nonexistent');
+
+    expect(Object.keys(tree.nodes).sort()).toEqual(before);
+  });
+});
+
 describe('integration: full branching scenario', () => {
   it('simulates send → regenerate → navigate workflow', () => {
     const tree = createEmptyTree();

package/src/helpers/conversationTree.ts

@@ -25,6 +25,7 @@ export function getLinearPath(tree: ConversationTree): BaseMessage[] {
       content: node.content,
       createdAt: node.createdAt,
       status: node.status,
+      ...(node.metadata !== undefined ? { metadata: node.metadata } : {}),
     };
 
     if (node.parentId !== null) {
@@ -61,7 +62,9 @@ export function addNode(tree: ConversationTree, node: TreeMessage): void {
 export function updateNode(
   tree: ConversationTree,
   id: string,
-  updates: Partial<Pick<TreeMessage, 'content' | 'status' | 'role'>>,
+  updates: Partial<
+    Pick<TreeMessage, 'content' | 'status' | 'role' | 'metadata'>
+  >,
 ): void {
   const node = tree.nodes[id];
   if (!node) return;
@@ -69,6 +72,9 @@ export function updateNode(
   if (updates.content !== undefined) node.content = updates.content;
   if (updates.status !== undefined) node.status = updates.status;
   if (updates.role !== undefined) node.role = updates.role;
+  // `metadata` is intentionally allowed through with an `undefined` value so
+  // callers can clear it; the call site checks for `'metadata' in updates`.
+  if ('metadata' in updates) node.metadata = updates.metadata;
 }
 
 export function updateNodeId(
@@ -99,6 +105,31 @@ export function updateNodeId(
   if (tree.rootId === oldId) tree.rootId = newId;
 }
 
+/**
+ * Remove every descendant of the node with the given id, keeping the node
+ * itself intact. The node's `childrenIds` is cleared and `activeChildIndex`
+ * is reset to 0, which makes the node the new leaf of its branch.
+ *
+ * Useful when editing a message and discarding the conversation that was
+ * generated after it.
+ */
+export function cutBranchAfter(tree: ConversationTree, id: string): void {
+  const node = tree.nodes[id];
+  if (!node) return;
+
+  const stack: string[] = [...node.childrenIds];
+  while (stack.length > 0) {
+    const currentId = stack.pop() as string;
+    const current = tree.nodes[currentId];
+    if (!current) continue;
+    stack.push(...current.childrenIds);
+    delete tree.nodes[currentId];
+  }
+
+  node.childrenIds = [];
+  node.activeChildIndex = 0;
+}
+
 /**
  * Navigate to a different branch by updating the parent's activeChildIndex
  * to point at the target sibling.

package/src/index.ts

@@ -107,6 +107,7 @@ export type {
   ConversationTree,
   MessageContent,
   MessageContentStatus,
+  MessageMetadata,
   MessageReasoningContent,
   MessageRole,
   MessageStatus,

package/src/types.ts

@@ -68,6 +68,19 @@ export type MessageContent =
   | MessageToolContent
   | MessageReasoningContent;
 
+/**
+ * Open-ended metadata bag carried alongside a message. chat-kit is
+ * intentionally agnostic about its contents — adapter authors
+ * (`responsesToTree`, …) and message handlers preserve the value verbatim
+ * through the tree → linear path projection, but never inspect it.
+ *
+ * Consumers attach their own well-known keys (attachments, citations,
+ * ratings, …) and define their own typed projections at the read site
+ * rather than extending this base type. That keeps the handler API stable
+ * even as the shape of what gets stored on a message evolves.
+ */
+export type MessageMetadata = Record<string, unknown>;
+
 export interface BaseMessage {
   id: string;
   role: MessageRole;
@@ -76,6 +89,7 @@ export interface BaseMessage {
   status: MessageStatus;
   siblingIds?: string[];
   siblingIndex?: number;
+  metadata?: MessageMetadata;
 }
 
 export interface TreeMessage {
@@ -87,6 +101,7 @@ export interface TreeMessage {
   content: MessageContent[];
   status: MessageStatus;
   createdAt: Date;
+  metadata?: MessageMetadata;
 }
 
 export interface ConversationTree {