@assistant-ui/core 0.3.16 → 0.3.18

package/src/runtimes/tool-invocations/ToolInvocationTracker.ts

@@ -14,7 +14,8 @@ import {
   type ReadonlyJSONValue,
 } from "assistant-stream/utils";
 import { isJSONValueEqual } from "../../utils/json/is-json-equal";
-import type { ThreadMessage } from "../../types/message";
+import type { ThreadMessage, ToolCallMessagePart } from "../../types/message";
+import { walkToolCallTree } from "../../runtime/utils/tool-call-tree";
 
 const TOOL_EXECUTION_ID = Symbol.for("assistant-stream.tool-execution-id");
 
@@ -54,6 +55,7 @@ type ToolCallEntry = {
       /** Active phase — chunks are flowing through `controller`. */
       controller: ToolCallStreamController;
       argsComplete: boolean;
+      clientOwned: boolean;
     }
 );
 
@@ -116,6 +118,9 @@ const getToolExecutionId = (value: object): symbol | undefined =>
 export class ToolInvocationTracker {
   private readonly _getTools: () => Record<string, Tool> | undefined;
   private readonly _callbacks: ToolInvocationTracker.Callbacks;
+  private readonly _isClientToolCall:
+    | ((toolCall: ToolCallMessagePart) => boolean | undefined)
+    | undefined;
 
   private readonly _entries = new Map<string, ToolCallEntry>();
   private readonly _humanInput = new Map<
@@ -127,6 +132,13 @@ export class ToolInvocationTracker {
     }
   >();
   private readonly _executing = new Set<symbol>();
+  /**
+   * Tool calls whose turn ended before they reached the executor. Held here
+   * rather than on the entry because an entry is rebuilt whenever a snapshot
+   * re-creates the call, and this is the one reason to skip that no later
+   * snapshot carries.
+   */
+  private readonly _discardedToolCallIds = new Set<string>();
   private readonly _settledResolvers: SettledResolver[] = [];
 
   private _statuses = new Map<string, ToolExecutionStatus>();
@@ -143,10 +155,12 @@ export class ToolInvocationTracker {
   /**
    * Set when the assistant-stream pipeline has died (errored out via
    * `.pipeTo(...).catch(...)`). The next `setState` re-initializes the
-   * pipeline and demotes all active entries to restored so they survive
-   * across the restart without re-firing `streamCall` (preserves the
-   * "exactly once" contract). Capped at a single auto-restart per session
-   * — repeated failures keep the tracker dead with a more visible error.
+   * pipeline and demotes each active entry that reached the executor to
+   * restored, so it survives the restart without re-firing `streamCall`.
+   * A restart is an execution boundary like `reset()`: an entry that had
+   * not reached it starts over and fires once there (F.4). Capped at a
+   * single auto-restart per session — repeated failures keep the tracker
+   * dead with a more visible error.
    */
   private _pipelineDead = false;
   private _pipelineRestartUsed = false;
@@ -154,9 +168,11 @@ export class ToolInvocationTracker {
   constructor(
     getTools: () => Record<string, Tool> | undefined,
     callbacks: ToolInvocationTracker.Callbacks,
+    isClientToolCall?: (toolCall: ToolCallMessagePart) => boolean | undefined,
   ) {
     this._getTools = getTools;
     this._callbacks = callbacks;
+    this._isClientToolCall = isClientToolCall;
 
     this._initPipeline();
   }
@@ -222,10 +238,10 @@ export class ToolInvocationTracker {
    */
   public setState(snapshot: ToolInvocationTracker.Snapshot): void {
     try {
-      // Recover from a dead pipeline before processing anything. We demote
-      // all active entries to "restored" so the rebuilt pipeline does not
-      // re-fire `streamCall` for tool calls that already fired pre-death;
-      // preserves the "exactly once per toolCallId" contract.
+      // Recover from a dead pipeline before processing anything. Entries
+      // that reached the executor are demoted to "restored" so the rebuilt
+      // pipeline does not re-fire `streamCall` for them; the rest start over
+      // across the boundary the restart opens (F.4).
       if (this._pipelineDead) {
         if (this._pipelineRestartUsed) {
           // Already retried once and failed again. Stay dead.
@@ -291,6 +307,7 @@ export class ToolInvocationTracker {
     try {
       this._pendingRestore = true;
       this._entries.clear();
+      this._discardedToolCallIds.clear();
       this._lastSnapshot = null;
       void this.abort();
       // Statuses are cleared synchronously: discarded executions may never
@@ -310,8 +327,13 @@ export class ToolInvocationTracker {
   /**
    * Abort any in-flight `execute()` invocations. Resolves once all of them
    * have settled (or immediately if none are running).
+   *
+   * `discardPending` additionally kills the calls that never reached the
+   * executor, for a caller ending the turn rather than interrupting it. The
+   * signal cannot reach those: they are waiting on the run to settle (A.10),
+   * and the settled snapshot arrives after this installs a fresh controller.
    */
-  public abort(): Promise<void> {
+  public abort(options?: { discardPending?: boolean }): Promise<void> {
     try {
       this._humanInput.forEach(({ reject }) => {
         try {
@@ -323,6 +345,15 @@ export class ToolInvocationTracker {
       });
       this._humanInput.clear();
 
+      if (options?.discardPending) {
+        for (const [toolCallId, entry] of this._entries) {
+          if (!entry.controller) continue;
+          if (entry.argsComplete || entry.hasResult) continue;
+          this._discardedToolCallIds.add(toolCallId);
+          entry.skipExecute = true;
+        }
+      }
+
       this._ac.abort();
       this._ac = new AbortController();
 
@@ -473,6 +504,8 @@ export class ToolInvocationTracker {
     if (entry.skipExecute) return;
 
     this._executing.add(executionId!);
+    // execute can park human() before onExecutionStart; preserve this execution's interrupt.
+    if (this._humanInput.get(toolCallId)?.executionId === executionId) return;
     this._setStatus(toolCallId, { type: "executing" });
   }
 
@@ -577,31 +610,40 @@ export class ToolInvocationTracker {
 
   // ──────────────── internal: snapshot processing ────────────────
 
-  private _hasExecutableTool(toolName: string): boolean {
-    const tool = this._getTools()?.[toolName];
-    return tool?.execute !== undefined || tool?.streamCall !== undefined;
+  private _warnProviderOwnedSkip(toolName: string, toolCallId: string): void {
+    if (process.env.NODE_ENV === "production") return;
+    if (this._getTools()?.[toolName]?.execute === undefined) return;
+    console.warn(
+      "[ToolInvocationTracker] the runtime reports this tool call as provider-owned, so the registered execute is skipped; the provider has to hand the call to the client for it to run (see EDGE_CASES.md A.9)",
+      { toolCallId, toolName },
+    );
   }
 
+  /**
+   * Closing the args stream hands the call to the client executor, so it may
+   * only happen once the provider can no longer speak about that call. The run
+   * ending is the only universal signal for that; an adapter that reports the
+   * call as client-owned has said it earlier, per call.
+   */
   private _shouldCloseArgsStream({
-    toolName,
     argsText,
     hasResult,
+    clientOwned,
   }: {
-    toolName: string;
     argsText: string;
     hasResult: boolean;
+    clientOwned: boolean;
   }): boolean {
     if (hasResult) return true;
-    if (!this._hasExecutableTool(toolName)) {
-      return !this._isRunning && isArgsTextComplete(argsText);
-    }
-    return isArgsTextComplete(argsText);
+    if (!isArgsTextComplete(argsText)) return false;
+    return clientOwned || !this._isRunning;
   }
 
   private _startActiveEntry(
     toolCallId: string,
     toolName: string,
     skipExecute: boolean,
+    clientOwned: boolean,
   ): ToolCallEntry {
     const toolCallController = this._controller.addToolCallPart({
       toolName,
@@ -614,6 +656,7 @@ export class ToolInvocationTracker {
       hasResult: false,
       skipExecute,
       argsComplete: false,
+      clientOwned,
     };
     this._entries.set(toolCallId, entry);
     return entry;
@@ -629,6 +672,16 @@ export class ToolInvocationTracker {
   private _demoteEntriesToRestored(): void {
     for (const [toolCallId, entry] of this._entries) {
       if (!entry.controller) continue;
+      if (!entry.argsComplete && !entry.hasResult) {
+        // The call never reached the executor. A restored entry is promoted
+        // only when its signature changes, and a call waiting on the run to
+        // settle already holds its final args, so demoting it would strand it
+        // unexecuted. Dropping it lets the next snapshot start it over; a call
+        // whose turn was discarded is held by `_discardedToolCallIds`, not by
+        // the entry, so starting over does not revive it.
+        this._entries.delete(toolCallId);
+        continue;
+      }
       this._entries.set(toolCallId, {
         toolName: entry.toolName,
         argsText: entry.argsText,
@@ -683,9 +736,9 @@ export class ToolInvocationTracker {
           isEquivalentCompleteArgsText(entry.argsText, content.argsText)
         ) {
           const shouldClose = this._shouldCloseArgsStream({
-            toolName: content.toolName,
             argsText: content.argsText,
             hasResult,
+            clientOwned: entry.clientOwned,
           });
           if (shouldClose) entry.controller.argsText.close();
           entry.argsText = content.argsText;
@@ -717,9 +770,9 @@ export class ToolInvocationTracker {
         const delta = content.argsText.slice(entry.argsText.length);
         entry.controller.argsText.append(delta);
         const shouldClose = this._shouldCloseArgsStream({
-          toolName: content.toolName,
           argsText: content.argsText,
           hasResult,
+          clientOwned: entry.clientOwned,
         });
         if (shouldClose) entry.controller.argsText.close();
         entry.argsText = content.argsText;
@@ -732,9 +785,9 @@ export class ToolInvocationTracker {
       // gates on the streamed content; a divergent snapshot (A.2) can be
       // complete while the controller still holds an incomplete stale prefix.
       const shouldClose = this._shouldCloseArgsStream({
-        toolName: content.toolName,
         argsText: entry.argsText,
         hasResult,
+        clientOwned: entry.clientOwned,
       });
       if (shouldClose) {
         entry.controller.argsText.close();
@@ -746,81 +799,81 @@ export class ToolInvocationTracker {
   private _processMessages(messages: readonly ThreadMessage[]): void {
     const isRestore = this._pendingRestore;
 
-    for (const message of messages) {
-      if (!message || !Array.isArray((message as ThreadMessage).content)) {
+    for (const { part: content } of walkToolCallTree(messages)) {
+      const existing = this._entries.get(content.toolCallId);
+
+      if (isRestore) {
+        // Don't overwrite an already-active entry (e.g. live tool-call
+        // observed before this restore snapshot landed). Restore can
+        // only seed entries the runtime has never seen.
+        if (!existing?.controller) {
+          this._entries.set(content.toolCallId, {
+            toolName: content.toolName,
+            argsText: content.argsText,
+            hasResult: content.result !== undefined,
+          });
+        }
         continue;
       }
-      for (const content of message.content as readonly ThreadMessage["content"][number][]) {
-        if (!content || content.type !== "tool-call") continue;
-
-        const existing = this._entries.get(content.toolCallId);
-
-        if (isRestore) {
-          // Don't overwrite an already-active entry (e.g. live tool-call
-          // observed before this restore snapshot landed). Restore can
-          // only seed entries the runtime has never seen.
-          if (!existing?.controller) {
-            this._entries.set(content.toolCallId, {
-              toolName: content.toolName,
-              argsText: content.argsText,
-              hasResult: content.result !== undefined,
-            });
-          }
-          if (content.messages) this._processMessages(content.messages);
-          continue;
-        }
 
-        // Live snapshot.
-        let entry = existing;
-
-        if (entry && !entry.controller) {
-          // Restored entry observed in a live snapshot. Promote if its
-          // signature has changed; otherwise treat as still-historical.
-          const signatureChanged =
-            content.argsText !== entry.argsText ||
-            (content.result !== undefined) !== entry.hasResult;
-          if (!signatureChanged) {
-            if (content.messages) this._processMessages(content.messages);
-            continue;
-          }
-          this._entries.delete(content.toolCallId);
-          entry = undefined;
-        }
+      // Live snapshot.
+      let entry = existing;
+
+      // A discarded id is remembered only until the call is answered, which
+      // bounds the set to the open calls of a discarded turn.
+      if (content.result !== undefined)
+        this._discardedToolCallIds.delete(content.toolCallId);
+
+      if (entry && !entry.controller) {
+        // Restored entry observed in a live snapshot. Promote if its
+        // signature has changed; otherwise treat as still-historical.
+        const signatureChanged =
+          content.argsText !== entry.argsText ||
+          (content.result !== undefined) !== entry.hasResult;
+        if (!signatureChanged) continue;
+        this._entries.delete(content.toolCallId);
+        entry = undefined;
+      }
 
-        if (!entry) {
-          entry = this._startActiveEntry(
-            content.toolCallId,
-            content.toolName,
-            content.result !== undefined,
-          );
-        }
+      if (!entry) {
+        const ownership = this._isClientToolCall?.(content);
+        const providerOwned =
+          content.result === undefined && ownership === false;
+        if (providerOwned)
+          this._warnProviderOwnedSkip(content.toolName, content.toolCallId);
+        entry = this._startActiveEntry(
+          content.toolCallId,
+          content.toolName,
+          content.result !== undefined ||
+            providerOwned ||
+            this._discardedToolCallIds.has(content.toolCallId),
+          ownership === true,
+        );
+      }
 
-        if (content.approval !== undefined) entry.skipExecute = true;
-
-        this._processArgsText(entry, content);
-
-        if (content.result !== undefined && !entry.hasResult) {
-          // `entry` is in active phase from this point — either just
-          // created by `_startActiveEntry`, or pre-existing with a live
-          // controller. Narrow once instead of asserting at every use.
-          const { controller: activeController } = entry;
-          if (!activeController) continue;
-          entry.hasResult = true;
-          entry.argsComplete = true;
-          activeController.setResponse(
-            new ToolResponse({
-              result: content.result as ReadonlyJSONValue,
-              artifact: content.artifact as ReadonlyJSONValue | undefined,
-              isError: content.isError,
-              ...(content.modelContent !== undefined
-                ? { modelContent: content.modelContent }
-                : {}),
-            }),
-          );
-          activeController.close();
-        }
+      if (content.approval !== undefined) entry.skipExecute = true;
 
-        if (content.messages) this._processMessages(content.messages);
+      this._processArgsText(entry, content);
+
+      if (content.result !== undefined && !entry.hasResult) {
+        // `entry` is in active phase from this point — either just
+        // created by `_startActiveEntry`, or pre-existing with a live
+        // controller. Narrow once instead of asserting at every use.
+        const { controller: activeController } = entry;
+        if (!activeController) continue;
+        entry.hasResult = true;
+        entry.argsComplete = true;
+        activeController.setResponse(
+          new ToolResponse({
+            result: content.result as ReadonlyJSONValue,
+            artifact: content.artifact as ReadonlyJSONValue | undefined,
+            isError: content.isError,
+            ...(content.modelContent !== undefined
+              ? { modelContent: content.modelContent }
+              : {}),
+          }),
+        );
+        activeController.close();
       }
     }
   }

package/src/store/clients/external-thread.ts

@@ -21,7 +21,10 @@ import type {
   CreateAttachment,
   PendingAttachment,
 } from "../../types/attachment";
-import { isCreateAttachment } from "../../types/attachment";
+import {
+  isAttachmentComplete,
+  isCreateAttachment,
+} from "../../types/attachment";
 import type {
   AddToolResultOptions,
   RespondToToolApprovalOptions,
@@ -49,6 +52,7 @@ import {
   drainAttachmentAdd,
 } from "../../runtime/utils/attachment-add-operations";
 import { toMessagePartStatus } from "../../utils/normalizePartStatus";
+import { generateId } from "../../utils/id";
 import { ModelContext } from "./model-context-client";
 import { ThreadSuggestions } from "./suggestions";
 import { Tools } from "../../react/client/Tools";
@@ -118,7 +122,9 @@ export type ExternalThreadProps = {
   /** Branch adapter for runtimes that track sibling variants of messages. */
   branches?: ExternalThreadBranchAdapter;
   /** Callback for tool approval decisions. Absent: responding to an approval throws a capability error. */
-  onRespondToToolApproval?: (options: RespondToToolApprovalOptions) => void;
+  onRespondToToolApproval?: (
+    options: RespondToToolApprovalOptions,
+  ) => void | Promise<void>;
 };
 
 type MessageClientProps = {
@@ -130,7 +136,7 @@ type MessageClientProps = {
   queue?: ExternalThreadQueueAdapter | undefined;
   branches?: ExternalThreadBranchAdapter | undefined;
   onRespondToToolApproval?:
-    | ((options: RespondToToolApprovalOptions) => void)
+    | ((options: RespondToToolApprovalOptions) => void | Promise<void>)
     | undefined;
   onAddToolResult?: ((options: AddToolResultOptions) => void) | undefined;
   onResumeToolCall?: ((options: ResumeToolCallOptions) => void) | undefined;
@@ -313,7 +319,7 @@ type PartResourceProps = {
   status: ToolCallMessagePartStatus;
   messageId: string;
   onRespondToToolApproval?:
-    | ((options: RespondToToolApprovalOptions) => void)
+    | ((options: RespondToToolApprovalOptions) => void | Promise<void>)
     | undefined;
   onAddToolResult?: ((options: AddToolResultOptions) => void) | undefined;
   onResumeToolCall?: ((options: ResumeToolCallOptions) => void) | undefined;
@@ -382,9 +388,12 @@ const usePartResource = ({
       )
         throw new Error("Tool call has no pending approval");
 
-      onRespondToToolApproval(
-        resolveToolApprovalResponse(part.approval, response),
-      );
+      const options = resolveToolApprovalResponse(part.approval, response);
+      try {
+        return Promise.resolve(onRespondToToolApproval(options));
+      } catch (error) {
+        return Promise.reject(error);
+      }
     },
   };
 };
@@ -457,6 +466,35 @@ const useLiveState = <T>(initial: T) => {
   return [state, set, ref] as const;
 };
 
+// A try/catch inside the resource makes the React Compiler bail out of the
+// whole hook, so the adapter call and its failure status live at module scope.
+const removeAttachmentThroughAdapter = async (
+  attachment: Attachment,
+  attachmentAdapter: AttachmentAdapter | undefined,
+  setAttachments: (
+    next:
+      | readonly Attachment[]
+      | ((prev: readonly Attachment[]) => readonly Attachment[]),
+  ) => void,
+) => {
+  try {
+    await attachmentAdapter?.remove(attachment);
+  } catch (error) {
+    const message = error instanceof Error ? error.message : String(error);
+    setAttachments((prev) =>
+      prev.map((candidate) =>
+        candidate.id === attachment.id && !isAttachmentComplete(candidate)
+          ? {
+              ...candidate,
+              status: { type: "incomplete", reason: "error", message },
+            }
+          : candidate,
+      ),
+    );
+    throw error;
+  }
+};
+
 // Composer Client - minimal implementation
 const useComposerClientResource = ({
   type,
@@ -510,7 +548,13 @@ const useComposerClientResource = ({
           attachment,
           onRemove: async () => {
             attachmentAddOperations.cancel(attachment.id);
-            await attachmentAdapter?.remove(attachment);
+            if (!isAttachmentComplete(attachment)) {
+              await removeAttachmentThroughAdapter(
+                attachment,
+                attachmentAdapter,
+                setAttachments,
+              );
+            }
             setAttachments((prev) =>
               prev.filter((a) => a.id !== attachment.id),
             );
@@ -632,7 +676,7 @@ const useComposerClientResource = ({
         }
       } else if (!isCreateAttachment(fileOrAttachment)) {
         const newAttachment: Attachment = {
-          id: Math.random().toString(36).substring(7),
+          id: generateId(),
           type: "file",
           name: fileOrAttachment.name,
           contentType: fileOrAttachment.type,
@@ -643,7 +687,7 @@ const useComposerClientResource = ({
         setAttachments((prev) => [...prev, newAttachment]);
       } else {
         const newAttachment: Attachment = {
-          id: fileOrAttachment.id ?? Math.random().toString(36).substring(7),
+          id: fileOrAttachment.id ?? generateId(),
           type: fileOrAttachment.type ?? "document",
           name: fileOrAttachment.name,
           contentType: fileOrAttachment.contentType,
@@ -899,7 +943,11 @@ const useExternalThread = ({
       : undefined;
   };
 
+  // `messages` changes identity on every streamed token, and this body fans out
+  // to one client per message, so pruning during render would re-run the whole
+  // fan-out per token. The effect bails out on an unchanged map instead.
   useEffect(() => {
+    // eslint-disable-next-line react-hooks/set-state-in-effect
     setSubmittedFeedback((prev) => {
       const live = Object.entries(prev).filter(([id, entry]) => {
         const msg = messages.find((m) => m.id === id);
@@ -1107,7 +1155,10 @@ const useExternalThread = ({
             }
           : {
               createdAt: message.createdAt ?? new Date(),
-              parentId: message.parentId ?? messages.at(-1)?.id ?? null,
+              parentId:
+                message.parentId === undefined
+                  ? (messages.at(-1)?.id ?? null)
+                  : message.parentId,
               sourceId: message.sourceId ?? null,
               role: message.role ?? "user",
               content: message.content,

package/src/store/clients/model-context-client.ts

@@ -3,6 +3,7 @@ import { resource } from "@assistant-ui/tap";
 import type { ClientOutput } from "@assistant-ui/store";
 import { CompositeContextProvider } from "../../utils/composite-context-provider";
 import type { ModelContextState } from "../scopes/model-context";
+import { shallowEqual } from "@assistant-ui/store/internal";
 
 const EMPTY_TOOL_NAMES: readonly string[] = [];
 
@@ -12,7 +13,7 @@ const INITIAL_STATE: ModelContextState = {
 };
 
 const toolNamesEqual = (a: readonly string[], b: readonly string[]): boolean =>
-  a === b || (a.length === b.length && a.every((v, i) => v === b[i]));
+  a === b || shallowEqual(a, b);
 
 const deriveState = (
   composite: CompositeContextProvider,
@@ -36,6 +37,9 @@ const useModelContext = (): ClientOutput<"modelContext"> => {
   );
 
   useEffect(() => {
+    // The composite is an external store; this catches the snapshot up for
+    // providers registered between render and subscribe.
+    // eslint-disable-next-line react-hooks/set-state-in-effect
     setState((prev) => deriveState(composite, prev));
     return composite.subscribe(() => {
       setState((prev) => deriveState(composite, prev));

package/src/store/clients/single-thread-list.ts

@@ -59,6 +59,7 @@ const useSingleThreadList = ({
       mainThreadId: THREAD_ID,
       newThreadId: null,
       isLoading: false,
+      loadError: undefined,
       isLoadingMore: false,
       hasMore: false,
       threadIds: [THREAD_ID],