@assistant-ui/core 0.3.18 → 0.3.19

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

@@ -1,9 +1,94 @@
 import { describe, it, expect } from "vitest";
-import { convertExternalMessages } from "./external-message-converter";
+import {
+  convertExternalMessages,
+  createExternalMessageConversionCache,
+} from "./external-message-converter";
 import type { useExternalMessageConverter } from "./external-message-converter";
 import { isErrorMessageId } from "../../utils/id";
 
 describe("convertExternalMessages", () => {
+  it("keeps unchanged messages and rebuilds a replaced one with an opt-in cache", () => {
+    type Input = { id: string; role: "user" | "assistant"; text: string };
+    const metadata = {};
+    const cache = createExternalMessageConversionCache();
+    const callback: useExternalMessageConverter.Callback<Input> = (
+      message,
+    ) => ({
+      id: message.id,
+      role: message.role,
+      content: message.text,
+    });
+    const question: Input = { id: "question", role: "user", text: "hi" };
+    const answer: Input = { id: "answer", role: "assistant", text: "hel" };
+
+    const first = convertExternalMessages(
+      [question, answer],
+      callback,
+      true,
+      metadata,
+      cache,
+    );
+    const second = convertExternalMessages(
+      [question, { ...answer, text: "hello" }],
+      callback,
+      true,
+      metadata,
+      cache,
+    );
+
+    expect(second[0]).toBe(first[0]);
+    expect(second[1]).not.toBe(first[1]);
+    expect(second[1]?.content).toMatchObject([{ type: "text", text: "hello" }]);
+  });
+
+  it("invalidates cached messages when the callback or metadata changes", () => {
+    const input = { text: "nested" };
+    const cache = createExternalMessageConversionCache();
+    const withFormat =
+      (
+        format: (text: string) => string,
+      ): useExternalMessageConverter.Callback<typeof input> =>
+      (message, metadata) => ({
+        id: "nested",
+        role: "assistant",
+        content: format(message.text),
+        metadata: { timing: metadata.messageTiming?.["nested"] },
+      });
+    const callback = withFormat((text) => text);
+    const uppercaseCallback = withFormat((text) => text.toUpperCase());
+    const metadata = {};
+    const timing = { streamStartTime: 1, totalChunks: 1, toolCallCount: 0 };
+
+    const first = convertExternalMessages(
+      [input],
+      callback,
+      false,
+      metadata,
+      cache,
+    );
+    const second = convertExternalMessages(
+      [input],
+      uppercaseCallback,
+      false,
+      metadata,
+      cache,
+    );
+    const third = convertExternalMessages(
+      [input],
+      uppercaseCallback,
+      false,
+      { messageTiming: { nested: timing } },
+      cache,
+    );
+
+    expect(second[0]).not.toBe(first[0]);
+    expect(second[0]?.content).toMatchObject([
+      { type: "text", text: "NESTED" },
+    ]);
+    expect(third[0]).not.toBe(second[0]);
+    expect(third[0]?.metadata).toMatchObject({ timing });
+  });
+
   describe("reasoning part merging", () => {
     it("should merge reasoning parts with the same parentId", () => {
       const messages = [
@@ -205,6 +290,56 @@ describe("convertExternalMessages", () => {
       expect((toolCallParts[0] as any).argsText).toBe('{"query":"new"}');
     });
 
+    it("keeps a merged tool call in its original content position", () => {
+      const messages = [
+        {
+          role: "assistant" as const,
+          content: [
+            {
+              type: "tool-call" as const,
+              toolCallId: "tc1",
+              toolName: "search",
+              args: { query: "old" },
+            },
+            { type: "text" as const, text: "after tool" },
+          ],
+        },
+        {
+          role: "assistant" as const,
+          content: [
+            {
+              type: "tool-call" as const,
+              toolCallId: "tc1",
+              toolName: "search",
+              args: { query: "new" },
+            },
+          ],
+        },
+        {
+          role: "tool" as const,
+          toolCallId: "tc1",
+          result: "found",
+        },
+      ];
+
+      const result = convertExternalMessages(
+        messages,
+        (message) => message,
+        false,
+        {},
+      );
+
+      expect(result[0]!.content).toMatchObject([
+        {
+          type: "tool-call",
+          toolCallId: "tc1",
+          args: { query: "new" },
+          result: "found",
+        },
+        { type: "text", text: "after tool" },
+      ]);
+    });
+
     it("should ignore orphaned tool results without throwing", () => {
       const messages = [
         {

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

@@ -2,28 +2,29 @@
 
 import { useMemo, useState } from "react";
 import {
-  chunkExternalMessages,
-  completeExternalMessageConversion,
-  convertExternalMessageCallback,
-  convertExternalMessageChunk,
   convertExternalMessages as convertExternalMessagesInternal,
-  shallowArrayEqual,
+  createExternalMessageConversionCache as createExternalMessageConversionCacheInternal,
   type ExternalMessageConverterCallback,
-  type ExternalMessageConverterCallbackResult,
-  type ExternalMessageConverterChunk,
   type ExternalMessageConverterMessage,
   type ExternalMessageConverterMetadata,
+  type InternalExternalMessageConversionCache,
   type JoinStrategy,
 } from "../../runtime/utils/external-message-conversion";
 import { bindExternalStoreMessage } from "../../runtime/utils/external-store-message";
-import { ThreadMessageConverter } from "../../runtimes/external-store/thread-message-converter";
-import type { ThreadMessage } from "../../types/message";
-
-// Generatedness is tracked by identity, not by id shape: a caller-supplied id
-// that happens to match the generated pattern must never be rewritten.
-const generatedFallbackMessages = new WeakSet<object>();
 
 export type { JoinStrategy };
+export type ExternalMessageConversionCache = {
+  readonly __brand: unique symbol;
+};
+
+/**
+ * Creates a cache for plain external message conversion. Pass the same cache on every call for one message list, and a source message that has not changed since the previous call converts to the same `ThreadMessage` object.
+ *
+ * Entries are keyed by source message identity and rebuilt whenever `callback` or `metadata` is a different object than on the previous call, so keep both referentially stable.
+ */
+export const createExternalMessageConversionCache =
+  (): ExternalMessageConversionCache =>
+    createExternalMessageConversionCacheInternal() as unknown as ExternalMessageConversionCache;
 
 export namespace useExternalMessageConverter {
   export type Message = ExternalMessageConverterMessage;
@@ -31,17 +32,21 @@ export namespace useExternalMessageConverter {
   export type Callback<T> = ExternalMessageConverterCallback<T>;
 }
 
-export const convertExternalMessages: <T extends WeakKey>(
+export const convertExternalMessages = <T extends WeakKey>(
   messages: T[],
   callback: useExternalMessageConverter.Callback<T>,
   isRunning: boolean,
   metadata: useExternalMessageConverter.Metadata,
-) => ThreadMessage[] = convertExternalMessagesInternal;
-
-type CallbackCacheEntry<T> = ExternalMessageConverterCallbackResult<T> & {
-  metadata: useExternalMessageConverter.Metadata;
-  callback: useExternalMessageConverter.Callback<T>;
-};
+  cache?: ExternalMessageConversionCache,
+) =>
+  convertExternalMessagesInternal(
+    messages,
+    callback,
+    isRunning,
+    metadata,
+    undefined,
+    cache as unknown as InternalExternalMessageConversionCache<T> | undefined,
+  );
 
 export const useExternalMessageConverter = <T extends WeakKey>({
   callback,
@@ -56,89 +61,30 @@ 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<
-      ExternalMessageConverterMessage,
-      ExternalMessageConverterChunk<T>
-    >(),
-    converterCache: new ThreadMessageConverter(),
-  }));
+  // The cache lives for the component lifetime: React Compiler hoists allocations without reactive dependencies out of useMemo, so entries carry the callback and metadata that produced them instead of being flushed when those change.
+  const [cache] = useState(() =>
+    createExternalMessageConversionCacheInternal<T>(),
+  );
 
   const state = useMemo(
     () => ({
       metadata: metadata ?? {},
       callback,
-      ...caches,
     }),
-    [callback, metadata, caches],
+    [callback, metadata],
   );
 
   return useMemo(() => {
-    const callbackResults: ExternalMessageConverterCallbackResult<T>[] = [];
-    for (const message of messages) {
-      let result = state.callbackCache.get(message);
-      if (
-        !result ||
-        result.metadata !== state.metadata ||
-        result.callback !== state.callback
-      ) {
-        result = {
-          ...convertExternalMessageCallback(
-            message,
-            state.callback,
-            state.metadata,
-          ),
-          metadata: state.metadata,
-          callback: state.callback,
-        };
-        state.callbackCache.set(message, result);
-      }
-      callbackResults.push(result);
-    }
-
-    const chunks = chunkExternalMessages(callbackResults, joinStrategy).map(
-      (message) => {
-        const key = message.outputs[0];
-        if (!key) return message;
-
-        const cached = state.chunkCache.get(key);
-        if (cached && shallowArrayEqual(cached.outputs, message.outputs)) {
-          return cached;
-        }
-        state.chunkCache.set(key, message);
-        return message;
-      },
-    );
-
-    const threadMessages = state.converterCache.convertMessages(
-      chunks,
-      (cache, message, idx) =>
-        convertExternalMessageChunk(
-          message,
-          idx,
-          chunks.length,
-          isRunning,
-          state.metadata.error,
-          {
-            message: cache,
-            generatedFallbackMessages,
-          },
-          state.metadata.cancelledMessageIds,
-        ),
+    const threadMessages = convertExternalMessagesInternal(
+      messages,
+      state.callback,
+      isRunning,
+      state.metadata,
+      joinStrategy,
+      cache,
     );
 
     bindExternalStoreMessage(threadMessages, messages);
-    return completeExternalMessageConversion(
-      threadMessages,
-      state.metadata.error,
-    );
-  }, [state, messages, isRunning, joinStrategy]);
+    return threadMessages;
+  }, [state, messages, isRunning, joinStrategy, cache]);
 };

package/src/react/runtimes/useRemoteThreadListRuntime.ts

@@ -109,7 +109,7 @@ export const useRemoteThreadListRuntime = (
   );
 
   const onThreadIdChange = useEffectEvent((threadId: string | undefined) => {
-    options.onThreadIdChange?.(threadId);
+    return options.onThreadIdChange?.(threadId);
   });
 
   const stableOptions = useMemo<RemoteThreadListOptions>(

package/src/react/utils/groupParts.ts

@@ -47,8 +47,13 @@ export const GROUPBY_MEMO_KEY: unique symbol = Symbol.for(
  *   `display: "standalone"`). Resolving the registry-driven cases reads the
  *   {@link GroupByContext} passed to the `groupBy` function. Takes precedence
  *   over the `"tool-call"` entry.
+ *
+ * A `"tool-call:<name>"` entry matches tool calls by tool name and takes precedence over the plain `"tool-call"` entry, while `"standalone-tool-call"` still wins first.
  */
-type GroupPartType = PartState["type"] | "standalone-tool-call";
+type GroupPartType =
+  | PartState["type"]
+  | "standalone-tool-call"
+  | `tool-call:${string}`;
 
 /**
  * Build a `groupBy` from a `part.type → group-key path` lookup.
@@ -63,11 +68,14 @@ type GroupPartType = PartState["type"] | "standalone-tool-call";
  * {@link GroupByContext} that `<MessagePrimitive.GroupedParts>` passes to the
  * `groupBy` function — the helper needs nothing threaded into it.
  *
+ * A `"tool-call:<name>"` entry matches tool calls by tool name and takes precedence over the plain `"tool-call"` entry, while `"standalone-tool-call"` still wins first.
+ *
  * @example
  * ```tsx
  * <MessagePrimitive.GroupedParts
  *   groupBy={groupPartByType({
  *     reasoning: ["group-thought", "group-reasoning"],
+ *     "tool-call:task": ["group-subagents"],
  *     "tool-call": ["group-thought", "group-tool"],
  *     "standalone-tool-call": [],
  *   })}
@@ -91,6 +99,9 @@ export const groupPartByType = <TKey extends `group-${string}`>(
       if (isStandalone && lookup["standalone-tool-call"] !== undefined) {
         return lookup["standalone-tool-call"]!;
       }
+      const toolCallPath = lookup[`tool-call:${part.toolName}`];
+      if (toolCallPath !== undefined) return toolCallPath;
+      return lookup["tool-call"] ?? [];
     }
     return lookup[part.type] ?? [];
   }) as ((part: PartState, context?: GroupByContext) => readonly TKey[]) & {

package/src/react/utils/useJSONEqual.ts

@@ -1,6 +1,15 @@
 import { useRef } from "react";
+import type { ReadonlyJSONValue } from "assistant-stream/utils";
 import { isJSONValueEqual } from "../../utils/json/is-json-equal";
 
+export function useJSONEqualValue<T extends ReadonlyJSONValue>(value: T): T {
+  const prev = useRef(value);
+  if (prev.current !== value && !isJSONValueEqual(prev.current, value)) {
+    prev.current = value;
+  }
+  return prev.current;
+}
+
 /**
  * Like `useShallow`, but with JSON deep-equality. Use when a selector derives an
  * equal-but-fresh value on every store update — e.g. folding over

package/src/react/utils/useJSONSchemaDependency.ts

@@ -0,0 +1,22 @@
+import { useMemo } from "react";
+import { toJSONSchema } from "assistant-stream";
+import type { ReadonlyJSONValue } from "assistant-stream/utils";
+import { isJSONValue } from "../../utils/json/is-json";
+import type { Unstable_InteractableStateSchema } from "../types/scopes/interactables";
+import { useJSONEqualValue } from "./useJSONEqual";
+
+export function useJSONSchemaDependency(
+  schema: Unstable_InteractableStateSchema,
+): ReadonlyJSONValue {
+  const normalizedSchema = useMemo(() => {
+    try {
+      const jsonSchema = toJSONSchema(schema);
+      return isJSONValue(jsonSchema) ? jsonSchema : null;
+    } catch {
+      // Opaque schemas share one fallback so rebuilt instances do not churn registration.
+      return null;
+    }
+  }, [schema]);
+
+  return useJSONEqualValue(normalizedSchema);
+}

package/src/runtime/api/thread-runtime.test.ts

@@ -1,7 +1,13 @@
 import { describe, expect, it, vi } from "vitest";
 import { LocalRuntimeCore } from "../../runtimes/local/local-runtime-core";
 import { ExternalStoreRuntimeCore } from "../../runtimes/external-store/external-store-runtime-core";
+import { ReadonlyThreadRuntimeCore } from "../../runtimes/readonly/ReadonlyThreadRuntimeCore";
 import { AssistantRuntimeImpl } from "./assistant-runtime";
+import {
+  ThreadRuntimeImpl,
+  type ThreadListItemRuntimeBinding,
+  type ThreadRuntimeCoreBinding,
+} from "./thread-runtime";
 
 describe("ThreadRuntime.append", () => {
   it.each([
@@ -76,3 +82,55 @@ describe("ThreadRuntime.append with an external store", () => {
     expect(onNew).not.toHaveBeenCalled();
   });
 });
+
+describe("ThreadRuntime state subscriptions", () => {
+  it("tears down every source before reconnecting after an error", () => {
+    const core = new ReadonlyThreadRuntimeCore();
+    const cleanupError = new Error("thread cleanup failed");
+    const threadCleanup = vi.fn(() => {
+      throw cleanupError;
+    });
+    const itemCleanup = vi.fn();
+    let threadSubscriptions = 0;
+    let itemSubscriptions = 0;
+    const path = {
+      ref: "test.thread",
+      threadSelector: { type: "main" as const },
+    };
+    const runtime = new ThreadRuntimeImpl(
+      {
+        path,
+        getState: () => core,
+        subscribe: () => {
+          threadSubscriptions += 1;
+          return threadCleanup;
+        },
+        outerSubscribe: () => () => {},
+      } satisfies ThreadRuntimeCoreBinding,
+      {
+        path,
+        getState: () => ({
+          id: "test",
+          remoteId: undefined,
+          externalId: undefined,
+          isMain: true,
+          isRunning: false,
+          status: "regular",
+          title: undefined,
+        }),
+        subscribe: () => {
+          itemSubscriptions += 1;
+          return itemCleanup;
+        },
+      } satisfies ThreadListItemRuntimeBinding,
+    );
+
+    const unsubscribe = runtime.subscribe(() => {});
+    expect(() => unsubscribe()).toThrow(cleanupError);
+    expect(itemCleanup).toHaveBeenCalledOnce();
+
+    runtime.subscribe(() => {});
+    expect(threadSubscriptions).toBe(2);
+    expect(itemSubscriptions).toBe(2);
+  });
+});

package/src/runtime/api/thread-runtime.ts

@@ -18,6 +18,7 @@ import {
 } from "./message-runtime";
 import { NestedSubscriptionSubject } from "../../subscribable/subscribable";
 import {
+  runCleanups,
   ShallowMemoizeSubject,
   SKIP_UPDATE,
 } from "../../subscribable/subscribable";
@@ -372,10 +373,7 @@ export class ThreadRuntimeImpl implements ThreadRuntime {
       subscribe: (callback) => {
         const sub1 = threadBinding.subscribe(callback);
         const sub2 = threadListItemBinding.subscribe(callback);
-        return () => {
-          sub1();
-          sub2();
-        };
+        return () => runCleanups([sub1, sub2]);
       },
     });
 

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

@@ -167,7 +167,7 @@ export abstract class BaseComposerRuntimeCore
     const adapter = this.getAttachmentAdapter();
     if (adapter) {
       const pending = this._attachments.filter((a) => !isAttachmentComplete(a));
-      await Promise.all(pending.map((a) => adapter.remove(a)));
+      await Promise.all(pending.map(async (a) => adapter.remove(a)));
     }
   }
 
@@ -203,6 +203,10 @@ export abstract class BaseComposerRuntimeCore
 
   public async clearAttachments() {
     this._cancelAllAttachmentAdds();
+    if (this._isSending) {
+      for (const attachment of this._attachments)
+        this._removedDuringSend.add(attachment.id);
+    }
     const task = this._onClearAttachments();
     this.setAttachments([]);
 
@@ -213,8 +217,13 @@ export abstract class BaseComposerRuntimeCore
     if (!this.canSend || this._isSending) return;
 
     if (this._dictationSession) {
-      this._dictationSession.cancel();
-      this._cleanupDictation();
+      try {
+        this._dictationSession.cancel();
+      } catch (error) {
+        console.error("[assistant-ui] Dictation session cancel threw", error);
+      } finally {
+        this._cleanupDictation();
+      }
     }
 
     const adapter = this.getAttachmentAdapter();
@@ -636,14 +645,11 @@ export abstract class BaseComposerRuntimeCore
       throw new Error("Dictation adapter not configured");
     }
 
+    const isReplacing = this._dictationSession !== undefined;
     if (this._dictationSession) {
-      for (const unsub of this._dictationUnsubscribes) {
-        unsub();
-      }
-      this._dictationUnsubscribes = [];
       const oldSession = this._dictationSession;
-      oldSession.stop().catch(() => {});
-      this._dictationSession = undefined;
+      this._cleanupDictation({ notify: false });
+      this._stopDictationSession(oldSession);
     }
 
     const inputDisabled = adapter.disableInputDuringDictation ?? false;
@@ -651,7 +657,22 @@ export abstract class BaseComposerRuntimeCore
     this._dictationBaseText = this._text;
     this._currentInterimText = "";
 
-    const session = adapter.listen();
+    let session: DictationAdapter.Session;
+    try {
+      session = adapter.listen();
+    } catch (error) {
+      if (isReplacing) {
+        try {
+          this._notifySubscribers();
+        } catch (notifyError) {
+          console.error(
+            "[assistant-ui] Dictation replacement rollback notification threw",
+            notifyError,
+          );
+        }
+      }
+      throw error;
+    }
     this._dictationSession = session;
     const sessionId = ++this._dictationSessionIdCounter;
     this._activeDictationSessionId = sessionId;
@@ -729,27 +750,59 @@ export abstract class BaseComposerRuntimeCore
     const session = this._dictationSession;
     const sessionId = this._activeDictationSessionId;
     const cleanup = () => this._cleanupDictation({ sessionId });
-    void session.stop().then(cleanup, cleanup);
+    this._stopDictationSession(session, cleanup);
   }
 
-  private _cleanupDictation(options?: { sessionId: number | undefined }): void {
+  private _stopDictationSession(
+    session: DictationAdapter.Session,
+    onSettled: () => void = () => {},
+  ): void {
+    let task: Promise<void>;
+    try {
+      task = session.stop();
+    } catch (error) {
+      console.error("[assistant-ui] Dictation session stop threw", error);
+      onSettled();
+      return;
+    }
+
+    void task.then(onSettled, (error) => {
+      console.error("[assistant-ui] Dictation session stop rejected", error);
+      onSettled();
+    });
+  }
+
+  private _cleanupDictation(options?: {
+    sessionId?: number | undefined;
+    notify?: boolean | undefined;
+  }): void {
     const isStaleSession =
       options?.sessionId !== undefined &&
       options.sessionId !== this._activeDictationSessionId;
     if (isStaleSession || this._isCleaningDictation) return;
 
     this._isCleaningDictation = true;
-    try {
-      for (const unsub of this._dictationUnsubscribes) {
-        unsub();
+    const runCleanup = (cleanup: () => void) => {
+      try {
+        cleanup();
+      } catch (error) {
+        console.error("[assistant-ui] Dictation cleanup threw", error);
       }
+    };
+
+    try {
+      const unsubscribes = this._dictationUnsubscribes;
       this._dictationUnsubscribes = [];
       this._dictationSession = undefined;
       this._activeDictationSessionId = undefined;
       this._dictation = undefined;
       this._dictationBaseText = "";
       this._currentInterimText = "";
-      this._notifySubscribers();
+
+      for (const unsubscribe of unsubscribes) runCleanup(unsubscribe);
+      if (options?.notify !== false) {
+        runCleanup(() => this._notifySubscribers());
+      }
     } finally {
       this._isCleaningDictation = false;
     }