@assistant-ui/core 0.3.16 → 0.3.18

package/src/runtime/utils/tool-call-tree.ts

@@ -0,0 +1,97 @@
+import type {
+  ThreadAssistantMessage,
+  ThreadAssistantMessagePart,
+  ThreadMessage,
+  ToolCallMessagePart,
+} from "../../types/message";
+
+export type ToolCallTreeEntry = {
+  readonly part: ToolCallMessagePart;
+  /** Id of the message whose content directly holds `part`. */
+  readonly messageId: string;
+};
+
+/**
+ * Walk every tool-call part reachable from `messages`. A tool call projects a
+ * child run as nested messages on `ToolCallMessagePart.messages`, so a
+ * subagent's own calls live below the top-level content and a seam that only
+ * scans the top level cannot reach them.
+ *
+ * Parts arrive in document order, each one ahead of its own descendants.
+ * `messageId` names the message that directly holds the part, which for a
+ * nested call is the child run's message rather than the top-level message the
+ * tree hangs from. Only assistant messages are descended, the same rule
+ * {@link mapToolCallPartsDeep} rewrites under, so a part this reports is always
+ * a part that can be written back.
+ */
+export function* walkToolCallTree(
+  messages: readonly ThreadMessage[],
+): Generator<ToolCallTreeEntry> {
+  for (const message of messages) {
+    if (message?.role !== "assistant" || !Array.isArray(message.content)) {
+      continue;
+    }
+    for (const part of message.content) {
+      if (!part || part.type !== "tool-call") continue;
+      yield { part, messageId: message.id };
+      if (part.messages?.length) yield* walkToolCallTree(part.messages);
+    }
+  }
+}
+
+/**
+ * Content-level form of {@link walkToolCallTree}, for a caller that already
+ * holds one assistant message's content and does not need the owning id.
+ */
+export function* iterateToolCallParts(
+  content: readonly ThreadAssistantMessagePart[],
+): Generator<ToolCallMessagePart> {
+  for (const part of content) {
+    if (!part || part.type !== "tool-call") continue;
+    yield part;
+    if (part.messages?.length) {
+      for (const entry of walkToolCallTree(part.messages)) {
+        yield entry.part;
+      }
+    }
+  }
+}
+
+/**
+ * Rebuild `content` with `fn` applied to every tool-call part in the tree.
+ * Arrays and nested messages are reused wherever `fn` returned the part it was
+ * given, so an unchanged subtree keeps its identity and `changed` reports
+ * whether anything moved.
+ */
+export function mapToolCallPartsDeep(
+  content: readonly ThreadAssistantMessagePart[],
+  fn: (part: ToolCallMessagePart) => ToolCallMessagePart,
+): { content: readonly ThreadAssistantMessagePart[]; changed: boolean } {
+  let changed = false;
+  const next = content.map((part): ThreadAssistantMessagePart => {
+    if (part.type !== "tool-call") return part;
+    let mapped = fn(part);
+    if (mapped.messages !== undefined) {
+      let nestedChanged = false;
+      const nestedMessages = mapped.messages.map((nested) => {
+        if (nested.role !== "assistant" || !Array.isArray(nested.content)) {
+          return nested;
+        }
+        const assistant = nested as ThreadAssistantMessage;
+        const result = mapToolCallPartsDeep(assistant.content, fn);
+        if (!result.changed) return nested;
+        nestedChanged = true;
+        return { ...assistant, content: result.content };
+      });
+      if (nestedChanged) {
+        mapped =
+          mapped === part
+            ? { ...part, messages: nestedMessages }
+            : { ...mapped, messages: nestedMessages };
+      }
+    }
+    if (mapped !== part) changed = true;
+    return mapped;
+  });
+  return changed ? { content: next, changed } : { content, changed };
+}

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

