@assistant-ui/core 0.3.13 → 0.3.14

package/src/runtimes/tool-invocations/EDGE_CASES.md

@@ -10,6 +10,9 @@ observe via `setState(snapshot)` and what the current behavior is.
 > first observation — args regress, args change after first completion,
 > result is replaced, result is cleared, key order shuffles — the tracker
 > never invokes the host's tool callback a second time.
+>
+> `reset()` starts a new execution boundary. A reused `toolCallId` in the
+> next boundary gets a new execution identity and may fire once there.
 
 This guarantees host-side side effects (the typical reason `streamCall` /
 `execute` exists at all) can't double-run. The cost: post-completion
@@ -65,7 +68,7 @@ The tracker calls `setResponse` on the active controller and closes it.
 The backend result is emitted before the args stream closes, so a stale
 args parse failure cannot replace it. `reader.response.get()` resolves.
 If the tool also had a frontend `execute`, the executor is short-circuited
-via `_skipExecuteStreamIds`. Single fire.
+via the entry's per-execution skip marker. Single fire.
 
 ### A.6. Previously-resolved tool's `result` is replaced
 Silently ignored — `entry.hasResult` short-circuits both the
@@ -153,15 +156,16 @@ tracker's view of "what we last observed". The next snapshot retries.
 
 ### F.1. `reset()` called while `execute()` invocations are in flight
 `abort()` is invoked, in-flight executions reject with
-`Tool execution aborted`. Once they settle, the cleanup logic clears
-`_executing`. The settled-resolver promises fire so the abort promise
-resolves.
+`Tool execution aborted`. Once they settle, the cleanup logic removes their
+execution identities from `_executing`. The settled-resolver promises fire so
+the abort promise resolves.
 
 ### F.2. `setState` called during `reset()`'s in-flight abort
 The new snapshot is processed against an empty `_entries`. Tool calls
 in it are seeded as restored (because `reset()` re-armed
-`_pendingRestore`). Eventual cancellation `result` chunks for the
-aborted executions are dropped via `_skipExecuteStreamIds`.
+`_pendingRestore`). Eventual cancellation `result` chunks for the aborted
+executions are dropped when their execution identity no longer matches the
+current entry.
 
 ### F.3. `resume(toolCallId, payload)` for an unknown id
 Silently no-ops. (The pre-class hook *threw*; the tracker softens this
@@ -175,6 +179,13 @@ rebuilt pipeline does not re-fire `streamCall` for them) and the
 snapshot is processed against the fresh pipeline. Repeated failures
 keep the tracker dead with a visible error to avoid restart loops.
 
+### F.5. Reset followed by same-id reuse
+The assistant-stream pipeline assigns each streamed tool part an opaque
+execution identity. Human-input callbacks, lifecycle callbacks, result chunks,
+and pending execution cleanup are accepted only when that identity is still
+current for the `toolCallId`. A late callback from before `reset()` is therefore
+dropped even when the next session reuses the id.
+
 ## Known limitations
 
 ### Host callback throws

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

@@ -1172,3 +1172,246 @@ describe("ToolInvocationTracker", () => {
     }
   });
 });
