@assistant-ui/core 0.2.20 → 0.2.21

package/src/react/runtimes/external-message-converter.test.tsx

@@ -0,0 +1,102 @@
+// @vitest-environment jsdom
+
+import { renderHook } from "@testing-library/react";
+import { describe, expect, it } from "vitest";
+import { useExternalMessageConverter } from "./external-message-converter";
+
+type TestMessage = {
+  id: string;
+  role: "user" | "assistant";
+  text: string;
+};
+
+type TestMetadata = useExternalMessageConverter.Metadata & {
+  optimisticMessageId?: string;
+};
+
+const convert: useExternalMessageConverter.Callback<TestMessage> = (
+  message,
+  metadata,
+) => ({
+  role: message.role,
+  id: message.id,
+  content: [{ type: "text", text: message.text }],
+  metadata: {
+    ...(message.role === "assistant" &&
+      message.id === (metadata as TestMetadata).optimisticMessageId && {
+        isOptimistic: true,
+      }),
+  },
+});
+
+const MESSAGES: TestMessage[] = [
+  { id: "u1", role: "user", text: "hi" },
+  { id: "a1", role: "assistant", text: "hello" },
+];
+
+const EMPTY: TestMetadata = {};
+
+type Props = {
+  callback?: useExternalMessageConverter.Callback<TestMessage>;
+  metadata?: TestMetadata;
+};
+
+const renderConverter = (initialProps: Props = {}) =>
+  renderHook(
+    ({ callback = convert, metadata = EMPTY }: Props) =>
+      useExternalMessageConverter<TestMessage>({
+        callback,
+        messages: MESSAGES,
+        isRunning: false,
+        metadata,
+      }),
+    { initialProps },
+  );
+
+describe("useExternalMessageConverter", () => {
+  it("reuses converted messages across rerenders when inputs are unchanged", () => {
+    const { result, rerender } = renderConverter();
+
+    const first = result.current;
+    rerender({});
+
+    expect(result.current[0]).toBe(first[0]);
+    expect(result.current[1]).toBe(first[1]);
+  });
+
+  it("re-converts cached messages when metadata changes", () => {
+    const { result, rerender } = renderConverter({
+      metadata: { optimisticMessageId: "a1" },
+    });
+
+    expect(result.current.at(-1)?.metadata.isOptimistic).toBe(true);
+
+    rerender({ metadata: {} });
+
+    expect(result.current.at(-1)?.metadata.isOptimistic).toBeUndefined();
+  });
+
+  it("re-converts cached messages when the callback changes", () => {
+    const upper: useExternalMessageConverter.Callback<TestMessage> = (
+      message,
+    ) => ({
+      role: message.role,
+      id: message.id,
+      content: [{ type: "text", text: message.text.toUpperCase() }],
+    });
+
+    const { result, rerender } = renderConverter();
+
+    expect(result.current[1]?.content[0]).toMatchObject({
+      type: "text",
+      text: "hello",
+    });
+
+    rerender({ callback: upper });
+
+    expect(result.current[1]?.content[0]).toMatchObject({
+      type: "text",
+      text: "HELLO",
+    });
+  });
+});

package/src/react/runtimes/external-message-converter.ts

@@ -1,6 +1,6 @@
 "use client";
 
