@assistant-ui/core 0.2.20 → 0.2.22

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

@@ -1,4 +1,4 @@
-import { describe, expect, it } from "vitest";
+import { afterEach, describe, expect, it, vi } from "vitest";
 import { LocalRuntimeCore } from "./local-runtime-core";
 import type {
   ChatModelAdapter,
@@ -6,13 +6,28 @@ import type {
   ChatModelRunResult,
 } from "../../runtime/utils/chat-model-adapter";
 import type { AppendMessage } from "../../types/message";
+import type { LocalRuntimeOptionsBase } from "./local-runtime-options";
 
 const flush = () => new Promise((resolve) => setTimeout(resolve, 10));
 
-const createThread = (adapter: ChatModelAdapter) => {
+afterEach(() => {
+  vi.restoreAllMocks();
+});
+
+const createThread = (
+  adapter: ChatModelAdapter,
+  options?: {
+    suggestion?: LocalRuntimeOptionsBase["adapters"]["suggestion"];
+  },
+) => {
   const core = new LocalRuntimeCore(
     {
-      adapters: { chatModel: adapter },
+      adapters: {
+        chatModel: adapter,
+        ...(options?.suggestion !== undefined && {
+          suggestion: options.suggestion,
+        }),
+      },
       unstable_humanToolNames: ["send_email"],
     },
     undefined,
@@ -63,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"));
@@ -117,6 +190,53 @@ describe("LocalThreadRuntimeCore human-in-the-loop tools", () => {
       "text",
     ]);
   });
+
+  it.each([
+    ["false", false],
+    ["zero", 0],
+    ["an empty string", ""],
+    ["null", null],
+  ])("resumes when the tool result is %s", async (_label, result) => {
+    const { thread, runs } = createApprovalThread(toolCallResult("send_email"));
+
+    await thread.append(userMessage("send an email"));
+    await flush();
+
+    thread.addToolResult({
+      messageId: thread.messages.at(-1)!.id,
+      toolCallId: "call-send_email",
+      toolName: "send_email",
+      result,
+      isError: false,
+    });
+    await flush();
+
+    expect(runs).toHaveLength(2);
+    const toolCall = runs[1]!
+      .unstable_getMessage()
+      .content.find((part) => part.type === "tool-call");
+    expect(toolCall?.result).toEqual(result);
+    expect(thread.messages.at(-1)?.status?.type).toBe("complete");
+  });
+});
+
+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", () => {
@@ -368,3 +488,156 @@ describe("LocalThreadRuntimeCore tool approvals", () => {
     ).toThrowError(/unstable_humanToolNames/);
   });
 });
+
+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(
+      ({ signal }: { signal?: AbortSignal }) =>
+        new Promise<readonly { prompt: string }[]>((resolve) => {
+          signal?.addEventListener("abort", () => {
+            resolve([{ prompt: "stale" }]);
+          });
+        }),
+    );
+
+    const thread = createThread(
+      {
+        async run() {
+          return { content: [{ type: "text", text: "hello" }] };
+        },
+      },
+      { suggestion: { generate } },
+    );
+
+    const appendPromise = thread.append(userMessage("hi"));
+    await appendPromise;
+    await new Promise((r) => setTimeout(r, 0));
+
+    expect(generate).toHaveBeenCalledTimes(1);
+    const signal = generate.mock.calls[0]![0].signal as AbortSignal;
+    expect(signal.aborted).toBe(false);
+
+    thread.cancelRun();
+    expect(signal.aborted).toBe(true);
+    expect(thread.suggestions).toEqual([]);
+  });
+
+  it("completes the run when suggestion generation rejects", async () => {
+    const generate = vi.fn().mockRejectedValue(new Error("suggestion failed"));
+
+    const thread = createThread(
+      {
+        async run() {
+          return { content: [{ type: "text", text: "hello" }] };
+        },
+      },
+      { suggestion: { generate } },
+    );
+
+    await thread.append(userMessage("hi"));
+    await new Promise((r) => setTimeout(r, 0));
+
+    expect(generate).toHaveBeenCalledTimes(1);
+    expect(thread.messages.at(-1)?.status?.type).toBe("complete");
+    expect(thread.suggestions).toEqual([]);
+  });
+
+  it("resolves append before suggestion generation completes", async () => {
+    let resolveSuggestions!: (value: readonly { prompt: string }[]) => void;
+    const suggestionsDeferred = new Promise<readonly { prompt: string }[]>(
+      (resolve) => {
+        resolveSuggestions = resolve;
+      },
+    );
+    const generate = vi.fn().mockReturnValue(suggestionsDeferred);
+
+    const thread = createThread(
+      {
+        async run() {
+          return { content: [{ type: "text", text: "hello" }] };
+        },
+      },
+      { suggestion: { generate } },
+    );
+
+    await thread.append(userMessage("hi"));
+    expect(thread.messages.at(-1)?.status?.type).toBe("complete");
+    expect(thread.suggestions).toEqual([]);
+
+    await new Promise((r) => setTimeout(r, 0));
+    expect(generate).toHaveBeenCalledTimes(1);
+    expect(thread.suggestions).toEqual([]);
+
+    resolveSuggestions([{ prompt: "follow up" }]);
+    await new Promise((r) => setTimeout(r, 0));
+    expect(thread.suggestions).toEqual([{ prompt: "follow up" }]);
+  });
+});

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