+
+describe("ToolInvocationTracker reset", () => {
+  const trackerWith = (tools: Record<string, Tool>) => {
+    let statuses: Record<string, ToolExecutionStatus> = {};
+    const tracker = new ToolInvocationTracker(() => tools, {
+      onResult: vi.fn(),
+      onStatusesChange: (s: ReadonlyMap<string, ToolExecutionStatus>) => {
+        statuses = Object.fromEntries(s);
+      },
+    });
+    return { tracker, statuses: () => statuses };
+  };
+
+  it("clears statuses for discarded executions that never settle", async () => {
+    const { tracker, statuses } = trackerWith({
+      weatherSearch: {
+        parameters: { type: "object", properties: {} },
+        execute: vi.fn(() => new Promise(() => {})),
+      } satisfies Tool,
+    });
+    tracker.setState(createState([]));
+    tracker.setState(
+      createState([
+        createAssistantMessage('{"query":"London"}', { query: "London" }),
+      ]),
+    );
+    await waitFor(() => {
+      expect(statuses()["tool-1"]?.type).toBe("executing");
+    });
+
+    tracker.reset();
+
+    expect(statuses()).toEqual({});
+
+    // A new-session snapshot restoring tool activity must not drag the
+    // discarded id back in through the whole-map republication. The first
+    // post-reset snapshot is the fresh session's empty state, consuming the
+    // pending-restore mark.
+    tracker.setState(createState([]));
+    tracker.setState(
+      createState([
+        createAssistantMessage(
+          '{"query":"Paris"}',
+          { query: "Paris" },
+          { toolCallId: "tool-2" },
+        ),
+      ]),
+    );
+    await waitFor(() => {
+      expect(statuses()["tool-2"]?.type).toBe("executing");
+    });
+    expect(statuses()["tool-1"]).toBeUndefined();
+  });
+
+  it("does not republish sibling statuses when a discarded execution settles after reset", async () => {
+    const settled = Promise.withResolvers<void>();
+    const { tracker, statuses } = trackerWith({
+      weatherSearch: {
+        parameters: { type: "object", properties: {} },
+        execute: vi.fn(
+          (_args, { abortSignal }: { abortSignal: AbortSignal }) =>
+            new Promise((resolve) => {
+              abortSignal.addEventListener("abort", () => {
+                resolve({ ok: true });
+                settled.resolve();
+              });
+            }),
+        ),
+      } satisfies Tool,
+      pressureSearch: {
+        parameters: { type: "object", properties: {} },
+        execute: vi.fn(() => new Promise(() => {})),
+      } satisfies Tool,
+    });
+    tracker.setState(createState([]));
+    tracker.setState(
+      createState([
+        createAssistantMessage('{"query":"London"}', { query: "London" }),
+        createAssistantMessage(
+          '{"query":"London"}',
+          { query: "London" },
+          {
+            toolCallId: "tool-2",
+            toolName: "pressureSearch",
+          },
+        ),
+      ]),
+    );
+    await waitFor(() => {
+      expect(statuses()["tool-1"]?.type).toBe("executing");
+      expect(statuses()["tool-2"]?.type).toBe("executing");
+    });
+
+    tracker.reset();
+    expect(statuses()).toEqual({});
+
+    // The settle-after-abort of tool-1 is what could republish tool-2;
+    // await it deterministically instead of sleeping.
+    await settled.promise;
+    await new Promise((r) => setTimeout(r, 0));
+    expect(statuses()).toEqual({});
+  });
+
+  it("rejects a human-input request from a discarded execution without resurrecting its status", async () => {
+    let humanFn: ((payload: unknown) => Promise<unknown>) | undefined;
+    const { tracker, statuses } = trackerWith({
+      weatherSearch: {
+        parameters: { type: "object", properties: {} },
+        execute: vi.fn(() => new Promise(() => {})),
+        streamCall: vi.fn((_reader, { human }) => {
+          humanFn = human;
+        }),
+      } satisfies Tool,
+    });
+    tracker.setState(createState([]));
+    tracker.setState(
+      createState([
+        createAssistantMessage('{"query":"London"}', { query: "London" }),
+      ]),
+    );
+    await waitFor(() => {
+      expect(statuses()["tool-1"]?.type).toBe("executing");
+      expect(humanFn).toBeDefined();
+    });
+
+    tracker.reset();
+
+    await expect(humanFn!({ request: "approve" })).rejects.toThrow(
+      "Tool execution aborted",
+    );
+    expect(statuses()).toEqual({});
+  });
+
+  it("rejects a late human-input request when reset reuses the tool call id", async () => {
+    const humanFns: ((payload: unknown) => Promise<unknown>)[] = [];
+    const { tracker, statuses } = trackerWith({
+      weatherSearch: {
+        parameters: { type: "object", properties: {} },
+        execute: vi.fn(() => new Promise(() => {})),
+        streamCall: vi.fn((_reader, { human }) => {
+          humanFns.push(human);
+        }),
+      } satisfies Tool,
+    });
+
+    tracker.setState(createState([]));
+    tracker.setState(
+      createState([
+        createAssistantMessage('{"query":"London"}', { query: "London" }),
+      ]),
+    );
+    await waitFor(() => {
+      expect(humanFns).toHaveLength(1);
+      expect(statuses()["tool-1"]?.type).toBe("executing");
+    });
+
+    tracker.reset();
+    tracker.setState(createState([]));
+    tracker.setState(
+      createState([
+        createAssistantMessage(
+          '{"query":"Paris"}',
+          { query: "Paris" },
+          { toolCallId: "tool-1" },
+        ),
+      ]),
+    );
+    await waitFor(() => expect(humanFns).toHaveLength(2));
+
+    await expect(humanFns[0]!({ stale: true })).rejects.toThrow(
+      "Tool execution aborted",
+    );
+    expect(statuses()["tool-1"]?.type).toBe("executing");
+
+    tracker.reset();
+  });
+
+  it("drops a late result when reset reuses the tool call id", async () => {
+    const executions: Array<{
+      resolve: (value: { session: string }) => void;
+      promise: Promise<{ session: string }>;
+    }> = [];
+    const onResult = vi.fn();
+    const tracker = new ToolInvocationTracker(
+      () => ({
+        weatherSearch: {
+          parameters: { type: "object", properties: {} },
+          execute: vi.fn(
+            (_args: unknown, { abortSignal }: { abortSignal: AbortSignal }) => {
+              let resolveExecution!: (value: { session: string }) => void;
+              const promise = new Promise<{ session: string }>((resolve) => {
+                resolveExecution = resolve;
+              });
+              const execution = { promise, resolve: resolveExecution };
+              abortSignal.addEventListener("abort", () => {
+                execution.resolve({ session: "old" });
+              });
+              executions.push(execution);
+              return promise;
+            },
+          ),
+        } satisfies Tool,
+      }),
+      { onResult, onStatusesChange: vi.fn() },
+    );
+
+    tracker.setState(createState([]));
+    tracker.setState(
+      createState([
+        createAssistantMessage('{"query":"London"}', { query: "London" }),
+      ]),
+    );
+    await waitFor(() => expect(executions).toHaveLength(1));
+
+    tracker.reset();
+    tracker.setState(createState([]));
+    tracker.setState(
+      createState([
+        createAssistantMessage(
+          '{"query":"Paris"}',
+          { query: "Paris" },
+          { toolCallId: "tool-1" },
+        ),
+      ]),
+    );
+    await waitFor(() => expect(executions).toHaveLength(2));
+    await new Promise((resolve) => setTimeout(resolve, 0));
+    expect(onResult).not.toHaveBeenCalled();
+
+    executions[1]!.resolve({ session: "new" });
+    await waitFor(() => {
+      expect(onResult).toHaveBeenCalledOnce();
+      expect(onResult).toHaveBeenCalledWith(
+        expect.objectContaining({
+          toolCallId: "tool-1",
+          result: { session: "new" },
+        }),
+      );
+    });
+
+    tracker.reset();
+  });
+});

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

