@ably/ai-transport 0.0.1 → 0.1.0

package/src/core/transport/{conversation-tree.ts → tree.ts}

@@ -1,5 +1,5 @@
 /**
- * ConversationTree — materializes a branching conversation from a flat
+ * Tree — materializes a branching conversation from a flat
  * oplog of Ably messages using serial-first ordering.
  *
  * Serial order (the total order assigned by Ably) is the primary mechanism
@@ -9,39 +9,91 @@
  *
  * `upsert()` is the sole mutation method. Messages can arrive in any order
  * (live subscription, history pages, seed data) and the tree produces the
- * correct `flatten()` output once all messages are present.
+ * correct `flattenNodes()` output once all messages are present.
  *
- * The tree owns conversation state. `flatten()` returns the linear message
+ * The tree owns conversation state. `flattenNodes()` returns the linear node
  * list for the currently selected branches — this is what the transport's
  * `getMessages()` delegates to.
  */
 
+import type * as Ably from 'ably';
+
 import { HEADER_FORK_OF, HEADER_PARENT } from '../../constants.js';
+import { EventEmitter } from '../../event-emitter.js';
 import type { Logger } from '../../logger.js';
-import type { ConversationNode, ConversationTree } from './types.js';
+import type { MessageNode, Tree, TurnLifecycleEvent } from './types.js';
 
 // ---------------------------------------------------------------------------
 // Internal node type
 // ---------------------------------------------------------------------------
 
 interface InternalNode<TMessage> {
-  node: ConversationNode<TMessage>;
+  node: MessageNode<TMessage>;
   /** Insertion sequence — tiebreaker for null-serial messages. */
   insertSeq: number;
 }
 
+// ---------------------------------------------------------------------------
+// Internal interface — extended surface consumed by View
+// ---------------------------------------------------------------------------
+
+/** Internal tree surface used by View — not part of the public Tree API. */
+export interface TreeInternal<TMessage> extends Tree<TMessage> {
+  /**
+   * Monotonic counter that increments on structural changes (node insert,
+   * delete, serial promotion/reorder) but NOT on content-only updates
+   * (existing node's message replaced). Allows the View to skip full
+   * tree walks when only message content changed.
+   */
+  readonly structuralVersion: number;
+
+  /**
+   * Flatten the tree along selected branches into a linear node list.
+   * The `selections` map provides the selected sibling's msgId at each
+   * fork point, keyed by group root msgId. Fork points not present in
+   * the map default to the latest sibling. If a selectedMsgId is not
+   * found in the sibling group (stale/deleted), falls back to latest.
+   */
+  flattenNodes(selections: Map<string, string>): MessageNode<TMessage>[];
+
+  /**
+   * Get the "group root" msgId for a sibling group — the original message
+   * that all forks in the group trace back to.
+   */
+  getGroupRoot(msgId: string): string;
+
+  /**
+   * Get the sibling group that `msgId` belongs to, as full MessageNode objects.
+   * Allows callers to resolve index ↔ msgId without losing identity.
+   */
+  getSiblingNodes(msgId: string): MessageNode<TMessage>[];
+
+  /** Forward a raw Ably message event to tree subscribers. */
+  emitAblyMessage(msg: Ably.InboundMessage): void;
+  /** Forward a turn lifecycle event to tree subscribers. */
+  emitTurn(event: TurnLifecycleEvent): void;
+  /** Register an active turn. */
+  trackTurn(turnId: string, clientId: string): void;
+  /** Unregister an active turn. */
+  untrackTurn(turnId: string): void;
+}
+
 // ---------------------------------------------------------------------------
 // Implementation
 // ---------------------------------------------------------------------------
 
+/** EventEmitter events map for the tree. */
+interface TreeEventsMap {
+  update: undefined;
+  'ably-message': Ably.InboundMessage;
+  turn: TurnLifecycleEvent;
+}
+
 // Spec: AIT-CT13