-import { useMemo } from "react";
+import { useMemo, useState } from "react";
 import { ThreadMessageConverter } from "../../runtimes/external-store/thread-message-converter";
 import {
   getExternalStoreMessages,
@@ -76,6 +76,11 @@ type CallbackResult<T> = {
   outputs: useExternalMessageConverter.Message[];
 };
 
+type CallbackCacheEntry<T> = CallbackResult<T> & {
+  metadata: useExternalMessageConverter.Metadata;
+  callback: useExternalMessageConverter.Callback<T>;
+};
+
 type ChunkResult<T> = {
   inputs: T[];
   outputs: useExternalMessageConverter.Message[];
@@ -404,28 +409,48 @@ export const useExternalMessageConverter = <T extends WeakKey>({
   joinStrategy?: JoinStrategy | undefined;
   metadata?: useExternalMessageConverter.Metadata | undefined;
 }) => {
+  // The caches live for the component lifetime; React Compiler hoists
+  // allocations without reactive dependencies out of useMemo, so re-creating
+  // them on dependency change would not survive compilation. Staleness is
+  // instead tracked per entry via the metadata/callback that produced it,
+  // keeping correctness independent of memoization (which React treats as
+  // droppable). A "use no memo" opt-out would restore the useMemo semantics
+  // but leave cache flushing coupled to memo firing.
+  const [caches] = useState(() => ({
+    callbackCache: new WeakMap<T, CallbackCacheEntry<T>>(),
+    chunkCache: new WeakMap<
+      useExternalMessageConverter.Message,
+      ChunkResult<T>
+    >(),
+    converterCache: new ThreadMessageConverter(),
+  }));
+
   const state = useMemo(
     () => ({
       metadata: metadata ?? {},
       callback,
-      callbackCache: new WeakMap<T, CallbackResult<T>>(),
-      chunkCache: new WeakMap<
-        useExternalMessageConverter.Message,
-        ChunkResult<T>
-      >(),
-      converterCache: new ThreadMessageConverter(),
+      ...caches,
     }),
-    [callback, metadata],
+    [callback, metadata, caches],
   );
 
   return useMemo(() => {
     const callbackResults: CallbackResult<T>[] = [];
     for (const message of messages) {
       let result = state.callbackCache.get(message);
-      if (!result) {
+      if (
+        !result ||
+        result.metadata !== state.metadata ||
+        result.callback !== state.callback
+      ) {
         const output = state.callback(message, state.metadata);
         const outputs = Array.isArray(output) ? output : [output];
-        result = { input: message, outputs };
+        result = {
+          input: message,
+          outputs,
+          metadata: state.metadata,
+          callback: state.callback,
+        };
         state.callbackCache.set(message, result);
       }
       callbackResults.push(result);

package/src/runtime/base/base-composer-runtime-core.ts

@@ -195,16 +195,30 @@ export abstract class BaseComposerRuntimeCore
           )
         : [];
 
+    const originalAttachments = this.attachments;
     const text = this.text;
     const quote = this._quote;
     this._quote = undefined;
     this._emptyTextAndAttachments();
 
+    let resolvedAttachments: Awaited<typeof attachments>;
+    try {
+      resolvedAttachments = await attachments;
+    } catch (e) {
+      if (this.isEmpty && this._quote === undefined) {
+        this._attachments = originalAttachments;
+        this._text = text;
+        this._quote = quote;
+        this._notifySubscribers();
+      }
+      throw e;
+    }
+
     const message: Omit<AppendMessage, "parentId" | "sourceId"> = {
       createdAt: new Date(),
       role: this.role,
       content: text ? [{ type: "text", text }] : [],
-      attachments: await attachments,
+      attachments: resolvedAttachments,
       runConfig: this.runConfig,
       metadata: { custom: { ...(quote ? { quote } : {}) } },
     };
@@ -321,7 +335,11 @@ export abstract class BaseComposerRuntimeCore
       if (lastAttachment) {
         upsertAttachment({
           ...lastAttachment,
-          status: { type: "incomplete", reason: "error" },
+          status: {
+            type: "incomplete",
+            reason: "error",
+            message: e instanceof Error ? e.message : String(e),
+          },
         });
       }
       this._safeEmitAttachmentAddError(
@@ -333,14 +351,15 @@ export abstract class BaseComposerRuntimeCore
       throw e;
     }
 
-    const hasError =
+    if (
       lastAttachment?.status.type === "incomplete" &&
-      lastAttachment.status.reason === "error";
-    if (hasError) {
+      lastAttachment.status.reason === "error"
+    ) {
       this._safeEmitAttachmentAddError(
         "adapter-error",
-        "Attachment upload did not complete successfully.",
-        lastAttachment?.id,
+        lastAttachment.status.message ??
+          "Attachment upload did not complete successfully.",
+        lastAttachment.id,
       );
     } else {
       this._notifyEventSubscribers("attachmentAdd", {});

package/src/runtime/utils/message-repository.ts

@@ -436,13 +436,18 @@ export class MessageRepository {
   export(): ExportedMessageRepository {
     const exportItems: ExportedMessageRepository["messages"] = [];
 
-    // Optimistic messages are ephemeral and never persisted. They're always
-    // leaf nodes, so skipping them can't orphan a persisted child.
+    // Optimistic messages are ephemeral and never persisted. A persisted child
+    // of an optimistic node is re-parented onto its nearest persisted ancestor
+    // so the exported tree never references a skipped id.
     for (const [, message] of this.messages) {
       if (message.current.metadata?.isOptimistic) continue;
+      let prev = message.prev;
+      while (prev && prev.current.metadata?.isOptimistic) {
+        prev = prev.prev;
+      }
       exportItems.push({
         message: message.current,
-        parentId: message.prev?.current.id ?? null,
+        parentId: prev?.current.id ?? null,
       });
     }
 

package/src/runtime/utils/thread-message-like.ts

@@ -20,6 +20,7 @@ import type {
 import type { CompleteAttachment } from "../../types/attachment";
 import type {
   MessageTiming,
+  PartProviderMetadata,
   TextMessagePart,
   ToolApprovalOption,
   ToolCallTiming,
@@ -61,6 +62,7 @@ export type ThreadMessageLike = {
             readonly messages?: readonly ThreadMessage[] | undefined;
             readonly interrupt?: { type: "human"; payload: unknown };
             readonly timing?: ToolCallTiming;
+            readonly providerMetadata?: PartProviderMetadata;
             readonly approval?: {
               readonly id: string;
               readonly approved?: boolean;

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

@@ -1,4 +1,4 @@
-import { describe, expect, it } from "vitest";
+import { describe, expect, it, vi } from "vitest";
 import { LocalRuntimeCore } from "./local-runtime-core";
 import type {
   ChatModelAdapter,
@@ -6,13 +6,24 @@ 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) => {
+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,
@@ -117,6 +128,34 @@ 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 tool approvals", () => {
@@ -368,3 +407,88 @@ describe("LocalThreadRuntimeCore tool approvals", () => {
     ).toThrowError(/unstable_humanToolNames/);
   });
 });
+
+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 {}
+      })();
     }
   }
 
@@ -539,7 +541,6 @@ export class LocalThreadRuntimeCore
         });
       }
     } catch (e) {
-      // TODO this should be handled by the run result stream
       if (e instanceof AbortError) {
         updateMessage({
           status: { type: "incomplete", reason: "cancelled" },
@@ -553,10 +554,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 +582,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 +591,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/tests/MessageRepository.test.ts

@@ -309,6 +309,32 @@ describe("MessageRepository", () => {
       expect(exported.headId).toBe("u");
     });
 
+    it("re-parents children of a skipped optimistic message in export()", () => {
+      repository.addOrUpdateMessage(null, createTestMessage({ id: "u1" }));
+      repository.addOrUpdateMessage("u1", optimistic({ id: "a1" }));
+      repository.addOrUpdateMessage(
+        "a1",
+        createTestMessage({ id: "u2", role: "user" }),
+      );
+      repository.addOrUpdateMessage("u2", createTestMessage({ id: "a2" }));
+      repository.resetHead("a2");
+
+      const exported = repository.export();
+      const byId = new Map(exported.messages.map((m) => [m.message.id, m]));
+
+      expect([...byId.keys()].sort()).toEqual(["a2", "u1", "u2"]);
+      expect(byId.get("u2")?.parentId).toBe("u1");
+      expect(byId.get("a2")?.parentId).toBe("u2");
+
+      const restored = new MessageRepository();
+      restored.import(exported);
+      expect(restored.getMessages().map((m) => m.id)).toEqual([
+        "u1",
+        "u2",
+        "a2",
+      ]);
+    });
+
     it("round-trips through export/import without resurrecting the placeholder", () => {
       const parent = createTestMessage({ id: "u" });
       const real = createTestMessage({ id: "a" });

package/src/tests/RemoteThreadListThreadListRuntimeCore-errors.test.ts

@@ -0,0 +1,42 @@
+import { describe, it, expect, vi } from "vitest";
+import { createCore, makeAdapter } from "./remote-thread-list-test-helpers";
+import { InMemoryThreadListAdapter } from "../runtimes/remote-thread-list/adapter/in-memory";
+
+describe("RemoteThreadListThreadListRuntimeCore errors", () => {
+  it("includes the requested thread id when a thread is missing", () => {
+    const core = createCore(makeAdapter());
+
+    expect(() => core.rename("missing-thread", "New title")).toThrow(
+      'Thread "missing-thread" not found while renaming it.',
+    );
+  });
+
+  it("includes the requested thread id and status when an operation is invalid", async () => {
+    const adapter = makeAdapter({
+      list: vi.fn(async () => ({
+        threads: [
+          {
+            status: "archived" as const,
+            remoteId: "archived-thread",
+            externalId: "archived-thread",
+            title: "Archived",
+          },
+        ],
+      })),
+    });
+    const core = createCore(adapter);
+    await core.getLoadThreadsPromise();
+
+    await expect(core.archive("archived-thread")).rejects.toThrow(
+      'Thread "archived-thread" has status "archived", so it cannot be archived.',
+    );
+  });
+
+  it("includes the requested thread id when the in-memory adapter cannot fetch it", async () => {
+    const adapter = new InMemoryThreadListAdapter();
+
+    await expect(adapter.fetch("missing-thread")).rejects.toThrow(
+      'Thread "missing-thread" not found in in-memory thread list.',
+    );
+  });
+});