@assistant-ui/core 0.2.5 → 0.2.7

package/src/react/utils/groupParts.ts

@@ -1,3 +1,6 @@
+import { isMcpAppUri } from "../../types/message";
+import type { PartState } from "../../store/scopes/part";
+
 /**
  * Hierarchical adjacent-coalescing grouping for message parts.
  *
@@ -11,16 +14,71 @@
  */
 
 /**
- * Public group key type. Group keys must be prefixed with `group-` so
- * that a unified `switch (part.type)` in the renderer can distinguish
- * a group key (e.g. `"group-thought"`) from a real part type
- * (`"text"`, `"tool-call"`).
+ * Symbol attached to memoizable `groupBy` functions (e.g. those returned
+ * by {@link groupPartByType}). Carries a string fingerprint of the config
+ * so `MessagePrimitive.GroupedParts` can memo the tree on
+ * `[parts, memoKey]` across renders — even when the helper call site
+ * reconstructs the function each render.
+ */
+export const GROUPBY_MEMO_KEY: unique symbol = Symbol.for(
+  "@assistant-ui/groupBy.memoKey",
+);
+
+/**
+ * Synthetic part-type key recognized by {@link groupPartByType}: a
+ * tool-call whose `mcp.app.resourceUri` points at an assistant-ui MCP
+ * app. Map this key to control how MCP-app tool calls are grouped —
+ * separately from regular `"tool-call"` parts.
+ */
+type GroupPartType = PartState["type"] | "mcp-app";
+
+/**
+ * Build a `groupBy` from a `part.type → group-key path` lookup.
+ * Parts whose type isn't in the map are left ungrouped. The returned
+ * function carries a stable {@link GROUPBY_MEMO_KEY} fingerprint so
+ * `<MessagePrimitive.GroupedParts>` can memoize its tree across renders.
+ *
+ * Special key `"mcp-app"` matches tool-call parts that point at an
+ * assistant-ui MCP app resource (`ui://...`) and takes precedence over
+ * the `"tool-call"` entry for those parts.
+ *
+ * @example
+ * ```tsx
+ * <MessagePrimitive.GroupedParts
+ *   groupBy={groupPartByType({
+ *     reasoning: ["group-thought", "group-reasoning"],
+ *     "tool-call": ["group-thought", "group-tool"],
+ *     "mcp-app": [],
+ *   })}
+ * >
+ *   {({ part, children }) => { ... }}
+ * </MessagePrimitive.GroupedParts>
+ * ```
  */
-export type GroupKey<TKey extends `group-${string}` = `group-${string}`> =
-  | TKey
-  | readonly TKey[]
-  | null
-  | undefined;
+export const groupPartByType = <TKey extends `group-${string}`>(
+  map: Partial<Readonly<Record<GroupPartType, readonly TKey[]>>>,
+): ((part: PartState) => readonly TKey[]) => {
+  const lookup = map as Readonly<Record<string, readonly TKey[] | undefined>>;
+  const fn = ((part) => {
+    if (
+      part.type === "tool-call" &&
+      lookup["mcp-app"] !== undefined &&
+      isMcpAppUri(part.mcp?.app?.resourceUri)
+    ) {
+      return lookup["mcp-app"]!;
+    }
+    return lookup[part.type] ?? [];
+  }) as ((part: PartState) => readonly TKey[]) & {
+    [GROUPBY_MEMO_KEY]?: string;
+  };
+  // Sort keys so the fingerprint is insensitive to map insertion order —
+  // two maps with the same key/value pairs but different declaration order
+  // would otherwise hash differently and invalidate the memo unnecessarily.
+  const sortedKeys = Object.keys(map).sort();
+  const sortedEntries = sortedKeys.map((k) => [k, map[k as keyof typeof map]]);
+  fn[GROUPBY_MEMO_KEY] = `groupPartByType:${JSON.stringify(sortedEntries)}`;
+  return fn;
+};
 
 export type GroupNode = GroupNodeGroup | GroupNodePart;
 
@@ -43,19 +101,6 @@ export interface GroupNodePart {
   readonly nodeKey: string;
 }
 