-class DefaultConversationTree<TMessage> implements ConversationTree<TMessage> {
+export class DefaultTree<TMessage> implements TreeInternal<TMessage> {
   /** All nodes indexed by msgId (x-ably-msg-id). */
   private readonly _nodeIndex = new Map<string, InternalNode<TMessage>>();
 
-  /** Secondary index: codec message key to msgId. Bridges UIMessage.id to x-ably-msg-id. */
-  private readonly _codecKeyIndex = new Map<string, string>();
-
   /**
    * All nodes sorted by serial (lexicographic). Null-serial messages
    * (optimistic inserts, seed data) sort after all serial-bearing messages,
@@ -55,21 +107,25 @@ class DefaultConversationTree<TMessage> implements ConversationTree<TMessage> {
    */
   private readonly _parentIndex = new Map<string | undefined, Set<string>>();
 
-  /**
-   * Selected sibling index at each fork point, keyed by the msgId of
-   * the first sibling in the group (the fork target). Default: last.
-   */
-  private readonly _selections = new Map<string, number>();
-
-  private readonly _getKey: (message: TMessage) => string;
+  private readonly _emitter: EventEmitter<TreeEventsMap>;
   private readonly _logger: Logger;
 
+  /** Active turns: turnId → clientId. */
+  private readonly _turnClientIds = new Map<string, string>();
+
   /** Monotonically increasing counter for insertion sequence. */
   private _seqCounter = 0;
 
-  constructor(getKey: (message: TMessage) => string, logger: Logger) {
-    this._getKey = getKey;
+  /** Incremented on structural changes; unchanged on content-only updates. */
+  private _structuralVersion = 0;
+
+  get structuralVersion(): number {
+    return this._structuralVersion;
+  }
+
+  constructor(logger: Logger) {
     this._logger = logger;
+    this._emitter = new EventEmitter<TreeEventsMap>(logger);
   }
 
   // -------------------------------------------------------------------------
@@ -171,7 +227,7 @@ class DefaultConversationTree<TMessage> implements ConversationTree<TMessage> {
    * @returns The ordered list of sibling nodes.
    */
   // Spec: AIT-CT13b
-  private _getSiblingGroup(msgId: string): ConversationNode<TMessage>[] {
+  private _getSiblingGroup(msgId: string): MessageNode<TMessage>[] {
     const entry = this._nodeIndex.get(msgId);
     if (!entry) return [];
 
@@ -218,7 +274,7 @@ class DefaultConversationTree<TMessage> implements ConversationTree<TMessage> {
    * @param originalId - The group root to match against.
    * @returns True if the node belongs to the sibling group.
    */
-  private _isSiblingOf(node: ConversationNode<TMessage>, originalId: string): boolean {
+  private _isSiblingOf(node: MessageNode<TMessage>, originalId: string): boolean {
     if (node.msgId === originalId) return true;
     let current = node;
     const visited = new Set<string>([current.msgId]);
@@ -239,7 +295,7 @@ class DefaultConversationTree<TMessage> implements ConversationTree<TMessage> {
    * @param msgId - Any msg-id in the sibling group.
    * @returns The msg-id of the group root.
    */
-  private _getGroupRoot(msgId: string): string {
+  getGroupRoot(msgId: string): string {
     const entry = this._nodeIndex.get(msgId);
     if (!entry) return msgId;
 
@@ -259,8 +315,9 @@ class DefaultConversationTree<TMessage> implements ConversationTree<TMessage> {
   // Public query methods
   // -------------------------------------------------------------------------
 
-  flatten(): TMessage[] {
-    const result: TMessage[] = [];
+  flattenNodes(selections: Map<string, string>): MessageNode<TMessage>[] {
+    this._logger.trace('DefaultTree.flattenNodes();');
+    const result: MessageNode<TMessage>[] = [];
     const currentPath = new Set<string>();
     // Track which sibling groups we've already resolved to avoid
     // re-resolving for every member of the group.
@@ -278,14 +335,18 @@ class DefaultConversationTree<TMessage> implements ConversationTree<TMessage> {
       // Step 2: Check sibling selection.
       const group = this._getSiblingGroup(msgId);
       if (group.length > 1) {
-        const groupRootId = this._getGroupRoot(msgId);
+        const groupRootId = this.getGroupRoot(msgId);
         let selectedId = resolvedGroups.get(groupRootId);
         if (selectedId === undefined) {
-          const selectedIdx = this._selections.get(groupRootId) ?? group.length - 1;
-          const clamped = Math.max(0, Math.min(selectedIdx, group.length - 1));
-          const selected = group[clamped];
-          if (!selected) break; // unreachable: clamped is always in bounds
-          selectedId = selected.msgId;
+          const preferredId = selections.get(groupRootId);
+          // Verify the preferred msgId is in the group, otherwise default to latest
+          if (preferredId && group.some((n) => n.msgId === preferredId)) {
+            selectedId = preferredId;
+          } else {
+            const latest = group.at(-1);
+            if (!latest) break; // unreachable: group.length > 1
+            selectedId = latest.msgId;
+          }
           resolvedGroups.set(groupRootId, selectedId);
         }
         if (msgId !== selectedId) {
@@ -294,49 +355,32 @@ class DefaultConversationTree<TMessage> implements ConversationTree<TMessage> {
       }
 
       currentPath.add(msgId);
-      result.push(node.message);
+      result.push(node);
     }
 
     return result;
   }
 
   getSiblings(msgId: string): TMessage[] {
+    this._logger.trace('DefaultTree.getSiblings();', { msgId });
     return this._getSiblingGroup(msgId).map((n) => n.message);
   }
 
-  hasSiblings(msgId: string): boolean {
-    return this._getSiblingGroup(msgId).length > 1;
-  }
-
-  getSelectedIndex(msgId: string): number {
-    const group = this._getSiblingGroup(msgId);
-    if (group.length <= 1) return 0;
-    const groupRootId = this._getGroupRoot(msgId);
-    const stored = this._selections.get(groupRootId);
-    if (stored !== undefined) return Math.max(0, Math.min(stored, group.length - 1));
-    return group.length - 1; // default: latest
-  }
-
-  // Spec: AIT-CT13c
-  select(msgId: string, index: number): void {
-    this._logger.debug('ConversationTree.select();', { msgId, index });
-    const group = this._getSiblingGroup(msgId);
-    if (group.length <= 1) return;
-    const groupRootId = this._getGroupRoot(msgId);
-    this._selections.set(groupRootId, Math.max(0, Math.min(index, group.length - 1)));
+  getSiblingNodes(msgId: string): MessageNode<TMessage>[] {
+    return this._getSiblingGroup(msgId);
   }
 
-  getNode(msgId: string): ConversationNode<TMessage> | undefined {
-    return this._nodeIndex.get(msgId)?.node;
+  hasSiblings(msgId: string): boolean {
+    return this._getSiblingGroup(msgId).length > 1;
   }
 
-  getNodeByKey(key: string): ConversationNode<TMessage> | undefined {
-    const msgId = this._codecKeyIndex.get(key);
-    if (!msgId) return undefined;
+  getNode(msgId: string): MessageNode<TMessage> | undefined {
+    this._logger.trace('DefaultTree.getNode();', { msgId });
     return this._nodeIndex.get(msgId)?.node;
   }
 
   getHeaders(msgId: string): Record<string, string> | undefined {
+    this._logger.trace('DefaultTree.getHeaders();', { msgId });
     return this._nodeIndex.get(msgId)?.node.headers;
   }
 
@@ -348,9 +392,6 @@ class DefaultConversationTree<TMessage> implements ConversationTree<TMessage> {
     const parentId = headers[HEADER_PARENT] ?? undefined;
     const forkOf = headers[HEADER_FORK_OF] ?? undefined;
 
-    // Maintain codec key → msgId secondary index
-    this._codecKeyIndex.set(this._getKey(message), msgId);
-
     const existing = this._nodeIndex.get(msgId);
     if (existing) {
       // Update in place — message content may have changed (e.g. streaming).
@@ -363,18 +404,21 @@ class DefaultConversationTree<TMessage> implements ConversationTree<TMessage> {
       // Spec: AIT-CT13d
       // Promote serial: optimistic (null) → server-assigned on relay.
       if (serial && !existing.node.serial) {
-        this._logger.debug('ConversationTree.upsert(); promoting serial', { msgId, serial });
+        this._logger.debug('Tree.upsert(); promoting serial', { msgId, serial });
         existing.node.serial = serial;
         // Re-sort: remove from current position, re-insert at correct position.
         this._removeSorted(existing);
         this._insertSorted(existing);
+        this._structuralVersion++;
       }
+      this._emitter.emit('update');
       return;
     }
 
-    this._logger.trace('ConversationTree.upsert(); inserting new node', { msgId, parentId, forkOf });
+    this._logger.trace('Tree.upsert(); inserting new node', { msgId, parentId, forkOf });
 
-    const node: ConversationNode<TMessage> = {
+    const node: MessageNode<TMessage> = {
+      kind: 'message',
       message,
       msgId,
       parentId,
@@ -387,22 +431,18 @@ class DefaultConversationTree<TMessage> implements ConversationTree<TMessage> {
     this._nodeIndex.set(msgId, internal);
     this._addToParentIndex(parentId, msgId);
     this._insertSorted(internal);
+    this._structuralVersion++;
+    this._emitter.emit('update');
   }
 
   delete(msgId: string): void {
     const entry = this._nodeIndex.get(msgId);
     if (!entry) return;
 
-    this._logger.debug('ConversationTree.delete();', { msgId });
+    this._logger.debug('Tree.delete();', { msgId });
 
     const { node } = entry;
 
-    // Clean up secondary index
-    const codecKey = this._getKey(node.message);
-    if (this._codecKeyIndex.get(codecKey) === msgId) {
-      this._codecKeyIndex.delete(codecKey);
-    }
-
     // Remove from parent index
     this._removeFromParentIndex(node.parentId, msgId);
 
@@ -411,10 +451,87 @@ class DefaultConversationTree<TMessage> implements ConversationTree<TMessage> {
 
     // Remove from primary index
     this._nodeIndex.delete(msgId);
-    this._selections.delete(msgId);
 
-    // Children are NOT deleted — they become unreachable in flatten()
+    // Children are NOT deleted — they become unreachable in flattenNodes()
     // because their parent is no longer on the active path.
+    this._structuralVersion++;
+    this._emitter.emit('update');
+  }
+
+  // -------------------------------------------------------------------------
+  // Events
+  // -------------------------------------------------------------------------
+
+  // Spec: AIT-CT17
+  getActiveTurnIds(): Map<string, Set<string>> {
+    this._logger.trace('DefaultTree.getActiveTurnIds();');
+    const result = new Map<string, Set<string>>();
+    for (const [turnId, clientId] of this._turnClientIds) {
+      let set = result.get(clientId);
+      if (!set) {
+        set = new Set<string>();
+        result.set(clientId, set);
+      }
+      set.add(turnId);
+    }
+    return result;
+  }
+
+  // Spec: AIT-CT8b, AIT-CT8e
+  on(event: 'update', handler: () => void): () => void;
+  on(event: 'ably-message', handler: (msg: Ably.InboundMessage) => void): () => void;
+  on(event: 'turn', handler: (event: TurnLifecycleEvent) => void): () => void;
+  on(
+    event: 'update' | 'ably-message' | 'turn',
+    handler: (() => void) | ((msg: Ably.InboundMessage) => void) | ((event: TurnLifecycleEvent) => void),
+  ): () => void {
+    // CAST: overload signatures enforce correct handler types per event name.
+    const cb = handler as (arg: TreeEventsMap[keyof TreeEventsMap]) => void;
+    this._emitter.on(event, cb);
+    return () => {
+      this._emitter.off(event, cb);
+    };
+  }
+
+  // -------------------------------------------------------------------------
+  // Internal methods (called by the transport, not part of Tree interface)
+  // -------------------------------------------------------------------------
+
+  /**
+   * Forward a raw Ably message event to tree subscribers.
+   * @param msg - The raw Ably message to emit.
+   */
+  emitAblyMessage(msg: Ably.InboundMessage): void {
+    this._logger.trace('DefaultTree.emitAblyMessage();');
+    this._emitter.emit('ably-message', msg);
+  }
+
+  /**
+   * Forward a turn lifecycle event to tree subscribers.
+   * @param event - The turn lifecycle event to emit.
+   */
+  emitTurn(event: TurnLifecycleEvent): void {
+    this._logger.trace('DefaultTree.emitTurn();', { turnId: event.turnId });
+    this._emitter.emit('turn', event);
+  }
+
+  /**
+   * Register an active turn.
+   * @param turnId - The turn's unique identifier.
+   * @param clientId - The client that owns the turn.
+   */
+  trackTurn(turnId: string, clientId: string): void {
+    this._logger.trace('DefaultTree.trackTurn();', { turnId, clientId });
+    this._turnClientIds.set(turnId, clientId);
+  }
+
+  /**
+   * Unregister an active turn.
+   * @param turnId - The turn to untrack.
+   */
+  untrackTurn(turnId: string): void {
+    this._logger.trace('DefaultTree.untrackTurn();', { turnId });
+    this._turnClientIds.delete(turnId);
   }
 }
 
@@ -423,12 +540,10 @@ class DefaultConversationTree<TMessage> implements ConversationTree<TMessage> {
 // ---------------------------------------------------------------------------
 
 /**
- * Create a ConversationTree that materializes branching history from a flat oplog.
- * @param getKey - Codec function that returns a stable key for a domain message.
+ * Create a Tree that materializes branching history from a flat oplog.
  * @param logger - Logger for diagnostic output.
- * @returns A new {@link ConversationTree} instance.
+ * @returns A new {@link DefaultTree} instance. The transport uses DefaultTree
+ *   directly for internal methods (emitAblyMessage, emitTurn, trackTurn, untrackTurn).
+ *   Public consumers see the narrower {@link Tree} interface.
  */
-export const createConversationTree = <TMessage>(
-  getKey: (message: TMessage) => string,
-  logger: Logger,
-): ConversationTree<TMessage> => new DefaultConversationTree(getKey, logger);
+export const createTree = <TMessage>(logger: Logger): DefaultTree<TMessage> => new DefaultTree(logger);

package/src/core/transport/turn-manager.ts

@@ -11,6 +11,8 @@ import type * as Ably from 'ably';
 import {
   EVENT_TURN_END,
   EVENT_TURN_START,
+  HEADER_FORK_OF,
+  HEADER_PARENT,
   HEADER_TURN_CLIENT_ID,
   HEADER_TURN_ID,
   HEADER_TURN_REASON,
@@ -25,7 +27,12 @@ import type { TurnEndReason } from './types.js';
 /** Manages active turns and publishes turn lifecycle events on the channel. */
 export interface TurnManager {
   /** Register a new turn. Publishes turn-start on the channel. Returns AbortSignal. */
-  startTurn(turnId: string, clientId?: string, controller?: AbortController): Promise<AbortSignal>;
+  startTurn(
+    turnId: string,
+    clientId?: string,
+    controller?: AbortController,
+    metadata?: { parent?: string; forkOf?: string },
+  ): Promise<AbortSignal>;
   /** End a turn. Publishes turn-end on the channel. Cleans up internal state. */
   endTurn(turnId: string, reason: TurnEndReason): Promise<void>;
   /** Get the AbortSignal for a turn. */
@@ -44,7 +51,7 @@ export interface TurnManager {
 // Internal state
 // ---------------------------------------------------------------------------
 
-interface TurnState {
+interface ActiveTurnsEntry {
   controller: AbortController;
   clientId: string;
 }
@@ -56,28 +63,39 @@ interface TurnState {
 class DefaultTurnManager implements TurnManager {
   private readonly _channel: Ably.RealtimeChannel;
   private readonly _logger: Logger | undefined;
-  private readonly _activeTurns = new Map<string, TurnState>();
+  private readonly _activeTurns = new Map<string, ActiveTurnsEntry>();
 
   constructor(channel: Ably.RealtimeChannel, logger?: Logger) {
     this._channel = channel;
     this._logger = logger?.withContext({ component: 'TurnManager' });
   }
 
-  async startTurn(turnId: string, clientId?: string, externalController?: AbortController): Promise<AbortSignal> {
+  async startTurn(
+    turnId: string,
+    clientId?: string,
+    externalController?: AbortController,
+    metadata?: { parent?: string; forkOf?: string },
+  ): Promise<AbortSignal> {
     this._logger?.trace('DefaultTurnManager.startTurn();', { turnId, clientId });
 
     const controller = externalController ?? new AbortController();
     const resolvedClientId = clientId ?? '';
     this._activeTurns.set(turnId, { controller, clientId: resolvedClientId });
 
+    const headers: Record<string, string> = {
+      [HEADER_TURN_ID]: turnId,
+      [HEADER_TURN_CLIENT_ID]: resolvedClientId,
+    };
+    if (metadata?.parent !== undefined) {
+      headers[HEADER_PARENT] = metadata.parent;
+    }
+    if (metadata?.forkOf !== undefined) {
+      headers[HEADER_FORK_OF] = metadata.forkOf;
+    }
+
     await this._channel.publish({
       name: EVENT_TURN_START,
-      extras: {
-        headers: {
-          [HEADER_TURN_ID]: turnId,
-          [HEADER_TURN_CLIENT_ID]: resolvedClientId,
-        },
-      },
+      extras: { headers },
     });
 
     this._logger?.debug('DefaultTurnManager.startTurn(); turn started', { turnId });