@assistant-ui/core 0.2.6 → 0.2.8

package/src/runtime/utils/message-repository.ts

@@ -1,6 +1,6 @@
 import type { ThreadMessage } from "../../types/message";
 import type { RunConfig } from "../../types/message";
-import { generateId, generateOptimisticId } from "../../utils/id";
+import { generateId } from "../../utils/id";
 import type { ThreadMessageLike } from "./thread-message-like";
 import { getAutoStatus } from "./auto-status";
 import { fromThreadMessageLike } from "./thread-message-like";
@@ -259,20 +259,6 @@ export class MessageRepository {
     };
   }
 
-  appendOptimisticMessage(parentId: string | null, message: ThreadMessageLike) {
-    let optimisticId: string;
-    do {
-      optimisticId = generateOptimisticId();
-    } while (this.messages.has(optimisticId));
-
-    this.addOrUpdateMessage(
-      parentId,
-      fromThreadMessageLike(message, optimisticId, { type: "running" }),
-    );
-
-    return optimisticId;
-  }
-
   deleteMessage(messageId: string, replacementId?: string | null | undefined) {
     const message = this.messages.get(messageId);
 
@@ -322,6 +308,45 @@ export class MessageRepository {
     return children;
   }
 
+  /**
+   * Evicts optimistic messages (`metadata.isOptimistic`) the head just moved
+   * away from. Since eviction runs on every head move, the only optimistic
+   * messages in the repository live on the branch the head previously pointed
+   * at — so we walk just that branch rather than the whole repository. Keeps a
+   * client→server id swap from leaving a phantom sibling, and drops off-branch
+   * placeholders.
+   */
+  private evictOffBranchOptimisticMessages(
+    previousHead: RepositoryMessage | null,
+    currentHead: RepositoryMessage | null,
+  ) {
+    if (!previousHead) return;
+
+    const onHeadBranch = new Set<string>();
+    for (let current = currentHead; current; current = current.prev) {
+      onHeadBranch.add(current.current.id);
+    }
+
+    const stale: string[] = [];
+    for (
+      let current: RepositoryMessage | null = previousHead;
+      current;
+      current = current.prev
+    ) {
+      // Stop at the first node shared with the current head branch: every
+      // ancestor above it is shared too, so nothing further can be off-branch.
+      if (onHeadBranch.has(current.current.id)) break;
+      if (current.current.metadata?.isOptimistic) {
+        stale.push(current.current.id);
+      }
+    }
+
+    for (const id of stale) {
+      // A prior deletion may have already removed this node.
+      if (this.messages.has(id)) this.deleteMessage(id);
+    }
+  }
+
   switchToBranch(messageId: string) {
     const message = this.messages.get(messageId);
     if (!message)
@@ -329,11 +354,14 @@ export class MessageRepository {
         "MessageRepository(switchToBranch): Branch not found. This is likely an internal bug in assistant-ui.",
       );
 
+    const previousHead = this.head;
     const prevOrRoot = message.prev ?? this.root;
     prevOrRoot.next = message;
 
     this.head = findHead(message);
 
+    this.evictOffBranchOptimisticMessages(previousHead, this.head);
+
     this._messages.dirty();
   }
 
@@ -349,6 +377,8 @@ export class MessageRepository {
         "MessageRepository(resetHead): Branch not found. This is likely an internal bug in assistant-ui.",
       );
 
+    const previousHead = this.head;
+
     if (message.children.length > 0) {
       const deleteDescendants = (msg: RepositoryMessage) => {
         for (const childId of msg.children) {
@@ -378,6 +408,8 @@ export class MessageRepository {
       }
     }
 
+    this.evictOffBranchOptimisticMessages(previousHead, this.head);
+
     this._messages.dirty();
   }
 
@@ -394,15 +426,24 @@ export class MessageRepository {
   export(): ExportedMessageRepository {
     const exportItems: ExportedMessageRepository["messages"] = [];
 
+    // Optimistic messages are ephemeral and never persisted. They're always
+    // leaf nodes, so skipping them can't orphan a persisted child.
     for (const [, message] of this.messages) {
+      if (message.current.metadata?.isOptimistic) continue;
       exportItems.push({
         message: message.current,
         parentId: message.prev?.current.id ?? null,
       });
     }
 
+    // The head may itself be optimistic; walk up to the nearest persisted ancestor.
+    let head = this.head;
+    while (head?.current.metadata?.isOptimistic) {
+      head = head.prev;
+    }
+
     return {
-      headId: this.head?.current.id ?? null,
+      headId: head?.current.id ?? null,
       messages: exportItems,
     };
   }

package/src/runtime/utils/thread-message-like.ts

@@ -81,6 +81,7 @@ export type ThreadMessageLike = {
         readonly steps?: readonly ThreadStep[] | undefined;
         readonly timing?: MessageTiming | undefined;
         readonly submittedFeedback?: { readonly type: "positive" | "negative" };
+        readonly isOptimistic?: boolean | undefined;
         readonly custom?: Record<string, unknown> | undefined;
       }
     | undefined;
@@ -210,6 +211,7 @@ export const fromThreadMessageLike = (
           ...(metadata?.submittedFeedback && {
             submittedFeedback: metadata.submittedFeedback,
           }),
+          ...(metadata?.isOptimistic && { isOptimistic: true }),
         },
       } satisfies ThreadAssistantMessage;
 

package/src/runtimes/external-store/external-store-adapter.ts

@@ -16,6 +16,7 @@ import type {
 } from "../../runtime/interfaces/thread-runtime-core";
 import type { ExportedMessageRepository } from "../../runtime/utils/message-repository";
 import type { ReadonlyJSONValue } from "assistant-stream/utils";
+import type { ToolExecutionStatus } from "../tool-invocations/ToolInvocationTracker";
 
 export type ExternalStoreThreadData<TState extends "regular" | "archived"> = {
   status: TState;
@@ -131,6 +132,38 @@ type ExternalStoreAdapterBase<T> = {
         copy?: boolean | undefined;
       }
     | undefined;
+  /**
+   * Opt in to the built-in client-side tool-invocations pipeline
+   * (`streamCall` / `execute` / tool-status tracking) for this thread.
+   *
+   * Defaults to `false` — the runtime does *not* drive client-side tool
+   * callbacks on its own. Set to `true` to have the runtime construct a
+   * `ToolInvocationTracker` and feed every snapshot through it, so tool
+   * callbacks fire automatically for tool-call parts in `messages`.
+   *
+   * Opt-in by default because most external-store runtimes either run
+   * tools entirely server-side, or already wire their own client-side
+   * dispatch path. Enabling the embedded tracker on top of an existing
+   * dispatch path would cause tool callbacks to run twice.
+   *
+   * When enabled, client-side tool results (from `execute()` returning,
+   * or from `streamCall` resolving) flow back through
+   * `adapter.onAddToolResult` like any other tool result, with
+   * `modelContent` populated when present.
+   */
+  unstable_enableToolInvocations?: boolean | undefined;
+  /**
+   * Receives the current per-tool-call execution status map whenever it
+   * changes. Only invoked when `unstable_enableToolInvocations` is `true`
+   * — the runtime maintains the map via the embedded tracker.
+   *
+   * Wire this into local React state and feed it into the converter's
+   * `metadata.toolStatuses` so the UI can render `executing` spinners
+   * and human-input prompts.
+   */
+  setToolStatuses?:
+    | ((statuses: Record<string, ToolExecutionStatus>) => void)
+    | undefined;
 };
 
 export type ExternalStoreAdapter<T = ThreadMessage> =

package/src/runtimes/external-store/external-store-shared-options.ts

@@ -0,0 +1,18 @@
+import type { ExternalStoreAdapter } from "./external-store-adapter";
+
+export type ExternalStoreSharedOptions = Pick<
+  ExternalStoreAdapter,
+  "isDisabled" | "isSendDisabled" | "unstable_capabilities" | "suggestions"
+>;
+
+export const pickExternalStoreSharedOptions = (
+  options: ExternalStoreSharedOptions,
+): ExternalStoreSharedOptions =>
+  ({
+    isDisabled: options.isDisabled,
+    isSendDisabled: options.isSendDisabled,
+    unstable_capabilities: options.unstable_capabilities,
+    suggestions: options.suggestions,
+  }) satisfies {
+    [K in keyof Required<ExternalStoreSharedOptions>]: ExternalStoreSharedOptions[K];
+  };

package/src/runtimes/external-store/external-store-thread-list-runtime-core.ts

@@ -22,9 +22,7 @@ const DEFAULT_THREAD_DATA = Object.freeze({
   [DEFAULT_THREAD_ID]: DEFAULT_THREAD,
 });
 
-export class ExternalStoreThreadListRuntimeCore
-  implements ThreadListRuntimeCore
-{
+export class ExternalStoreThreadListRuntimeCore implements ThreadListRuntimeCore {
   private _mainThreadId: string = DEFAULT_THREAD_ID;
   private _threads: readonly string[] = DEFAULT_THREADS;
   private _archivedThreads: readonly string[] = EMPTY_ARRAY;

package/src/runtimes/external-store/external-store-thread-runtime-core.ts

@@ -30,6 +30,8 @@ import {
   ExportedMessageRepository,
   MessageRepository,
 } from "../../runtime/utils/message-repository";
+import { generateId } from "../../utils/id";
+import { ToolInvocationTracker } from "../tool-invocations/ToolInvocationTracker";
 
 const EMPTY_ARRAY: readonly ThreadSuggestion[] = Object.freeze([]);
 
@@ -53,9 +55,6 @@ export class ExternalStoreThreadRuntimeCore
   extends BaseThreadRuntimeCore
   implements ThreadRuntimeCore
 {
-  private _assistantOptimisticId: string | null = null;
-  private _lastSyncedMessageIds = new Set<string>();
-
   private _capabilities: RuntimeCapabilities = {
     switchToBranch: false,
     switchBranchDuringRun: false,
@@ -105,6 +104,12 @@ export class ExternalStoreThreadRuntimeCore
 
   private _store!: ExternalStoreAdapter<any>;
 
+  /**
+   * Client-side tool-invocations pipeline. Constructed lazily on first
+   * snapshot — only when `adapter.unstable_enableToolInvocations === true`.
+   */
+  private _toolInvocations: ToolInvocationTracker | null = null;
+
   public override beginEdit(messageId: string) {
     if (!this._store.onEdit)
       throw new Error("Runtime does not support editing.");
@@ -169,10 +174,7 @@ export class ExternalStoreThreadRuntimeCore
         return;
       }
 
-      // Clear and import the message repository
       this.repository.clear();
-      this._assistantOptimisticId = null;
-      this._lastSyncedMessageIds = new Set();
       this.repository.import(store.messageRepository);
 
       messages = this.repository.getMessages();
@@ -225,12 +227,6 @@ export class ExternalStoreThreadRuntimeCore
             return newMessage;
           });
 
-      const nextIds = new Set(messages.map((m) => m.id));
-      for (const prevId of this._lastSyncedMessageIds) {
-        if (!nextIds.has(prevId)) this.repository.deleteMessage(prevId);
-      }
-      this._lastSyncedMessageIds = nextIds;
-
       for (let i = 0; i < messages.length; i++) {
         const message = messages[i]!;
         const parent = messages[i - 1];
@@ -253,29 +249,132 @@ export class ExternalStoreThreadRuntimeCore
       }
     }
 
-    if (this._assistantOptimisticId) {
-      this.repository.deleteMessage(this._assistantOptimisticId);
-      this._assistantOptimisticId = null;
-    }
-
+    // Append an optimistic placeholder while running but before a trailing
+    // assistant message exists. resetHead evicts off-branch optimistic messages
+    // (prior placeholders, mid-run id-swap siblings); export() never persists them.
+    let optimisticId: string | null = null;
     if (hasUpcomingMessage(isRunning, messages)) {
-      this._assistantOptimisticId = this.repository.appendOptimisticMessage(
+      optimisticId = generateId();
+      this.repository.addOrUpdateMessage(
         messages.at(-1)?.id ?? null,
-        {
-          role: "assistant",
-          content: [],
-        },
+        fromThreadMessageLike(
+          { role: "assistant", content: [], metadata: { isOptimistic: true } },
+          optimisticId,
+          { type: "running" },
+        ),
       );
     }
 
-    this.repository.resetHead(
-      this._assistantOptimisticId ?? messages.at(-1)?.id ?? null,
-    );
+    this.repository.resetHead(optimisticId ?? messages.at(-1)?.id ?? null);
 
     this._messages = this.repository.getMessages();
+
+    this._driveToolInvocations();
+
     this._notifySubscribers();
   }
 
+  /**
+   * Feed the current message snapshot into the tool-invocations tracker.
+   * Opt-in via `adapter.unstable_enableToolInvocations: true`. The tracker
+   * itself is fail-silent — see ToolInvocationTracker for the
+   * state-transition contract.
+   */
+  private _driveToolInvocations(): void {
+    if (!this._store.unstable_enableToolInvocations) {
+      // Adapter did not opt in (default). If a tracker was previously
+      // constructed (e.g. the adapter just toggled the flag off via a
+      // dynamic swap), drop it so subsequent snapshots are no-ops.
+      if (this._toolInvocations) {
+        this._toolInvocations.reset();
+        this._toolInvocations = null;
+        this._store.setToolStatuses?.({});
+      }
+      return;
+    }
+
+    if (!this._toolInvocations) {
+      this._toolInvocations = new ToolInvocationTracker(
+        () => this.getModelContext().tools,
+        {
+          onResult: (command) => {
+            try {
+              const messageId = this._findMessageIdForToolCall(
+                command.toolCallId,
+              );
+              if (messageId === undefined) {
+                // The tool call no longer exists in the snapshot (e.g.
+                // rolled back). Drop the result.
+                return;
+              }
+              this._store.onAddToolResult?.({
+                messageId,
+                toolCallId: command.toolCallId,
+                toolName: command.toolName,
+                result: command.result,
+                isError: command.isError,
+                ...(command.artifact !== undefined && {
+                  artifact: command.artifact,
+                }),
+                ...(command.modelContent !== undefined && {
+                  modelContent: command.modelContent,
+                }),
+              });
+            } catch (err) {
+              console.error(
+                "[ExternalStoreThreadRuntimeCore] onAddToolResult dispatch failed",
+                err,
+              );
+            }
+          },
+          onStatusesChange: (statuses) => {
+            this._store.setToolStatuses?.(Object.fromEntries(statuses));
+          },
+        },
+      );
+    }
+
+    this._toolInvocations.setState({
+      messages: this._messages,
+      isRunning: this._store.isRunning ?? false,
+      ...(this._store.isLoading !== undefined && {
+        isLoading: this._store.isLoading,
+      }),
+    });
+  }
+
+  /**
+   * Lookup table from `toolCallId` to the owning assistant message's `id`,
+   * rebuilt lazily when `_messages` changes (see `_messagesForToolCallIndex`).
+   */
+  private _toolCallToMessageId = new Map<string, string>();
+  private _messagesForToolCallIndex: readonly ThreadMessage[] | null = null;
+
+  /**
+   * Look up the assistant message that owns a tool-call part. Lazily builds
+   * (and caches) a `toolCallId → messageId` map keyed off the current
+   * `_messages` reference, so onResult dispatches stay O(1) instead of
+   * walking the full thread on every result.
+   */
+  private _findMessageIdForToolCall(toolCallId: string): string | undefined {
+    if (this._messagesForToolCallIndex !== this._messages) {
+      this._toolCallToMessageId.clear();
+      const visit = (messages: readonly ThreadMessage[]): void => {
+        for (const message of messages) {
+          if (!Array.isArray(message.content)) continue;
+          for (const part of message.content) {
+            if (!part || part.type !== "tool-call") continue;
+            this._toolCallToMessageId.set(part.toolCallId, message.id);
+            if (part.messages) visit(part.messages);
+          }
+        }
+      };
+      visit(this._messages);
+      this._messagesForToolCallIndex = this._messages;
+    }
+    return this._toolCallToMessageId.get(toolCallId);
+  }
+
   public override switchToBranch(branchId: string): void {
     if (!this._store.setMessages)
       throw new Error("Runtime does not support switching branches.");
@@ -290,6 +389,16 @@ export class ExternalStoreThreadRuntimeCore
   }
 
   public async append(message: AppendMessage): Promise<void> {
+    // Auto-abort in-flight client-side tool executions when a new run is
+    // about to start. Without this, a tool that finishes after the new turn
+    // begins would feed a stale result into `onAddToolResult`, racing with
+    // the new turn the user just initiated. `startRun` defaults to true for
+    // user messages — matches the satellites' historical opt-in cancel
+    // behavior, which is now built in.
+    if (message.startRun ?? message.role === "user") {
+      await this._toolInvocations?.abort();
+    }
+
     if (message.parentId !== (this.messages.at(-1)?.id ?? null)) {
       if (!this._store.onEdit)
         throw new Error("Runtime does not support editing messages.");
@@ -303,6 +412,11 @@ export class ExternalStoreThreadRuntimeCore
     if (!this._store.onReload)
       throw new Error("Runtime does not support reloading messages.");
 
+    // Auto-abort in-flight client-side tool executions when a run reloads;
+    // any results that land afterward would target a turn that no longer
+    // exists. See `append` above for full rationale.
+    await this._toolInvocations?.abort();
+
     await this._store.onReload(config.parentId, config);
   }
 
@@ -324,6 +438,18 @@ export class ExternalStoreThreadRuntimeCore
     if (!this._store.onLoadExternalState)
       throw new Error("Runtime does not support importing external states.");
 
+    // Re-arm the tracker so the next adapter snapshot (containing the
+    // imported state) is treated as historical — no streamCall/execute
+    // fires for the loaded tool calls. The adapter is expected to update
+    // its messages in response to onLoadExternalState; that update flows
+    // back here via __internal_setAdapter. We only clear adapter-side
+    // tool statuses when the tracker is the source of truth — otherwise
+    // we'd wipe statuses the adapter is managing on its own.
+    if (this._toolInvocations) {
+      this._toolInvocations.reset();
+      this._store.setToolStatuses?.({});
+    }
+
     this._store.onLoadExternalState(state);
   }
 
@@ -331,11 +457,18 @@ export class ExternalStoreThreadRuntimeCore
     if (!this._store.onCancel)
       throw new Error("Runtime does not support cancelling runs.");
 
+    // Abort any in-flight client-side tool executions. Fire-and-forget —
+    // the abort resolves once executions settle, but we don't gate the
+    // cancel on it.
+    void this._toolInvocations?.abort();
+
     this._store.onCancel();
 
-    if (this._assistantOptimisticId) {
-      this.repository.deleteMessage(this._assistantOptimisticId);
-      this._assistantOptimisticId = null;
+    // Drop an empty optimistic head (placeholder or pre-stream message); a
+    // partially-streamed one is kept and re-supplied by the store on resync.
+    const head = this.repository.getMessages().at(-1);
+    if (head && head.metadata.isOptimistic && head.content.length === 0) {
+      this.repository.deleteMessage(head.id);
     }
 
     let messages = this.repository.getMessages();
@@ -345,7 +478,6 @@ export class ExternalStoreThreadRuntimeCore
       previousMessage.id === messages.at(-1)?.id // ensure the previous message is a leaf node
     ) {
       this.repository.deleteMessage(previousMessage.id);
-      this._lastSyncedMessageIds.delete(previousMessage.id);
       if (!this.composer.text.trim()) {
         this.composer.setText(getThreadMessageText(previousMessage));
       }
@@ -362,15 +494,29 @@ export class ExternalStoreThreadRuntimeCore
   }
 
   public addToolResult(options: AddToolResultOptions) {
-    if (!this._store.onAddToolResult && !this._store.onAddToolResult)
+    if (!this._store.onAddToolResult)
       throw new Error("Runtime does not support tool results.");
     this._store.onAddToolResult?.(options);
   }
 
   public resumeToolCall(options: ResumeToolCallOptions) {
-    if (!this._store.onResumeToolCall)
-      throw new Error("Runtime does not support resuming tool calls.");
-    this._store.onResumeToolCall(options);
+    // Tracker owns its own human-input handlers — let it resume in-process
+    // tool calls without round-tripping through the adapter. Falls back to
+    // the adapter's onResumeToolCall (if any) for tool calls the tracker
+    // doesn't know about.
+    const handled =
+      this._toolInvocations?.resume(options.toolCallId, options.payload) ??
+      false;
+    if (handled) return;
+
+    if (this._store.onResumeToolCall) {
+      this._store.onResumeToolCall(options);
+      return;
+    }
+
+    throw new Error(
+      `Tool call ${options.toolCallId} is not waiting for resume.`,
+    );
   }
 
   public respondToToolApproval(options: RespondToToolApprovalOptions) {
@@ -380,16 +526,12 @@ export class ExternalStoreThreadRuntimeCore
   }
 
   public override reset(initialMessages?: readonly ThreadMessageLike[]) {
-    this._lastSyncedMessageIds = new Set();
     const repo = new MessageRepository();
     repo.import(ExportedMessageRepository.fromArray(initialMessages ?? []));
     this.updateMessages(repo.getMessages());
   }
 
   public override import(data: ExportedMessageRepository) {
-    this._assistantOptimisticId = null;
-    this._lastSyncedMessageIds = new Set();
-
     super.import(data);
 
     if (this._store.onImport) {