@@ -1,4 +1,8 @@
-import type { AppendMessage, ThreadMessage } from "../../types/message";
+import type {
+  AppendMessage,
+  ThreadMessage,
+  ToolCallMessagePart,
+} from "../../types/message";
 import type { ThreadMessageLike } from "../../runtime/utils/thread-message-like";
 import type { AttachmentAdapter } from "../../adapters/attachment";
 import type {
@@ -14,7 +18,10 @@ import type {
   ResumeRunConfig,
   ThreadSuggestion,
 } from "../../runtime/interfaces/thread-runtime-core";
-import type { ExportedMessageRepository } from "../../runtime/utils/message-repository";
+import type {
+  ExportedMessageRepository,
+  MessageRepository,
+} from "../../runtime/utils/message-repository";
 import type { ReadonlyJSONValue } from "assistant-stream/utils";
 import type { ToolExecutionStatus } from "../tool-invocations/ToolInvocationTracker";
 import type { ExternalThreadQueueAdapter } from "../../runtime/queue/external-thread-queue-adapter";
@@ -105,6 +112,15 @@ type ExternalStoreAdapterBase<T> = {
   isLoading?: boolean | undefined;
   messages?: readonly T[];
   messageRepository?: ExportedMessageRepository;
+  /**
+   * An externally owned message repository instance. When provided, the
+   * thread runtime adopts it as its branch store and swaps to it atomically
+   * whenever a different instance is passed, so hosts that route multiple
+   * conversations through one runtime keep each conversation's history and
+   * branches isolated in its own instance. Omit it to keep the runtime's own
+   * repository.
+   */
+  unstable_messageRepositoryInstance?: MessageRepository | undefined;
   suggestions?: readonly ThreadSuggestion[] | undefined;
   state?: ReadonlyJSONValue | undefined;
   extras?: unknown;
@@ -207,6 +223,26 @@ type ExternalStoreAdapterBase<T> = {
    * `modelContent` populated when present.
    */
   unstable_enableToolInvocations?: boolean | undefined;
+  /**
+   * Decides whether a tool call's result is produced on the client. Only
+   * consulted when `unstable_enableToolInvocations` is `true`.
+   *
+   * A provider that runs tools itself answers its own calls, and its result
+   * arrives one or more snapshots after the call's arguments complete. In
+   * that window the call is complete and result-less, so a registered tool
+   * of the same name would otherwise execute locally and produce a result
+   * the provider never asked for. An adapter that can tell the two apart
+   * supplies this predicate; it is read once per tool call, when the call is
+   * first observed live.
+   *
+   * The predicate is also what licenses running a frontend tool while the
+   * provider's run is still open. Without it, ownership is unknown until the
+   * run ends, so a registered tool executes only once the run's outcome is
+   * known and cannot fire on a call the provider was about to answer or gate.
+   */
+  unstable_isClientToolCall?:
+    | ((toolCall: ToolCallMessagePart) => boolean)
+    | undefined;
   /**
    * Receives the current per-tool-call execution status map whenever it
    * changes. Only invoked when `unstable_enableToolInvocations` is `true`

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

@@ -6,6 +6,7 @@ import type {
 import type { ExternalStoreThreadListAdapter } from "./external-store-adapter";
 import { invalidateThreadRuntime } from "../../runtime/utils/thread-runtime-lifecycle";
 import { BaseSubscribable } from "../../subscribable/subscribable";
+import { nullProtoRecord } from "../../utils/record";
 
 export type ExternalStoreThreadFactory = () => ExternalStoreThreadRuntimeCore;
 
@@ -19,9 +20,11 @@ const DEFAULT_THREAD = Object.freeze({
   status: "regular",
 });
 const RESOLVED_PROMISE = Promise.resolve();
-const DEFAULT_THREAD_DATA = Object.freeze({
-  [DEFAULT_THREAD_ID]: DEFAULT_THREAD,
-});
+const DEFAULT_THREAD_DATA = Object.freeze(
+  nullProtoRecord<ThreadListItemCoreState>({
+    [DEFAULT_THREAD_ID]: DEFAULT_THREAD,
+  }),
+);
 
 export class ExternalStoreThreadListRuntimeCore
   extends BaseSubscribable
@@ -84,7 +87,9 @@ export class ExternalStoreThreadListRuntimeCore
   }
 
   public getItemById(threadId: string) {
-    return this._threadData[threadId];
+    return Object.hasOwn(this._threadData, threadId)
+      ? this._threadData[threadId]
+      : undefined;
   }
 
   public __internal_setAdapter(
@@ -105,6 +110,7 @@ export class ExternalStoreThreadListRuntimeCore
 
     if (
       !initialLoad &&
+      (previousAdapter.isLoading ?? false) === (adapter.isLoading ?? false) &&
       previousThreadId === newThreadId &&
       previousThreads === newThreads &&
       previousArchivedThreads === newArchivedThreads
@@ -117,9 +123,9 @@ export class ExternalStoreThreadListRuntimeCore
       previousArchivedThreads !== newArchivedThreads ||
       previousThreadId !== newThreadId
     ) {
-      this._threadData = {
-        ...DEFAULT_THREAD_DATA,
-        ...Object.fromEntries(
+      this._threadData = nullProtoRecord(
+        DEFAULT_THREAD_DATA,
+        Object.fromEntries(
           adapter.threads?.map((t) => [
             t.id,
             {
@@ -130,7 +136,7 @@ export class ExternalStoreThreadListRuntimeCore
             },
           ]) ?? [],
         ),
-        ...Object.fromEntries(
+        Object.fromEntries(
           adapter.archivedThreads?.map((t) => [
             t.id,
             {
@@ -141,7 +147,7 @@ export class ExternalStoreThreadListRuntimeCore
             },
           ]) ?? [],
         ),
-      };
+      );
     }
 
     if (previousThreads !== newThreads) {
@@ -160,16 +166,15 @@ export class ExternalStoreThreadListRuntimeCore
       this._mainThread = this.threadFactory();
     }
 
-    if (!this._threadData[this._mainThreadId]) {
-      this._threadData = {
-        ...this._threadData,
+    if (!Object.hasOwn(this._threadData, this._mainThreadId)) {
+      this._threadData = nullProtoRecord(this._threadData, {
         [this._mainThreadId]: {
           id: this._mainThreadId,
           remoteId: undefined,
           externalId: undefined,
           status: "regular",
         },
-      };
+      });
     }
 
     this._notifySubscribers();

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

@@ -1,3 +1,4 @@
+import { shallowEqual } from "@assistant-ui/store/client";
 import type { AppendMessage, ThreadMessage } from "../../types/message";
 import type { Attachment } from "../../types/attachment";
 import type {
@@ -28,6 +29,7 @@ import {
   type ThreadMessageLike,
 } from "../../runtime/utils/thread-message-like";
 import { getThreadMessageText } from "../../utils/text";
+import { shallowArrayEqual } from "../../runtime/utils/external-message-conversion";
 import type {
   RuntimeCapabilities,
   ThreadRuntimeCore,
@@ -43,6 +45,7 @@ import {
   MessageRepository,
 } from "../../runtime/utils/message-repository";
 import { generateId } from "../../utils/id";
+import { walkToolCallTree } from "../../runtime/utils/tool-call-tree";
 import {
   ToolInvocationTracker,
   type ToolExecutionStatus,
@@ -57,7 +60,7 @@ import {
 const EMPTY_ARRAY: readonly ThreadSuggestion[] = Object.freeze([]);
 
 const observeAdapterCallback = (
-  name: "onAddToolResult" | "onRespondToToolApproval" | "onCancel",
+  name: "onAddToolResult" | "onCancel",
   result: Promise<void> | void,
 ) => {
   void Promise.resolve(result).catch((error) => {
@@ -68,15 +71,6 @@ const observeAdapterCallback = (
   });
 };
 
-const shallowEqual = (a: object, b: object): boolean => {
-  const aKeys = Object.keys(a);
-  if (aKeys.length !== Object.keys(b).length) return false;
-  for (const key of aKeys) {
-    if ((a as any)[key] !== (b as any)[key]) return false;
-  }
-  return true;
-};
-
 export const hasUpcomingMessage = (
   isRunning: boolean,
   messages: readonly ThreadMessage[],
@@ -154,6 +148,11 @@ export class ExternalStoreThreadRuntimeCore
   // and a delete whose confirmation races a send keeps its eviction.
   private _pendingDeleteEvictions = new Set<string>();
 
+  // Placeholder id for the upcoming assistant message, reused across snapshot
+  // passes while the same tail message awaits its response so the placeholder
+  // keeps one identity per response.
+  private _optimistic: { id: string; parentId: string | null } | null = null;
+
   private _store!: ExternalStoreAdapter<any>;
 
   private _getInitializePromise?: () => Promise<unknown> | undefined;
@@ -246,6 +245,14 @@ export class ExternalStoreThreadRuntimeCore
     const oldStore = this._store as ExternalStoreAdapter<any> | undefined;
     this._store = store;
     const isRunning = this._getEffectiveIsRunning(store);
+    const repositoryInstance = store.unstable_messageRepositoryInstance;
+    const repositoryChanged =
+      repositoryInstance !== undefined &&
+      repositoryInstance !== this.repository;
+    if (repositoryChanged) {
+      this.repository = repositoryInstance;
+      this._pendingDeleteEvictions.clear();
+    }
     if (oldStore?.queue !== store.queue) {
       this._transformedQueue = undefined;
       store.queue?.__internal_setDispatchTransform?.((message) => {
@@ -296,6 +303,7 @@ export class ExternalStoreThreadRuntimeCore
       // Handle messageRepository
       if (
         oldStore &&
+        !repositoryChanged &&
         oldStore.isRunning === store.isRunning &&
         oldStore.messageRepository === store.messageRepository &&
         previousIsRunning === isRunning
@@ -308,7 +316,11 @@ export class ExternalStoreThreadRuntimeCore
       const headId =
         store.messageRepository.headId ?? incoming.at(-1)?.message.id ?? null;
 
-      if (oldStore && oldStore.messageRepository === store.messageRepository) {
+      if (
+        oldStore &&
+        !repositoryChanged &&
+        oldStore.messageRepository === store.messageRepository
+      ) {
         this.repository.resetHead(headId);
         messages = this.repository.getMessages();
       } else {
@@ -333,6 +345,7 @@ export class ExternalStoreThreadRuntimeCore
         if (oldStore.convertMessage !== store.convertMessage) {
           this._converter = new ThreadMessageConverter();
         } else if (
+          !repositoryChanged &&
           oldStore.isRunning === store.isRunning &&
           oldStore.messages === store.messages &&
           previousIsRunning === isRunning
@@ -436,9 +449,13 @@ export class ExternalStoreThreadRuntimeCore
     // (prior placeholders, mid-run id-swap siblings); export() never persists them.
     let optimisticId: string | null = null;
     if (hasUpcomingMessage(isRunning, messages)) {
-      optimisticId = generateId();
+      const parentId = messages.at(-1)?.id ?? null;
+      if (this._optimistic?.parentId !== parentId) {
+        this._optimistic = { id: generateId(), parentId };
+      }
+      optimisticId = this._optimistic.id;
       this.repository.addOrUpdateMessage(
-        messages.at(-1)?.id ?? null,
+        parentId,
         fromThreadMessageLike(
           { role: "assistant", content: [], metadata: { isOptimistic: true } },
           optimisticId,
@@ -447,10 +464,20 @@ export class ExternalStoreThreadRuntimeCore
       );
     }
 
+    if (optimisticId === null) this._optimistic = null;
     this.repository.resetHead(optimisticId ?? messages.at(-1)?.id ?? null);
 
-    this._messages = this.repository.getMessages();
+    const messagesSnapshot = this.repository.getMessages();
+    if (
+      !this._messages ||
+      !shallowArrayEqual(this._messages, messagesSnapshot)
+    ) {
+      this._messages = messagesSnapshot;
+    }
 
+    if (repositoryChanged) {
+      this._runTrackerUpdate(() => this._toolInvocations?.reset());
+    }
     this._runTrackerUpdate(() => this._driveToolInvocations());
 
     this._notifySubscribers();
@@ -529,6 +556,7 @@ export class ExternalStoreThreadRuntimeCore
             }
           },
         },
+        (toolCall) => this._store.unstable_isClientToolCall?.(toolCall),
       );
     }
 
@@ -557,17 +585,9 @@ export class ExternalStoreThreadRuntimeCore
   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);
+      for (const { part, messageId } of walkToolCallTree(this._messages)) {
+        this._toolCallToMessageId.set(part.toolCallId, messageId);
+      }
       this._messagesForToolCallIndex = this._messages;
     }
     return this._toolCallToMessageId.get(toolCallId);
@@ -671,7 +691,7 @@ export class ExternalStoreThreadRuntimeCore
     // 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();
+      await this._toolInvocations?.abort({ discardPending: true });
     }
     if (!isThreadRuntimeGenerationCurrent(this, generation)) return;
 
@@ -772,7 +792,7 @@ export class ExternalStoreThreadRuntimeCore
     // 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._toolInvocations?.abort({ discardPending: true });
 
     await this._store.onReload(config.parentId, config);
   }
@@ -826,7 +846,7 @@ export class ExternalStoreThreadRuntimeCore
     // 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();
+    void this._toolInvocations?.abort({ discardPending: true });
 
     // Before the run is aborted, so the settle it produces keeps the pending
     // items instead of dispatching the next one at the moment the user
@@ -936,13 +956,16 @@ export class ExternalStoreThreadRuntimeCore
     );
   }
 
-  public respondToToolApproval(options: RespondToToolApprovalOptions) {
+  public respondToToolApproval(
+    options: RespondToToolApprovalOptions,
+  ): Promise<void> {
     if (!this._store.onRespondToToolApproval)
       throw new Error("Runtime does not support tool approvals.");
-    observeAdapterCallback(
-      "onRespondToToolApproval",
-      this._store.onRespondToToolApproval(options),
-    );
+    try {
+      return Promise.resolve(this._store.onRespondToToolApproval(options));
+    } catch (error) {
+      return Promise.reject(error);
+    }
   }
 
   public override reset(initialMessages?: readonly ThreadMessageLike[]) {

package/src/runtimes/inert/InertThreadRuntimeCore.ts

@@ -136,7 +136,7 @@ export abstract class InertThreadRuntimeCore
     throw this.error;
   }
 
-  respondToToolApproval(): void {
+  respondToToolApproval(): Promise<void> {
     throw this.error;
   }
 

package/src/runtimes/internal.ts

@@ -31,6 +31,7 @@ export type {
 export {
   createThreadMappingId,
   getThreadData,
+  promoteNewThreadReducer,
   updateStatusReducer,
 } from "./remote-thread-list/remote-thread-state";
 export type {

package/src/runtimes/local/local-thread-runtime-core.test.ts

@@ -5,7 +5,7 @@ import type {
   ChatModelRunOptions,
   ChatModelRunResult,
 } from "../../runtime/utils/chat-model-adapter";
-import type { AppendMessage } from "../../types/message";
+import type { AppendMessage, ToolCallMessagePart } from "../../types/message";
 import type { LocalRuntimeOptionsBase } from "./local-runtime-options";
 import {
   ExportedMessageRepository,
@@ -61,7 +61,7 @@ const userMessage = (text: string): AppendMessage => ({
 
 const toolCallPart = (
   toolName: string,
-  approval?: { id: string; resolution?: "cancelled" | "expired" },
+  approval?: ToolCallMessagePart["approval"],
 ) => ({
   type: "tool-call" as const,
   toolCallId: `call-${toolName}`,
@@ -73,7 +73,7 @@ const toolCallPart = (
 
 const toolCallResult = (
   toolName: string,
-  approval?: { id: string; resolution?: "cancelled" | "expired" },
+  approval?: ToolCallMessagePart["approval"],
 ): ChatModelRunResult => ({
   content: [toolCallPart(toolName, approval)],
   status: { type: "requires-action", reason: "tool-calls" },
@@ -485,6 +485,32 @@ describe("LocalThreadRuntimeCore tool approvals", () => {
     });
   });
 
+  it("records a free-form answer alongside the decision", async () => {
+    const { thread, runs } = createApprovalThread(
+      toolCallResult("send_email", { id: "a1", display: "text" }),
+    );
+
+    await thread.append(userMessage("send an email"));
+    await flush();
+
+    await thread.respondToToolApproval({
+      approvalId: "a1",
+      approved: true,
+      text: "Quarterly update",
+    });
+    await flush();
+
+    const toolCall = runs[1]!
+      .unstable_getMessage()
+      .content.find((part) => part.type === "tool-call");
+    expect(toolCall?.approval).toEqual({
+      id: "a1",
+      display: "text",
+      approved: true,
+      text: "Quarterly update",
+    });
+  });
+
   it("treats a terminal resolution as non-pending and continues the run", async () => {
     const { thread, runs } = createApprovalThread(
       toolCallResult("deploy", { id: "a1", resolution: "expired" }),

package/src/runtimes/local/local-thread-runtime-core.ts

@@ -900,8 +900,9 @@ export class LocalThreadRuntimeCore
     approvalId,
     approved,
     optionId,
+    text,
     reason,
-  }: RespondToToolApprovalOptions) {
+  }: RespondToToolApprovalOptions): Promise<void> {
     let message = this.repository
       .getMessages()
       .findLast(
@@ -944,6 +945,7 @@ export class LocalThreadRuntimeCore
         ...targetApproval,
         approved,
         ...(optionId != null && { optionId }),
+        ...(text != null && { text }),
         ...(reason != null && { reason }),
       };
       if (approved) return { ...c, approval };
@@ -968,5 +970,7 @@ export class LocalThreadRuntimeCore
     } else {
       this._persistPausedMessage(parentId, message);
     }
+
+    return Promise.resolve();
   }
 }