@@ -6,6 +6,7 @@ import type {
 } from "../../runtime/utils/chat-model-adapter";
 import { shouldContinue } from "./should-continue";
 import type { LocalRuntimeOptionsBase } from "./local-runtime-options";
+import { consumeSuggestionResult } from "../../adapters/suggestion";
 import type {
   AddToolResultOptions,
   ResumeToolCallOptions,
@@ -21,6 +22,7 @@ import type {
   ThreadAssistantMessage,
 } from "../../types/message";
 import type { RunConfig } from "../../types/message";
+import { toAssistantError } from "../../types/error";
 import type { ModelContextProvider } from "../../model-context/types";
 import {
   createMessageQueue,
@@ -375,28 +377,28 @@ export class LocalThreadRuntimeCore
       }
     }
 
-    this._suggestionsController = new AbortController();
-    const signal = this._suggestionsController.signal;
     if (
       this.adapters.suggestion &&
       message.status?.type !== "requires-action"
     ) {
-      const promiseOrGenerator = this.adapters.suggestion?.generate({
-        messages: this.messages,
-      });
-
-      if (Symbol.asyncIterator in promiseOrGenerator) {
-        for await (const r of promiseOrGenerator) {
-          if (signal.aborted) break;
-          this._suggestions = r;
-          this._notifySubscribers();
-        }
-      } else {
-        const result = await promiseOrGenerator;
-        if (signal.aborted) return;
-        this._suggestions = result;
-        this._notifySubscribers();
-      }
+      this._suggestionsController = new AbortController();
+      const signal = this._suggestionsController.signal;
+      const adapter = this.adapters.suggestion;
+      void (async () => {
+        try {
+          const promiseOrGenerator = adapter.generate({
+            messages: this.messages,
+            signal,
+          });
+          await consumeSuggestionResult(promiseOrGenerator, {
+            signal,
+            onUpdate: (r) => {
+              this._suggestions = r;
+              this._notifySubscribers();
+            },
+          });
+        } catch {}
+      })();
     }
   }
 
@@ -440,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
@@ -535,11 +537,12 @@ 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) {
-      // TODO this should be handled by the run result stream
       if (e instanceof AbortError) {
         updateMessage({
           status: { type: "incomplete", reason: "cancelled" },
@@ -553,10 +556,7 @@ export class LocalThreadRuntimeCore
           status: {
             type: "incomplete",
             reason: "error",
-            error:
-              e instanceof Error
-                ? e.message
-                : `[${typeof e}] ${new String(e).toString()}`,
+            error: toAssistantError(e),
           },
         });
 
@@ -584,6 +584,8 @@ export class LocalThreadRuntimeCore
     const error = new AbortError(true);
     this.abortController?.abort(error);
     this.abortController = null;
+    this._suggestionsController?.abort();
+    this._suggestionsController = null;
   }
 
   public cancelRun() {
@@ -591,6 +593,8 @@ export class LocalThreadRuntimeCore
     const error = new AbortError(false);
     this.abortController?.abort(error);
     this.abortController = null;
+    this._suggestionsController?.abort();
+    this._suggestionsController = null;
   }
 
   public addToolResult({

package/src/runtimes/local/should-continue.ts

@@ -23,14 +23,17 @@ export const shouldContinue = (
   // TODO legacy behavior -- make specifying human tool names required
   if (humanToolNames === undefined) {
     return result.content.every(
-      (c) => c.type !== "tool-call" || !!c.result || c.approval !== undefined,
+      (c) =>
+        c.type !== "tool-call" ||
+        c.result !== undefined ||
+        c.approval !== undefined,
     );
   }
 
   return result.content.every(
     (c) =>
       c.type !== "tool-call" ||
-      !!c.result ||
+      c.result !== undefined ||
       c.approval !== undefined ||
       !humanToolNames.includes(c.toolName),
   );

package/src/runtimes/remote-thread-list/adapter/in-memory.test.ts

@@ -0,0 +1,14 @@
+import { describe, expect, it } from "vitest";
+import { InMemoryThreadListAdapter } from "./in-memory";
+
+describe("InMemoryThreadListAdapter", () => {
+  it("closes its title stream when title generation is unsupported", async () => {
+    const adapter = new InMemoryThreadListAdapter();
+    const stream = await adapter.generateTitle();
+
+    await expect(stream.getReader().read()).resolves.toEqual({
+      done: true,
+      value: undefined,
+    });
+  });
+});

package/src/runtimes/remote-thread-list/adapter/in-memory.ts

@@ -38,10 +38,18 @@ export class InMemoryThreadListAdapter implements RemoteThreadListAdapter {
   }
 
   generateTitle(): Promise<AssistantStream> {
-    return Promise.resolve(new ReadableStream<AssistantStreamChunk>());
-  }
-
-  fetch(_threadId: string): Promise<RemoteThreadMetadata> {
-    return Promise.reject(new Error("Thread not found"));
+    return Promise.resolve(
+      new ReadableStream<AssistantStreamChunk>({
+        start(controller) {
+          controller.close();
+        },
+      }),
+    );
+  }
+
+  fetch(threadId: string): Promise<RemoteThreadMetadata> {
+    return Promise.reject(
+      new Error(`Thread "${threadId}" not found in in-memory thread list.`),
+    );
   }
 }

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 = {