-const EMPTY_PATH: readonly string[] = Object.freeze([]);
-
-/**
- * Normalize a `groupBy` return value to a path array.
- * `null`/`undefined`/`[]` → `[]` (ungrouped).
- * `"foo"` → `["foo"]`. Arrays pass through.
- */
-export const normalizeGroupKey = (key: GroupKey): readonly string[] => {
-  if (key == null) return EMPTY_PATH;
-  if (typeof key === "string") return [key];
-  return key;
-};
-
 interface BuildFrame {
   key: string;
   nodeKey: string;

package/src/runtime/api/attachment-runtime.ts

@@ -42,8 +42,7 @@ export type AttachmentRuntime<
 
 export abstract class AttachmentRuntimeImpl<
   Source extends AttachmentRuntimeSource = AttachmentRuntimeSource,
-> implements AttachmentRuntime
-{
+> implements AttachmentRuntime {
   public get path() {
     return this._core.path;
   }

package/src/runtime/api/message-part-runtime.ts

@@ -26,6 +26,7 @@ type MessagePartSnapshotBinding = SubscribableWithState<
 export type MessagePartRuntime = {
   addToolResult(result: any | ToolResponse<any>): void;
   resumeToolCall(payload: unknown): void;
+  respondToToolApproval(response: { approved: boolean; reason?: string }): void;
 
   readonly path: MessagePartRuntimePath;
   getState(): MessagePartState;
@@ -48,6 +49,7 @@ export class MessagePartRuntimeImpl implements MessagePartRuntime {
   protected __internal_bindMethods() {
     this.addToolResult = this.addToolResult.bind(this);
     this.resumeToolCall = this.resumeToolCall.bind(this);
+    this.respondToToolApproval = this.respondToToolApproval.bind(this);
     this.getState = this.getState.bind(this);
     this.subscribe = this.subscribe.bind(this);
   }
@@ -102,6 +104,30 @@ export class MessagePartRuntimeImpl implements MessagePartRuntime {
     });
   }
 
+  public respondToToolApproval(response: {
+    approved: boolean;
+    reason?: string;
+  }) {
+    const state = this.contentBinding.getState();
+    if (!state) throw new Error("Message part is not available");
+
+    if (state.type !== "tool-call")
+      throw new Error(
+        "Tried to respond to tool approval on non-tool message part",
+      );
+
+    if (!state.approval || state.approval.approved !== undefined)
+      throw new Error("Tool call has no pending approval");
+
+    if (!this.threadApi) throw new Error("Thread API is not available");
+
+    this.threadApi.getState().respondToToolApproval({
+      approvalId: state.approval.id,
+      approved: response.approved,
+      ...(response.reason != null && { reason: response.reason }),
+    });
+  }
+
   public subscribe(callback: () => void) {
     return this.contentBinding.subscribe(callback);
   }

package/src/runtime/base/base-thread-runtime-core.ts

@@ -15,6 +15,7 @@ import { DefaultThreadComposerRuntimeCore } from "./default-thread-composer-runt
 import type {
   AddToolResultOptions,
   ResumeToolCallOptions,
+  RespondToToolApprovalOptions,
   ThreadSuggestion,
   SubmitFeedbackOptions,
   ThreadRuntimeCore,
@@ -59,6 +60,9 @@ export abstract class BaseThreadRuntimeCore implements ThreadRuntimeCore {
   public abstract resumeRun(config: ResumeRunConfig): void;
   public abstract addToolResult(options: AddToolResultOptions): void;
   public abstract resumeToolCall(options: ResumeToolCallOptions): void;
+  public abstract respondToToolApproval(
+    options: RespondToToolApprovalOptions,
+  ): void;
   public abstract cancelRun(): void;
   public abstract exportExternalState(): any;
   public abstract importExternalState(state: any): void;

package/src/runtime/interfaces/thread-runtime-core.ts

@@ -1,3 +1,4 @@
+import type { ToolModelContentPart } from "assistant-stream";
 import type { ReadonlyJSONValue } from "assistant-stream/utils";
 import type { ModelContext } from "../../model-context/types";
 import type { Unsubscribe } from "../../types/unsubscribe";
@@ -38,6 +39,13 @@ export type AddToolResultOptions = {
   result: ReadonlyJSONValue;
   isError: boolean;
   artifact?: ReadonlyJSONValue | undefined;
+  /**
+   * Optional model-content payload produced by the tool. Populated when a
+   * client-side `execute()` or `streamCall` returns a `ToolResponse` with
+   * `modelContent`. Forwarded through `adapter.onAddToolResult` so the
+   * adapter can include it when sending the result back to its backend.
+   */
+  modelContent?: readonly ToolModelContentPart[] | undefined;
 };
 
 export type ResumeToolCallOptions = {
@@ -45,6 +53,12 @@ export type ResumeToolCallOptions = {
   payload: unknown;
 };
 
+export type RespondToToolApprovalOptions = {
+  approvalId: string;
+  approved: boolean;
+  reason?: string;
+};
+
 export type SubmitFeedbackOptions = {
   messageId: string;
   type: "negative" | "positive";
@@ -137,6 +151,7 @@ export type ThreadRuntimeCore = Readonly<{
 
   addToolResult: (options: AddToolResultOptions) => void;
   resumeToolCall: (options: ResumeToolCallOptions) => void;
+  respondToToolApproval: (options: RespondToToolApprovalOptions) => void;
 
   speak: (messageId: string) => void;
   stopSpeaking: () => void;

package/src/runtime/internal.ts

@@ -18,10 +18,7 @@ export { DefaultEditComposerRuntimeCore } from "./base/default-edit-composer-run
 // Runtime Impl Classes
 export { AssistantRuntimeImpl } from "./api/assistant-runtime";
 
-export {
-  getThreadState,
-  ThreadRuntimeImpl,
-} from "./api/thread-runtime";
+export { getThreadState, ThreadRuntimeImpl } from "./api/thread-runtime";
 export type {
   ThreadRuntimeCoreBinding,
   ThreadListItemRuntimeBinding,

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

@@ -54,6 +54,13 @@ export type ThreadMessageLike = {
             readonly isError?: boolean | undefined;
             readonly parentId?: string | undefined;
             readonly messages?: readonly ThreadMessage[] | undefined;
+            readonly interrupt?: { type: "human"; payload: unknown };
+            readonly approval?: {
+              readonly id: string;
+              readonly approved?: boolean;
+              readonly reason?: string;
+              readonly isAutomatic?: boolean;
+            };
           }
       )[];
   readonly id?: string | undefined;

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

@@ -9,12 +9,14 @@ import type { RealtimeVoiceAdapter } from "../../adapters/voice";
 import type { FeedbackAdapter } from "../../adapters/feedback";
 import type {
   AddToolResultOptions,
+  RespondToToolApprovalOptions,
   StartRunConfig,
   ResumeRunConfig,
   ThreadSuggestion,
 } 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;
@@ -108,6 +110,9 @@ type ExternalStoreAdapterBase<T> = {
   onResumeToolCall?:
     | ((options: { toolCallId: string; payload: unknown }) => void)
     | undefined;
+  onRespondToToolApproval?:
+    | ((options: RespondToToolApprovalOptions) => Promise<void> | void)
+    | undefined;
   convertMessage?: ExternalStoreMessageConverter<T> | undefined;
   adapters?:
     | {
@@ -127,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-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

@@ -3,6 +3,7 @@ import type {
   AddToolResultOptions,
   ResumeRunConfig,
   ResumeToolCallOptions,
+  RespondToToolApprovalOptions,
   StartRunConfig,
   ThreadSuggestion,
 } from "../../runtime/interfaces/thread-runtime-core";
@@ -29,6 +30,7 @@ import {
   ExportedMessageRepository,
   MessageRepository,
 } from "../../runtime/utils/message-repository";
+import { ToolInvocationTracker } from "../tool-invocations/ToolInvocationTracker";
 
 const EMPTY_ARRAY: readonly ThreadSuggestion[] = Object.freeze([]);
 
@@ -104,6 +106,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.");
@@ -272,9 +280,113 @@ export class ExternalStoreThreadRuntimeCore
     );
 
     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.");
@@ -289,6 +401,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.");
@@ -302,6 +424,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);
   }
 
@@ -323,6 +450,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);
   }
 
@@ -330,6 +469,11 @@ 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) {
@@ -361,15 +505,35 @@ 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) {
+    if (!this._store.onRespondToToolApproval)
+      throw new Error("Runtime does not support tool approvals.");
+    this._store.onRespondToToolApproval(options);
   }
 
   public override reset(initialMessages?: readonly ThreadMessageLike[]) {

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

@@ -9,6 +9,7 @@ import type { LocalRuntimeOptionsBase } from "./local-runtime-options";
 import type {
   AddToolResultOptions,
   ResumeToolCallOptions,
+  RespondToToolApprovalOptions,
   ThreadSuggestion,
   ThreadRuntimeCore,
   StartRunConfig,
@@ -535,4 +536,8 @@ export class LocalThreadRuntimeCore
   public resumeToolCall(_options: ResumeToolCallOptions) {
     throw new Error("Local runtime does not support resuming tool calls.");
   }
+
+  public respondToToolApproval(_options: RespondToToolApprovalOptions) {
+    throw new Error("Local runtime does not support tool approvals.");
+  }
 }

package/src/runtimes/readonly/ReadonlyThreadRuntimeCore.ts

@@ -65,6 +65,10 @@ export class ReadonlyThreadRuntimeCore
     throw READONLY_THREAD_ERROR;
   }
 
+  respondToToolApproval(): void {
+    throw READONLY_THREAD_ERROR;
+  }
+
   speak(): void {
     throw READONLY_THREAD_ERROR;
   }

package/src/runtimes/remote-thread-list/empty-thread-core.ts

@@ -40,6 +40,10 @@ export const EMPTY_THREAD_CORE: ThreadRuntimeCore = {
     throw EMPTY_THREAD_ERROR;
   },
 
+  respondToToolApproval() {
+    throw EMPTY_THREAD_ERROR;
+  },
+
   speak() {
     throw EMPTY_THREAD_ERROR;
   },