@assistant-ui/core 0.2.21 → 0.2.22

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

@@ -1,4 +1,4 @@
-import { describe, expect, it, vi } from "vitest";
+import { afterEach, describe, expect, it, vi } from "vitest";
 import { LocalRuntimeCore } from "./local-runtime-core";
 import type {
   ChatModelAdapter,
@@ -10,6 +10,10 @@ import type { LocalRuntimeOptionsBase } from "./local-runtime-options";
 
 const flush = () => new Promise((resolve) => setTimeout(resolve, 10));
 
+afterEach(() => {
+  vi.restoreAllMocks();
+});
+
 const createThread = (
   adapter: ChatModelAdapter,
   options?: {
@@ -74,6 +78,64 @@ const createApprovalThread = (firstResult: ChatModelRunResult) => {
   return { thread, runs };
 };
 
+describe("LocalThreadRuntimeCore events", () => {
+  it("isolates runEnd listener errors", async () => {
+    const listenerError = new Error("telemetry failed");
+    const consoleError = vi
+      .spyOn(console, "error")
+      .mockImplementation(() => {});
+    const laterListener = vi.fn();
+    const thread = createThread({
+      async run() {
+        return { content: [{ type: "text", text: "done" }] };
+      },
+    });
+
+    thread.unstable_on("runEnd", () => {
+      throw listenerError;
+    });
+    thread.unstable_on("runEnd", laterListener);
+
+    await expect(thread.append(userMessage("hello"))).resolves.toBeUndefined();
+
+    expect(laterListener).toHaveBeenCalledOnce();
+    expect(thread.messages.at(-1)?.status?.type).toBe("complete");
+    expect(consoleError).toHaveBeenCalledWith(
+      '[assistant-ui] Thread runtime "runEnd" listener threw an error',
+      listenerError,
+    );
+  });
+
+  it("isolates async runEnd listener rejections", async () => {
+    const listenerError = new Error("async telemetry failed");
+    const consoleError = vi
+      .spyOn(console, "error")
+      .mockImplementation(() => {});
+    const laterListener = vi.fn();
+    const thread = createThread({
+      async run() {
+        return { content: [{ type: "text", text: "done" }] };
+      },
+    });
+
+    thread.unstable_on("runEnd", async () => {
+      throw listenerError;
+    });
+    thread.unstable_on("runEnd", laterListener);
+
+    await expect(thread.append(userMessage("hello"))).resolves.toBeUndefined();
+
+    expect(laterListener).toHaveBeenCalledOnce();
+    expect(thread.messages.at(-1)?.status?.type).toBe("complete");
+    await vi.waitFor(() => {
+      expect(consoleError).toHaveBeenCalledWith(
+        '[assistant-ui] Thread runtime "runEnd" listener threw an error',
+        listenerError,
+      );
+    });
+  });
+});
+
 describe("LocalThreadRuntimeCore human-in-the-loop tools", () => {
   it("pauses on requires-action while a listed tool call has no result", async () => {
     const { thread, runs } = createApprovalThread(toolCallResult("send_email"));
@@ -158,6 +220,25 @@ describe("LocalThreadRuntimeCore human-in-the-loop tools", () => {
   });
 });
 
+describe("LocalThreadRuntimeCore state", () => {
+  it.each([
+    ["false", false],
+    ["zero", 0],
+    ["an empty string", ""],
+  ])("preserves %s model state", async (_label, state) => {
+    const thread = createThread({
+      async run() {
+        return { metadata: { unstable_state: state } };
+      },
+    });
+
+    await thread.append(userMessage("update state"));
+    await flush();
+
+    expect(thread.messages.at(-1)?.metadata.unstable_state).toBe(state);
+  });
+});
+
 describe("LocalThreadRuntimeCore tool approvals", () => {
   it("pauses the run while an approval is pending, even for unlisted tools", async () => {
     const { thread, runs } = createApprovalThread(
@@ -408,6 +489,74 @@ describe("LocalThreadRuntimeCore tool approvals", () => {
   });
 });
 
+describe("LocalThreadRuntimeCore cancellation", () => {
+  it("marks the message cancelled when a streaming adapter returns after abort", async () => {
+    let released!: () => void;
+    const streaming = new Promise<void>((resolve) => {
+      released = resolve;
+    });
+
+    const thread = createThread({
+      async *run({ abortSignal }) {
+        yield { content: [{ type: "text", text: "partial" }] };
+        await streaming;
+        if (abortSignal.aborted) return;
+        yield { content: [{ type: "text", text: "partial answer" }] };
+      },
+    });
+
+    const appendPromise = thread.append(userMessage("hi"));
+    await flush();
+
+    thread.cancelRun();
+    released();
+    await appendPromise;
+
+    expect(thread.messages.at(-1)?.status).toEqual({
+      type: "incomplete",
+      reason: "cancelled",
+    });
+  });
+
+  it("marks the message cancelled when a non-streaming adapter resolves after abort", async () => {
+    let released!: () => void;
+    const pending = new Promise<void>((resolve) => {
+      released = resolve;
+    });
+
+    const thread = createThread({
+      async run() {
+        await pending;
+        return { content: [{ type: "text", text: "hello" }] };
+      },
+    });
+
+    const appendPromise = thread.append(userMessage("hi"));
+    await flush();
+
+    thread.cancelRun();
+    released();
+    await appendPromise;
+
+    expect(thread.messages.at(-1)?.status).toEqual({
+      type: "incomplete",
+      reason: "cancelled",
+    });
+  });
+
+  it("keeps a completed run complete", async () => {
+    const thread = createThread({
+      async *run() {
+        yield { content: [{ type: "text", text: "hello" }] };
+      },
+    });
+
+    await thread.append(userMessage("hi"));
+
+    expect(thread.messages.at(-1)?.status?.type).toBe("complete");
+  });
+});
+
 describe("LocalThreadRuntimeCore suggestions", () => {
   it("cancelRun aborts pending suggestion generation", async () => {
     const generate = vi.fn().mockImplementation(

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

@@ -442,7 +442,7 @@ export class LocalThreadRuntimeCore
           ? {
               metadata: {
                 ...message.metadata,
-                ...(m.metadata.unstable_state
+                ...(m.metadata.unstable_state !== undefined
                   ? { unstable_state: m.metadata.unstable_state }
                   : undefined),
                 ...(annotations
@@ -537,7 +537,9 @@ export class LocalThreadRuntimeCore
 
       if (message.status.type === "running") {
         updateMessage({
-          status: { type: "complete", reason: "unknown" },
+          status: abortSignal.aborted
+            ? { type: "incomplete", reason: "cancelled" }
+            : { type: "complete", reason: "unknown" },
         });
       }
     } catch (e) {

package/src/runtimes/remote-thread-list/optimistic-state.test.ts

@@ -0,0 +1,63 @@
+import { describe, expect, it } from "vitest";
+import { OptimisticState } from "./optimistic-state";
+
+const deferred = () => {
+  let resolve!: () => void;
+  const promise = new Promise<void>((resolvePromise) => {
+    resolve = resolvePromise;
+  });
+
+  return { promise, resolve };
+};
+
+describe("OptimisticState", () => {
+  it("preserves invocation order when optimistic updates resolve in order", async () => {
+    const state = new OptimisticState({ title: "Untitled" });
+    const firstRequest = deferred();
+    const secondRequest = deferred();
+
+    const firstUpdate = state.optimisticUpdate({
+      execute: () => firstRequest.promise,
+      optimistic: (value) => ({ ...value, title: "Project Alpha" }),
+    });
+    const secondUpdate = state.optimisticUpdate({
+      execute: () => secondRequest.promise,
+      optimistic: (value) => ({ ...value, title: "Project Beta" }),
+    });
+
+    expect(state.value.title).toBe("Project Beta");
+
+    firstRequest.resolve();
+    await firstUpdate;
+    expect(state.value.title).toBe("Project Beta");
+
+    secondRequest.resolve();
+    await secondUpdate;
+
+    expect(state.value.title).toBe("Project Beta");
+  });
+
+  it("preserves invocation order when optimistic updates resolve out of order", async () => {
+    const state = new OptimisticState({ title: "Untitled" });
+    const firstRequest = deferred();
+    const secondRequest = deferred();
+
+    const firstUpdate = state.optimisticUpdate({
+      execute: () => firstRequest.promise,
+      optimistic: (value) => ({ ...value, title: "Project Alpha" }),
+    });
+    const secondUpdate = state.optimisticUpdate({
+      execute: () => secondRequest.promise,
+      optimistic: (value) => ({ ...value, title: "Project Beta" }),
+    });
+
+    secondRequest.resolve();
+    await secondUpdate;
+    expect(state.value.title).toBe("Project Beta");
+
+    firstRequest.resolve();
+    await firstUpdate;
+
+    expect(state.value.title).toBe("Project Beta");
+  });
+});

package/src/runtimes/remote-thread-list/optimistic-state.ts

@@ -14,9 +14,15 @@ type Transform<TState, TResult> = {
 };
 
 type PendingTransform<TState, TResult> = Transform<TState, TResult> & {
+  order: number;
   task: Promise<TResult>;
 };
 
+type CompletedOptimistic<TState> = {
+  order: number;
+  optimistic: (state: TState) => TState;
+};
+
 const pipeTransforms = <TState, TExtra>(
   initialState: TState,
   extraParam: TExtra,
@@ -32,14 +38,16 @@ export class OptimisticState<TState> extends BaseSubscribable {
     [];
 
   /**
-   * `optimistic` callbacks from transforms that have already resolved.
-   * Re-applied after every `then` callback so that a wholesale state
-   * replacement (e.g. list()) cannot erase earlier completed effects
-   * (e.g. delete). Cleared when no pending transforms remain.
+   * Completed optimistic callbacks remain active while any transform is
+   * pending, so later state replacements cannot erase them. Invocation order
+   * determines which overlapping optimistic update wins.
    *
    * Correctness requirement: `optimistic` callbacks must be idempotent.
    */
-  private readonly _completedOptimistics: Array<(state: TState) => TState> = [];
+  private readonly _completedOptimistics: Array<CompletedOptimistic<TState>> =
+    [];
+
+  private _nextTransformOrder = 0;
 
   private _baseValue: TState;
   private _cachedValue: TState;
@@ -51,12 +59,25 @@ export class OptimisticState<TState> extends BaseSubscribable {
   }
 
   private _updateState(): void {
-    this._cachedValue = this._pendingTransforms.reduce((state, transform) => {
-      return pipeTransforms(state, transform.task, [
-        transform.loading,
-        transform.optimistic,
-      ]);
-    }, this._baseValue);
+    const activeTransforms = [
+      ...this._pendingTransforms.map((transform) => ({
+        order: transform.order,
+        apply: (state: TState) =>
+          pipeTransforms(state, transform.task, [
+            transform.loading,
+            transform.optimistic,
+          ]),
+      })),
+      ...this._completedOptimistics.map(({ order, optimistic }) => ({
+        order,
+        apply: optimistic,
+      })),
+    ].sort((a, b) => a.order - b.order);
+
+    this._cachedValue = activeTransforms.reduce(
+      (state, transform) => transform.apply(state),
+      this._baseValue,
+    );
 
     this._notifySubscribers();
   }
@@ -77,8 +98,13 @@ export class OptimisticState<TState> extends BaseSubscribable {
   public async optimisticUpdate<TResult>(
     transform: Transform<TState, TResult>,
   ): Promise<TResult> {
+    const order = this._nextTransformOrder++;
     const task = transform.execute();
-    const pendingTransform = { ...transform, task };
+    const pendingTransform = {
+      ...transform,
+      order,
+      task,
+    };
     try {
       this._pendingTransforms.push(pendingTransform);
       this._updateState();
@@ -89,14 +115,20 @@ export class OptimisticState<TState> extends BaseSubscribable {
         transform.then,
       ]);
 
-      // Re-apply previously completed optimistic callbacks so that a
-      // then() that does wholesale replacement cannot erase their effects.
-      for (const fn of this._completedOptimistics) {
-        this._baseValue = fn(this._baseValue);
+      // `then` can replace state with a stale snapshot, so replay every
+      // completed effect; otherwise replay only newer overlapping effects.
+      for (const completed of this._completedOptimistics) {
+        if (transform.then || completed.order > pendingTransform.order) {
+          this._baseValue = completed.optimistic(this._baseValue);
+        }
       }
 
       if (transform.optimistic) {
-        this._completedOptimistics.push(transform.optimistic);
+        this._completedOptimistics.push({
+          order: pendingTransform.order,
+          optimistic: transform.optimistic,
+        });
+        this._completedOptimistics.sort((a, b) => a.order - b.order);
       }
 
       return result;

package/src/runtimes/remote-thread-list/types.ts

@@ -5,7 +5,7 @@ import type { AssistantStream } from "assistant-stream";
 
 export type RemoteThreadInitializeResponse = {
   remoteId: string;
-  externalId: string | undefined;
+  externalId?: string | undefined;
 };
 
 export type RemoteThreadMetadata = {

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

@@ -41,6 +41,15 @@ Subsequent snapshots that *are* prefixes of the new (regressed) snapshot
 also won't be appended, because `entry.argsText` still points at the
 pre-regression value used for delta calculation.
 
+The args stream closes only when the controller's *streamed* content
+(`entry.argsText`) is complete, not when a later snapshot is. A divergent
+snapshot can be complete while the controller still holds an incomplete
+stale prefix; closing on the snapshot would parse that stale prefix and
+auto-submit a bogus parse-error result (resuming the host graph and
+abandoning a pending interrupt). Gating the close on the streamed content
+leaves the stream open until the prefix itself completes, so no stale
+parse runs and no error result is fabricated from divergent args.
+
 ### A.3. Args complete then equivalent-JSON key reorder
 Both old and new `argsText` parse to equivalent JSON values (e.g. keys
 reordered by the backend). The tracker updates its tracked `argsText`
@@ -53,9 +62,10 @@ existing `streamCall` keeps its original args view.
 
 ### A.5. First resolution (`result` becomes defined)
 The tracker calls `setResponse` on the active controller and closes it.
-`reader.response.get()` resolves. If the tool also had a frontend
-`execute`, the executor is short-circuited via `_skipExecuteStreamIds`.
-Single fire.
+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.
 
 ### A.6. Previously-resolved tool's `result` is replaced
 Silently ignored — `entry.hasResult` short-circuits both the
@@ -167,22 +177,6 @@ keep the tracker dead with a visible error to avoid restart loops.
 
 ## Known limitations
 
-### Result delivery after args regression (A.2 + A.5 in the same snapshot)
-When a snapshot has both a regressed `argsText` *and* a backend result
-on the same tool call, `activeController.setResponse(result)` closes
-`argsText` before enqueueing the result chunk. The args-text-finish
-chunk reaches `ToolExecutionStream` first, attempts to parse the
-(stale) accumulated argsText, fails, and emits a parse-error result
-that beats the backend result to the reader's response promise.
-
-The tracker's `entry.hasResult` short-circuit *does* suppress both
-result chunks at the `onResult` callback level (no double-fire), but
-the reader's `response.get()` already resolved with the parse error.
-
-Fixable upstream in `ToolCallStreamControllerImpl.setResponse` by
-enqueueing the result chunk before closing argsText. Tracked separately;
-out of scope for the tracker layer.
-
 ### Host callback throws
 `onResult` and `onStatusesChange` are invoked through wrappers that
 catch and log. The tracker continues to function; the host's bad

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

@@ -152,6 +152,131 @@ describe("ToolInvocationTracker", () => {
     }
   });
 
+  it("does not auto-submit a parse-error result when divergent argsText closes without a backend result (#5130)", async () => {
+    // Mirrors the #5098 production shape: 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 (#5130)", 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 +1073,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

@@ -687,14 +687,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;
       }
     }

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();
+  });
+});