@@ -7,6 +7,7 @@ import {
   unstable_toolResultStream,
   type Tool,
   type ToolModelContentPart,
+  type ToolResultStreamOptions,
 } from "assistant-stream";
 import {
   AssistantMetaTransformStream,
@@ -15,6 +16,8 @@ import {
 import { isJSONValueEqual } from "../../utils/json/is-json-equal";
 import type { ThreadMessage } from "../../types/message";
 
+const TOOL_EXECUTION_ID = Symbol.for("assistant-stream.tool-execution-id");
+
 /**
  * Streaming execution state for a frontend tool.
  */
@@ -39,6 +42,8 @@ type ToolCallEntry = {
   toolName: string;
   argsText: string;
   hasResult: boolean;
+  executionId?: symbol;
+  skipExecute?: boolean;
 } & (
   | {
       /** Restored phase — observed during a history-load snapshot. */
@@ -52,6 +57,11 @@ type ToolCallEntry = {
     }
 );
 
+type SettledResolver = {
+  executionIds: ReadonlySet<symbol>;
+  resolve: () => void;
+};
+
 const isArgsTextComplete = (argsText: string) => {
   try {
     JSON.parse(argsText);
@@ -76,6 +86,11 @@ const isEquivalentCompleteArgsText = (previous: string, next: string) => {
   return isJSONValueEqual(previousValue, nextValue);
 };
 
+const getToolExecutionId = (value: object): symbol | undefined =>
+  (value as Record<PropertyKey, unknown>)[TOOL_EXECUTION_ID] as
+    | symbol
+    | undefined;
+
 /**
  * Plain-class port of the former `useToolInvocations` React hook. Owns the
  * assistant-stream pipeline that drives client-side `streamCall` / `execute`
@@ -103,27 +118,16 @@ export class ToolInvocationTracker {
   private readonly _callbacks: ToolInvocationTracker.Callbacks;
 
   private readonly _entries = new Map<string, ToolCallEntry>();
-  /**
-   * Tool call ids whose `execute` should be short-circuited in the wrapper.
-   * Populated when an entry is created with a result already attached
-   * (history reload, mid-run resume, etc.) — `execute` is suppressed so
-   * client-side side effects don't double-run. Membership outlives the
-   * entry: `reset()` deliberately does *not* clear this so post-abort
-   * cancellation `result` chunks for pre-resolved entries can still be
-   * recognized and dropped. Growth is bounded by the number of pre-resolved
-   * tool calls observed in the session.
-   */
-  private readonly _skipExecuteStreamIds = new Set<string>();
   private readonly _humanInput = new Map<
     string,
     {
+      executionId: symbol;
       resolve: (payload: unknown) => void;
       reject: (reason: unknown) => void;
     }
   >();
-  /** In-flight `execute` invocations keyed by tool call id. */
-  private readonly _executing = new Set<string>();
-  private readonly _settledResolvers: Array<() => void> = [];
+  private readonly _executing = new Set<symbol>();
+  private readonly _settledResolvers: SettledResolver[] = [];
 
   private _statuses = new Map<string, ToolExecutionStatus>();
 
@@ -166,14 +170,21 @@ export class ToolInvocationTracker {
     const [stream, controller] = createAssistantStreamController();
     this._controller = controller;
 
+    const human = (
+      toolCallId: string,
+      payload: unknown,
+      executionId?: symbol,
+    ) => this._onHumanInput(toolCallId, payload, executionId);
     const transform = unstable_toolResultStream(
       () => this._getWrappedTools(),
       () => this._ac.signal,
-      (toolCallId, payload) => this._onHumanInput(toolCallId, payload),
+      human,
       {
-        onExecutionStart: (id) => this._onExecutionStart(id),
-        onExecutionEnd: (id) => this._onExecutionEnd(id),
-      },
+        onExecutionStart: (id: string, _name: string, executionId?: symbol) =>
+          this._onExecutionStart(id, executionId),
+        onExecutionEnd: (id: string, _name: string, executionId?: symbol) =>
+          this._onExecutionEnd(id, executionId),
+      } as ToolResultStreamOptions,
     );
 
     stream
@@ -281,10 +292,16 @@ export class ToolInvocationTracker {
       this._pendingRestore = true;
       this._entries.clear();
       this._lastSnapshot = null;
-      // `_skipExecuteStreamIds` is intentionally not cleared — see field doc.
-      void this.abort().finally(() => {
-        this._executing.clear();
-      });
+      void this.abort();
+      // Statuses are cleared synchronously: discarded executions may never
+      // settle (aborting hands the signal to the tool, it does not force
+      // settlement), and a late settlement of one sibling must not republish
+      // the others. `_deleteStatus` no-ops for cleared ids, so post-reset
+      // settlements stay silent.
+      if (this._statuses.size > 0) {
+        this._statuses = new Map();
+        this._invokeOnStatusesChange();
+      }
     } catch (err) {
       console.error("[ToolInvocationTracker] reset failed", err);
     }
@@ -312,8 +329,9 @@ export class ToolInvocationTracker {
       if (this._executing.size === 0) {
         return Promise.resolve();
       }
+      const executionIds = new Set(this._executing);
       return new Promise<void>((resolve) => {
-        this._settledResolvers.push(resolve);
+        this._settledResolvers.push({ executionIds, resolve });
       });
     } catch (err) {
       console.error("[ToolInvocationTracker] abort failed", err);
@@ -360,34 +378,70 @@ export class ToolInvocationTracker {
     return Object.fromEntries(
       Object.entries(tools).map(([name, tool]) => {
         const execute = tool.execute;
-        if (execute === undefined) return [name, tool];
+        const streamCall = tool.streamCall;
+        if (execute === undefined && streamCall === undefined)
+          return [name, tool];
 
         const wrappedTool = {
           ...tool,
-          execute: (
-            ...[args, context]: Parameters<NonNullable<typeof execute>>
-          ) => {
-            if (this._skipExecuteStreamIds.has(context.toolCallId)) {
-              // Pre-resolved tool call: never invoke the host's execute.
-              // Returning a never-settling Promise keeps the executor's
-              // pending entry alive but enqueues nothing.
-              return new Promise(() => {}) as never;
-            }
-            return execute(args, context);
-          },
+          ...(execute !== undefined && {
+            execute: (...[args, context]: Parameters<typeof execute>) => {
+              const executionId = getToolExecutionId(context);
+              const entry = this._captureExecution(
+                context.toolCallId,
+                executionId,
+              );
+              if (!entry || entry.skipExecute) {
+                return new Promise(() => {}) as never;
+              }
+              return execute(args, context);
+            },
+          }),
+          ...(streamCall !== undefined && {
+            streamCall: (
+              ...[reader, context]: Parameters<typeof streamCall>
+            ) => {
+              const executionId = getToolExecutionId(context);
+              const entry = this._captureExecution(
+                context.toolCallId,
+                executionId,
+              );
+              if (!entry) return;
+              return streamCall(reader, context);
+            },
+          }),
         } as Tool;
         return [name, wrappedTool];
       }),
     ) as Record<string, Tool>;
   }
 
+  private _captureExecution(
+    toolCallId: string,
+    executionId: symbol | undefined,
+  ): ToolCallEntry | undefined {
+    if (executionId === undefined) return undefined;
+    const entry = this._entries.get(toolCallId);
+    if (!entry?.controller) return undefined;
+    if (entry.executionId === undefined) entry.executionId = executionId;
+    return entry.executionId === executionId ? entry : undefined;
+  }
+
   // ──────────────── internal: execution lifecycle callbacks ────────────────
 
   private _onHumanInput(
     toolCallId: string,
     payload: unknown,
+    executionId?: symbol,
   ): Promise<unknown> {
     return new Promise<unknown>((resolve, reject) => {
+      // A discarded execution must not resurrect a status entry or park an
+      // unanswerable request after its id has been reused.
+      const entry = this._entries.get(toolCallId);
+      if (!entry?.controller || entry.executionId !== executionId) {
+        reject(new Error("Tool execution aborted"));
+        return;
+      }
       const previous = this._humanInput.get(toolCallId);
       if (previous) {
         try {
@@ -398,7 +452,11 @@ export class ToolInvocationTracker {
           // host rejection handler threw; ignore and proceed
         }
       }
-      this._humanInput.set(toolCallId, { resolve, reject });
+      this._humanInput.set(toolCallId, {
+        executionId: executionId!,
+        resolve,
+        reject,
+      });
       this._setStatus(toolCallId, {
         type: "interrupt",
         payload: { type: "human", payload },
@@ -406,28 +464,42 @@ export class ToolInvocationTracker {
     });
   }
 
-  private _onExecutionStart(toolCallId: string): void {
-    if (this._skipExecuteStreamIds.has(toolCallId)) return;
+  private _onExecutionStart(
+    toolCallId: string,
+    executionId: symbol | undefined,
+  ): void {
+    if (!this._captureExecution(toolCallId, executionId)) return;
+    const entry = this._entries.get(toolCallId)!;
+    if (entry.skipExecute) return;
 
-    this._executing.add(toolCallId);
+    this._executing.add(executionId!);
     this._setStatus(toolCallId, { type: "executing" });
   }
 
-  private _onExecutionEnd(toolCallId: string): void {
-    if (!this._executing.delete(toolCallId)) return;
+  private _onExecutionEnd(
+    toolCallId: string,
+    executionId: symbol | undefined,
+  ): void {
+    if (executionId === undefined || !this._executing.delete(executionId))
+      return;
 
-    this._deleteStatus(toolCallId);
+    const entry = this._entries.get(toolCallId);
+    if (entry?.executionId === executionId) this._deleteStatus(toolCallId);
 
-    if (this._executing.size === 0) {
-      const resolvers = this._settledResolvers.splice(0);
-      resolvers.forEach((resolve) => {
-        try {
-          resolve();
-        } catch {
-          // ignore — settled-resolver consumer threw
-        }
-      });
-    }
+    const pending: SettledResolver[] = [];
+    this._settledResolvers.forEach(({ executionIds, resolve }) => {
+      if ([...executionIds].some((id) => this._executing.has(id))) {
+        pending.push({ executionIds, resolve });
+        return;
+      }
+      try {
+        resolve();
+      } catch {
+        // ignore — settled-resolver consumer threw
+      }
+    });
+    this._settledResolvers.length = 0;
+    this._settledResolvers.push(...pending);
   }
 
   private _handleResultChunk(chunk: {
@@ -439,14 +511,10 @@ export class ToolInvocationTracker {
     meta: { toolCallId: string; toolName: string };
   }): void {
     const toolCallId = chunk.meta.toolCallId;
+    const executionId = getToolExecutionId(chunk);
     const entry = this._entries.get(toolCallId);
 
-    // Pre-resolved tool call whose entry has been cleared by `reset()`.
-    // The post-abort cancellation chunk lands here after the entry is
-    // gone; suppress via the long-lived skip-execute marker.
-    if (!entry && this._skipExecuteStreamIds.has(toolCallId)) {
-      return;
-    }
+    if (!entry || entry.executionId !== executionId) return;
 
     // The host already set the result (via the live snapshot's
     // `setResponse` path). Suppress the executor's redundant emit.
@@ -538,14 +606,12 @@ export class ToolInvocationTracker {
       toolName,
       toolCallId,
     });
-    if (skipExecute) {
-      this._skipExecuteStreamIds.add(toolCallId);
-    }
     const entry: ToolCallEntry = {
       toolName,
       controller: toolCallController,
       argsText: "",
       hasResult: false,
+      skipExecute,
       argsComplete: false,
     };
     this._entries.set(toolCallId, entry);

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

@@ -39,10 +39,12 @@ describe("ModelContext", () => {
   it("reflects modelName from a registered provider", async () => {
     const { sub, unmount } = render();
     try {
-      sub.getValue().register(provider({ config: { modelName: "gpt-4" } }));
+      sub
+        .getValue()
+        .register(provider({ config: { modelName: "gpt-5.6-luna" } }));
       await tick();
 
-      expect(sub.getValue().getState().modelName).toBe("gpt-4");
+      expect(sub.getValue().getState().modelName).toBe("gpt-5.6-luna");
     } finally {
       unmount();
     }
@@ -67,11 +69,15 @@ describe("ModelContext", () => {
   it("keeps the same state reference when an extra provider does not change the merged values", async () => {
     const { sub, unmount } = render();
     try {
-      sub.getValue().register(provider({ config: { modelName: "gpt-4" } }));
+      sub
+        .getValue()
+        .register(provider({ config: { modelName: "gpt-5.6-luna" } }));
       await tick();
       const before = sub.getValue().getState();
 
-      sub.getValue().register(provider({ config: { modelName: "gpt-4" } }));
+      sub
+        .getValue()
+        .register(provider({ config: { modelName: "gpt-5.6-luna" } }));
       await tick();
       const after = sub.getValue().getState();
 
@@ -86,9 +92,9 @@ describe("ModelContext", () => {
     try {
       const unsubscribe = sub
         .getValue()
-        .register(provider({ config: { modelName: "gpt-4" } }));
+        .register(provider({ config: { modelName: "gpt-5.6-luna" } }));
       await tick();
-      expect(sub.getValue().getState().modelName).toBe("gpt-4");
+      expect(sub.getValue().getState().modelName).toBe("gpt-5.6-luna");
 
       unsubscribe();
       await tick();
@@ -102,9 +108,11 @@ describe("ModelContext", () => {
   it("reflects modelName synchronously after register without awaiting", () => {
     const { sub, unmount } = render();
     try {
-      sub.getValue().register(provider({ config: { modelName: "gpt-4" } }));
+      sub
+        .getValue()
+        .register(provider({ config: { modelName: "gpt-5.6-luna" } }));
 
-      expect(sub.getValue().getState().modelName).toBe("gpt-4");
+      expect(sub.getValue().getState().modelName).toBe("gpt-5.6-luna");
     } finally {
       unmount();
     }