@assistant-ui/core 0.3.16 → 0.3.17

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

@@ -49,6 +49,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 +119,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 +133,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 +316,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 +385,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);
+      }
     },
   };
 };
@@ -632,7 +638,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 +649,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,

package/src/store/clients/suggestions.test.ts

@@ -1,9 +1,68 @@
 import { useState } from "react";
 import { createTapRoot, flushTapSync, useResource } from "@assistant-ui/tap";
 import { describe, expect, it } from "vitest";
-import { Suggestions, type SuggestionConfig } from "./suggestions";
+import type { ThreadSuggestion } from "../../runtime/interfaces/thread-runtime-core";
+import {
+  Suggestions,
+  ThreadSuggestions,
+  type SuggestionConfig,
+} from "./suggestions";
 
 describe("Suggestions", () => {
+  it("keeps equal inline string suggestions stable across owner renders", () => {
+    let rerender!: () => void;
+    const root = createTapRoot(function SuggestionsRoot() {
+      const [, setTick] = useState(0);
+      rerender = () => setTick((tick) => tick + 1);
+      return useResource(Suggestions(["account-a", "account-b"]));
+    });
+
+    try {
+      const firstState = root.getValue().getState();
+      const firstSuggestion = root
+        .getValue()
+        .suggestion({ index: 0 })
+        .getState();
+
+      flushTapSync(rerender);
+
+      expect(root.getValue().getState()).toBe(firstState);
+      expect(root.getValue().suggestion({ index: 0 }).getState()).toBe(
+        firstSuggestion,
+      );
+    } finally {
+      root.unmount();
+    }
+  });
+
+  it("keeps equal inline object suggestions stable across owner renders", () => {
+    let rerender!: () => void;
+    const root = createTapRoot(function SuggestionsRoot() {
+      const [, setTick] = useState(0);
+      rerender = () => setTick((tick) => tick + 1);
+      return useResource(
+        Suggestions([{ title: "Title", label: "Label", prompt: "Prompt" }]),
+      );
+    });
+
+    try {
+      const firstState = root.getValue().getState();
+      const firstSuggestion = root
+        .getValue()
+        .suggestion({ index: 0 })
+        .getState();
+
+      flushTapSync(rerender);
+
+      expect(root.getValue().getState()).toBe(firstState);
+      expect(root.getValue().suggestion({ index: 0 }).getState()).toBe(
+        firstSuggestion,
+      );
+    } finally {
+      root.unmount();
+    }
+  });
+
   it("updates when the configured suggestions change", () => {
     let setSuggestions!: (suggestions: SuggestionConfig[]) => void;
     const root = createTapRoot(function SuggestionsRoot() {
@@ -28,4 +87,99 @@ describe("Suggestions", () => {
       root.unmount();
     }
   });
+
+  it("reuses unchanged suggestion state when one entry changes", () => {
+    let setSuggestions!: (suggestions: SuggestionConfig[]) => void;
+    const root = createTapRoot(function SuggestionsRoot() {
+      const [suggestions, setValue] = useState<SuggestionConfig[]>([
+        { title: "Title A", label: "Label A", prompt: "Prompt A" },
+        { title: "Title B", label: "Label B", prompt: "Prompt B" },
+      ]);
+      setSuggestions = setValue;
+      return useResource(Suggestions(suggestions));
+    });
+
+    try {
+      const firstState = root.getValue().getState();
+      const unchangedSuggestion = root
+        .getValue()
+        .suggestion({ index: 1 })
+        .getState();
+
+      flushTapSync(() =>
+        setSuggestions([
+          { title: "Changed", label: "Label A", prompt: "Prompt A" },
+          { title: "Title B", label: "Label B", prompt: "Prompt B" },
+        ]),
+      );
+
+      expect(root.getValue().getState()).not.toBe(firstState);
+      expect(root.getValue().getState().suggestions[0]?.title).toBe("Changed");
+      expect(root.getValue().suggestion({ index: 1 }).getState()).toBe(
+        unchangedSuggestion,
+      );
+    } finally {
+      root.unmount();
+    }
+  });
+});
+
+describe("ThreadSuggestions", () => {
+  it("keeps equal runtime suggestions stable across owner renders", () => {
+    let rerender!: () => void;
+    const root = createTapRoot(function ThreadSuggestionsRoot() {
+      const [, setTick] = useState(0);
+      rerender = () => setTick((tick) => tick + 1);
+      return useResource(ThreadSuggestions([{ prompt: "Prompt" }]));
+    });
+
+    try {
+      const firstState = root.getValue().getState();
+      const firstSuggestion = root
+        .getValue()
+        .suggestion({ index: 0 })
+        .getState();
+
+      flushTapSync(rerender);
+
+      expect(root.getValue().getState()).toBe(firstState);
+      expect(root.getValue().suggestion({ index: 0 }).getState()).toBe(
+        firstSuggestion,
+      );
+    } finally {
+      root.unmount();
+    }
+  });
+
+  it("reuses unchanged suggestion state when one entry changes", () => {
+    let setSuggestions!: (suggestions: ThreadSuggestion[]) => void;
+    const root = createTapRoot(function ThreadSuggestionsRoot() {
+      const [suggestions, setValue] = useState<ThreadSuggestion[]>([
+        { prompt: "Prompt A" },
+        { prompt: "Prompt B" },
+      ]);
+      setSuggestions = setValue;
+      return useResource(ThreadSuggestions(suggestions));
+    });
+
+    try {
+      const firstState = root.getValue().getState();
+      const unchangedSuggestion = root
+        .getValue()
+        .suggestion({ index: 1 })
+        .getState();
+
+      flushTapSync(() =>
+        setSuggestions([{ prompt: "Changed" }, { prompt: "Prompt B" }]),
+      );
+
+      expect(root.getValue().getState()).not.toBe(firstState);
+      expect(root.getValue().getState().suggestions[0]?.prompt).toBe("Changed");
+      expect(root.getValue().suggestion({ index: 1 }).getState()).toBe(
+        unchangedSuggestion,
+      );
+    } finally {
+      root.unmount();
+    }
+  });
 });