@assistant-ui/core 0.2.21 → 0.2.23

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

@@ -3,7 +3,6 @@ import { describe, expect, it, vi } from "vitest";
 import {
   ToolInvocationTracker,
   type ToolExecutionStatus,
-  type ToolInvocationTrackerSnapshot,
 } from "./ToolInvocationTracker";
 import type {
   ThreadAssistantMessage,
@@ -34,7 +33,7 @@ async function waitFor(
 const createState = (
   messages: ThreadAssistantMessage[],
   isRunning: boolean = true,
-): ToolInvocationTrackerSnapshot => ({
+): ToolInvocationTracker.Snapshot => ({
   messages: messages as readonly ThreadMessage[],
   isRunning,
 });
@@ -152,6 +151,131 @@ describe("ToolInvocationTracker", () => {
     }
   });
 
+  it("does not auto-submit a parse-error result when divergent argsText closes without a backend result", async () => {
+    // A human-in-the-loop tool whose argsText diverges mid-stream and never
+    // re-converges, with no backend result at close time. The args stream must
+    // not close on the divergent snapshot (which holds a stale prefix the
+    // execution path would parse), so no bogus parse-error result is
+    // auto-submitted to resume the host graph and abandon the pending
+    // interrupt.
+    const execute = vi.fn(async () => ({ forecast: "ok" }));
+    const streamCall = vi.fn((_reader, { human }) => {
+      // Request human input immediately — sets up the pending interrupt.
+      void human({ request: "approve" });
+    });
+    const getTools = () => ({
+      weatherSearch: {
+        parameters: { type: "object", properties: {} },
+        execute,
+        streamCall,
+      } satisfies Tool,
+    });
+    const onResult = vi.fn();
+    let statuses: Record<string, ToolExecutionStatus> = {};
+    const onStatusesChange = (s: ReadonlyMap<string, ToolExecutionStatus>) => {
+      statuses = Object.fromEntries(s);
+    };
+    const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => {});
+
+    try {
+      const tracker = new ToolInvocationTracker(getTools, {
+        onResult,
+        onStatusesChange,
+      });
+      tracker.setState(createState([]));
+
+      tracker.setState(
+        createState([
+          createAssistantMessage('{"query":"London","longitude":0', {
+            query: "London",
+            longitude: 0,
+          }),
+        ]),
+      );
+
+      // The pending interrupt is set up via streamCall → human().
+      await waitFor(() => {
+        expect(statuses["tool-1"]?.type).toBe("interrupt");
+      });
+
+      // Divergent regression (not a prefix of the streamed text).
+      tracker.setState(
+        createState([
+          createAssistantMessage('{"query":"London","longitude":-0.125', {
+            query: "London",
+            longitude: -0.125,
+          }),
+        ]),
+      );
+
+      // Complete valid JSON, divergent from the streamed prefix, no backend
+      // result. Previously this closed the args stream on the snapshot's
+      // completeness, parsed the stale prefix, and auto-submitted an error.
+      tracker.setState(
+        createState([
+          createAssistantMessage(
+            '{"query":"London","longitude":-0.125,"latitude":51.5072}',
+            { query: "London", longitude: -0.125, latitude: 51.5072 },
+          ),
+        ]),
+      );
+
+      for (let i = 0; i < 5; i++) {
+        await new Promise((r) => setTimeout(r, 0));
+      }
+
+      // No bogus parse-error result is auto-submitted, so the host graph is
+      // not resumed with a fake tool failure and the pending interrupt is
+      // preserved.
+      expect(onResult).not.toHaveBeenCalled();
+      // The frontend execute never ran: the stale prefix was never parsed.
+      expect(execute).not.toHaveBeenCalled();
+    } finally {
+      warnSpy.mockRestore();
+    }
+  });
+
+  it("does not auto-submit a parse-error result for a non-executable tool whose divergent argsText closes", async () => {
+    // Same close-gating mismatch as the executable case, but for a tool with
+    // no frontend execute. Closing on the divergent complete snapshot would
+    // still parse the incomplete stale prefix and fabricate a parse-error
+    // result, so the close must gate on the controller's streamed content.
+    const getTools = () => ({
+      weatherSearch: {
+        parameters: { type: "object", properties: {} },
+      } satisfies Tool,
+    });
+    const onResult = vi.fn();
+    const onStatusesChange = () => {};
+    const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => {});
+
+    try {
+      const tracker = new ToolInvocationTracker(getTools, {
+        onResult,
+        onStatusesChange,
+      });
+      tracker.setState(createState([]));
+
+      // Incomplete prefix while the run is still streaming.
+      tracker.setState(
+        createState([createAssistantMessage('{"a":1', { a: 1 })], true),
+      );
+
+      // Run settles and the snapshot regresses to a complete divergent text.
+      tracker.setState(
+        createState([createAssistantMessage('{"a":2}', { a: 2 })], false),
+      );
+
+      for (let i = 0; i < 5; i++) {
+        await new Promise((r) => setTimeout(r, 0));
+      }
+
+      expect(onResult).not.toHaveBeenCalled();
+    } finally {
+      warnSpy.mockRestore();
+    }
+  });
+
   it("clears executing status under the logical toolCallId when reset() lands while execute is pending", async () => {
     // Tests the F.1 lifecycle: reset() aborts in-flight execute() invocations
     // and clears their executing status. The status key is the logical
@@ -948,12 +1072,8 @@ describe("ToolInvocationTracker", () => {
     // result replacement) and verifying streamCall fires exactly once.
     //
     // The pathological mid-stream regression case (A.2) is covered by
-    // the dedicated regression test above. Mixing A.2 with a backend
-    // result in the same snapshot exposes a separate issue inside
-    // assistant-stream's `ToolCallStreamController.setResponse` ordering
-    // (parse-failure result reaches the reader before the backend
-    // result); that's tracked separately and out of scope for the
-    // tracker-level contract.
+    // the assistant-stream ordering regression test in
+    // packages/assistant-stream/src/core/modules/tool-call.test.ts.
     const streamCall = vi.fn();
     const execute = vi.fn(async () => ({ forecast: "ok" }));
     const getTools = () => ({

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

@@ -35,34 +35,6 @@ export type AddToolResultCommand = {
   readonly modelContent?: readonly ToolModelContentPart[];
 };
 
-export type ToolInvocationTrackerSnapshot = {
-  readonly messages: readonly ThreadMessage[];
-  /** Whether the producing runtime is currently streaming new output. */
-  readonly isRunning: boolean;
-  /**
-   * Whether the producing runtime is still loading historical state.
-   * When `true`, every snapshot is treated as historical (no `streamCall` /
-   * `execute` fires). When `false`, processing resumes as live.
-   */
-  readonly isLoading?: boolean;
-};
-
-export type ToolInvocationTrackerCallbacks = {
-  /**
-   * Invoked when a client-side `execute()` returns a result and the runtime
-   * needs to feed it back into the conversation.
-   */
-  onResult: (command: AddToolResultCommand) => void;
-  /**
-   * Invoked whenever the per-tool-call status map changes (executing /
-   * interrupt / cleared). The callback receives a fresh map; mutating the
-   * argument is not supported.
-   */
-  onStatusesChange: (
-    statuses: ReadonlyMap<string, ToolExecutionStatus>,
-  ) => void;
-};
-
 type ToolCallEntry = {
   toolName: string;
   argsText: string;
@@ -123,10 +95,12 @@ const isEquivalentCompleteArgsText = (previous: string, next: string) => {
  * the hot message-processing path, so a malformed snapshot must never crash
  * the host runtime. See ./EDGE_CASES.md for the known non-trivial state
  * transitions and what each does today.
+ *
+ * @deprecated Internal — for framework bindings; not a public API. May change without notice.
  */
 export class ToolInvocationTracker {
   private readonly _getTools: () => Record<string, Tool> | undefined;
-  private readonly _callbacks: ToolInvocationTrackerCallbacks;
+  private readonly _callbacks: ToolInvocationTracker.Callbacks;
 
   private readonly _entries = new Map<string, ToolCallEntry>();
   /**
@@ -157,7 +131,7 @@ export class ToolInvocationTracker {
   private _pendingRestore = true;
 
   /** Cached last snapshot, used to skip processing on identical re-renders. */
-  private _lastSnapshot: ToolInvocationTrackerSnapshot | null = null;
+  private _lastSnapshot: ToolInvocationTracker.Snapshot | null = null;
   private _isRunning = false;
 
   private _controller!: ReturnType<typeof createAssistantStreamController>[1];
@@ -175,7 +149,7 @@ export class ToolInvocationTracker {
 
   constructor(
     getTools: () => Record<string, Tool> | undefined,
-    callbacks: ToolInvocationTrackerCallbacks,
+    callbacks: ToolInvocationTracker.Callbacks,
   ) {
     this._getTools = getTools;
     this._callbacks = callbacks;
@@ -235,7 +209,7 @@ export class ToolInvocationTracker {
    * Feed the next observed snapshot into the tracker. Called from the host
    * runtime whenever its message list / running state changes.
    */
-  public setState(snapshot: ToolInvocationTrackerSnapshot): void {
+  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
@@ -687,14 +661,16 @@ export class ToolInvocationTracker {
     }
 
     if (!entry.argsComplete && entry.controller) {
+      // ToolExecutionStream parses the streamed prefix on close, so the close
+      // 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: content.argsText,
+        argsText: entry.argsText,
         hasResult,
       });
       if (shouldClose) {
         entry.controller.argsText.close();
-        entry.argsText = content.argsText;
         entry.argsComplete = true;
       }
     }
@@ -780,3 +756,35 @@ export class ToolInvocationTracker {
     }
   }
 }
+
+export namespace ToolInvocationTracker {
+  export type ExecutionStatus = ToolExecutionStatus;
+
+  export type Snapshot = {
+    readonly messages: readonly ThreadMessage[];
+    /** Whether the producing runtime is currently streaming new output. */
+    readonly isRunning: boolean;
+    /**
+     * Whether the producing runtime is still loading historical state.
+     * When `true`, every snapshot is treated as historical (no `streamCall` /
+     * `execute` fires). When `false`, processing resumes as live.
+     */
+    readonly isLoading?: boolean;
+  };
+
+  export type Callbacks = {
+    /**
+     * Invoked when a client-side `execute()` returns a result and the runtime
+     * needs to feed it back into the conversation.
+     */
+    onResult: (command: AddToolResultCommand) => void;
+    /**
+     * Invoked whenever the per-tool-call status map changes (executing /
+     * interrupt / cleared). The callback receives a fresh map; mutating the
+     * argument is not supported.
+     */
+    onStatusesChange: (
+      statuses: ReadonlyMap<string, ToolExecutionStatus>,
+    ) => void;
+  };
+}

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

@@ -149,7 +149,7 @@ describe("mergeModelContexts", () => {
     });
   });
 
-  it("still rejects duplicate tools at the same priority", () => {
+  it("rejects duplicate tools at the same priority", () => {
     expect(() =>
       mergeModelContexts(
         new Set([
@@ -160,6 +160,26 @@ describe("mergeModelContexts", () => {
     ).toThrow(/already exists/);
   });
 
+  it("silently replaces a same-priority tool when the later registration sets overwrite", () => {
+    const merged = mergeModelContexts(
+      new Set([
+        provider({ tools: { duplicate: toolFixture() } }),
+        provider({
+          tools: {
+            duplicate: {
+              ...toolFixture(),
+              description: "latest",
+              overwrite: true,
+            } as Tool<any, any>,
+          },
+        }),
+      ]),
+    );
+
+    expect(merged.tools?.duplicate?.description).toBe("latest");
+    expect(merged.tools?.duplicate?.overwrite).toBeUndefined();
+  });
+
   it("preserves the highest priority when a lower-priority provider reuses the same tool object", () => {
     const shared = {
       ...toolFixture(),

package/src/store/runtime-clients/thread-runtime-client.ts

@@ -1,7 +1,7 @@
 import type { Unsubscribe } from "../../types/unsubscribe";
 import type { ThreadRuntimeEventType } from "../../runtime/interfaces/thread-runtime-core";
 import type { ThreadRuntime } from "../../runtime/api/thread-runtime";
-import { useMemo, useEffect, RefObject } from "react";
+import { useMemo, useEffect, type RefObject } from "react";
 import { useResource, resource, withKey } from "@assistant-ui/tap";
 import { liveRef } from "./liveRef";
 import {
@@ -113,6 +113,7 @@ const useThreadClient = ({
     deleteMessage: runtime.deleteMessage,
     startRun: runtime.startRun,
     resumeRun: runtime.resumeRun,
+    importExternalState: runtime.importExternalState,
     cancelRun: runtime.cancelRun,
     getModelContext: runtime.getModelContext,
     export: runtime.export,

package/src/store/scopes/thread.ts

@@ -108,6 +108,7 @@ export type ThreadMethods = {
    * @param initialMessages - Optional array of initial messages to populate the thread
    */
   reset(initialMessages?: readonly ThreadMessageLike[]): void;
+  importExternalState?(state: unknown): void;
   message(selector: { id: string } | { index: number }): MessageMethods;
   /** @deprecated This API is still under active development and might change without notice. */
   stopSpeaking(): void;

package/src/subscribable/subscribable.test.ts

@@ -0,0 +1,58 @@
+import { describe, expect, it, vi } from "vitest";
+import type { SubscribableWithState } from "./subscribable";
+import { ShallowMemoizeSubject } from "./subscribable";
+
+type TestState = {
+  status: string;
+  error?: string;
+};
+
+const createBinding = (initialState: TestState) => {
+  let state = initialState;
+  const subscribers = new Set<() => void>();
+
+  const binding: SubscribableWithState<TestState, null> = {
+    path: null,
+    getState: () => state,
+    subscribe: (callback) => {
+      subscribers.add(callback);
+      return () => subscribers.delete(callback);
+    },
+  };
+
+  return {
+    binding,
+    update(nextState: TestState) {
+      state = nextState;
+      for (const callback of subscribers) callback();
+    },
+  };
+};
+
+describe("ShallowMemoizeSubject", () => {
+  it("notifies subscribers when a state key is removed", () => {
+    const source = createBinding({
+      status: "running",
+      error: "Connection failed",
+    });
+    const subject = new ShallowMemoizeSubject(source.binding);
+    const subscriber = vi.fn();
+    subject.subscribe(subscriber);
+
+    source.update({ status: "running" });
+
+    expect(subscriber).toHaveBeenCalledOnce();
+    expect(subject.getState()).toEqual({ status: "running" });
+  });
+
+  it("does not notify subscribers for a shallow-equal state", () => {
+    const source = createBinding({ status: "running" });
+    const subject = new ShallowMemoizeSubject(source.binding);
+    const subscriber = vi.fn();
+    subject.subscribe(subscriber);
+
+    source.update({ status: "running" });
+
+    expect(subscriber).not.toHaveBeenCalled();
+  });
+});

package/src/subscribable/subscribable.ts

@@ -1,4 +1,5 @@
 import type { Unsubscribe } from "../types/unsubscribe";
+import { notifyEventListeners } from "../utils/notify-event-listeners";
 
 export const SKIP_UPDATE = Symbol("skip-update");
 export type SKIP_UPDATE = typeof SKIP_UPDATE;
@@ -39,7 +40,10 @@ function shallowEqual<T extends object>(
   if (objA === undefined) return false;
   if (objB === undefined) return false;
 
-  for (const key of Object.keys(objA)) {
+  const keysA = Object.keys(objA);
+  if (keysA.length !== Object.keys(objB).length) return false;
+
+  for (const key of keysA) {
     const valueA = objA[key as keyof T];
     const valueB = objB[key as keyof T];
     if (!Object.is(valueA, valueB)) return false;
@@ -99,7 +103,12 @@ export abstract class BaseSubject {
 
   protected abstract _connect(): Unsubscribe;
 
-  protected notifySubscribers(payload?: unknown) {
+  protected notifySubscribers(payload?: unknown, errorContext?: string) {
+    if (errorContext) {
+      notifyEventListeners(this._subscriptions, payload, errorContext);
+      return;
+    }
+
     for (const callback of this._subscriptions) callback(payload);
   }
 
@@ -132,10 +141,11 @@ export class ShallowMemoizeSubject<TState extends object, TPath>
     return this.binding.path;
   }
 
-  constructor(
-    private binding: SubscribableWithState<TState | SKIP_UPDATE, TPath>,
-  ) {
+  private binding: SubscribableWithState<TState | SKIP_UPDATE, TPath>;
+
+  constructor(binding: SubscribableWithState<TState | SKIP_UPDATE, TPath>) {
     super();
+    this.binding = binding;
     const state = binding.getState();
     if (state === SKIP_UPDATE)
       throw new Error("Entry not available in the store");
@@ -175,10 +185,11 @@ export class LazyMemoizeSubject<TState extends object, TPath>
     return this.binding.path;
   }
 
-  constructor(
-    private binding: SubscribableWithState<TState | SKIP_UPDATE, TPath>,
-  ) {
+  private binding: SubscribableWithState<TState | SKIP_UPDATE, TPath>;
+
+  constructor(binding: SubscribableWithState<TState | SKIP_UPDATE, TPath>) {
     super();
+    this.binding = binding;
   }
 
   private _previousStateDirty = true;
@@ -219,8 +230,11 @@ export class NestedSubscriptionSubject<
     return this.binding.path;
   }
 
-  constructor(private binding: NestedSubscribable<TState, TPath>) {
+  private binding: NestedSubscribable<TState, TPath>;
+
+  constructor(binding: NestedSubscribable<TState, TPath>) {
     super();
+    this.binding = binding;
   }
 
   public getState() {
@@ -260,8 +274,11 @@ export class NestedSubscriptionSubject<
 export class EventSubscriptionSubject<
   TEvent extends string,
 > extends BaseSubject {
-  constructor(private config: EventSubscribable<TEvent>) {
+  private config: EventSubscribable<TEvent>;
+
+  constructor(config: EventSubscribable<TEvent>) {
     super();
+    this.config = config;
   }
 
   public getState() {
@@ -273,8 +290,9 @@ export class EventSubscriptionSubject<
   }
 
   protected _connect(): Unsubscribe {
+    const errorContext = `Runtime event "${this.config.event}"`;
     const callback = (payload?: unknown) => {
-      this.notifySubscribers(payload);
+      this.notifySubscribers(payload, errorContext);
     };
 
     let lastState = this.config.binding.getState();

package/src/tests/MessageRepository.test.ts

@@ -110,6 +110,18 @@ describe("MessageRepository", () => {
       }).toThrow(/Parent message not found/);
     });
 
+    it("should throw when a message would be linked under an ancestor with the same id", () => {
+      const message = createTestMessage({ id: "message-id" });
+      repository.addOrUpdateMessage(null, message);
+
+      expect(() => {
+        repository.addOrUpdateMessage(
+          "message-id",
+          createTestMessage({ id: "message-id" }),
+        );
+      }).toThrow(/same id already exists in the parent tree/);
+    });
+
     it("should retrieve all messages in the current branch", () => {
       const parent = createTestMessage({ id: "parent-id" });
       const child = createTestMessage({ id: "child-id" });

package/src/tests/OptimisticState-list-race.test.ts

@@ -79,6 +79,39 @@ const applyListResult = (
 };
 
 describe("list + delete race condition", () => {
+  it("stale list() does not re-add a thread deleted before list started", async () => {
+    const state = new OptimisticState<RemoteThreadState>(
+      applyListResult(EMPTY_STATE, {
+        threads: [{ remoteId: "thread-A", status: "regular", title: "A" }],
+      }),
+    );
+    const deleteDeferred = deferred<void>();
+    const listDeferred = deferred<ListResult>();
+
+    const deletePromise = state.optimisticUpdate({
+      execute: () => deleteDeferred.promise,
+      optimistic: (s) => updateStatusReducer(s, "thread-A", "deleted"),
+    });
+    const listPromise = state.optimisticUpdate({
+      execute: () => listDeferred.promise,
+      loading: (s) => ({ ...s, isLoading: true }),
+      then: applyListResult,
+    });
+
+    deleteDeferred.resolve();
+    await deletePromise;
+
+    listDeferred.resolve({
+      threads: [{ remoteId: "thread-A", status: "regular", title: "A" }],
+    });
+    await listPromise;
+
+    expect(state.value.threadIds).not.toContain("thread-A");
+    expect(
+      state.value.threadData[createThreadMappingId("thread-A")],
+    ).toBeUndefined();
+  });
+
   it("stale list() does not re-add a thread deleted while list was in flight", async () => {
     const state = new OptimisticState<RemoteThreadState>(EMPTY_STATE);