@assistant-ui/core 0.3.18 → 0.3.20

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 = [
         {
@@ -431,7 +566,7 @@ describe("convertExternalMessages", () => {
           role: "assistant" as const,
           content: [{ type: "text" as const, text: "first" }],
           metadata: {
-            timing: { streamStartTime: 1 },
+            timing: { streamStartTime: 1, totalChunks: 1, toolCallCount: 0 },
           },
         },
         {
@@ -439,7 +574,7 @@ describe("convertExternalMessages", () => {
           role: "assistant" as const,
           content: [{ type: "text" as const, text: "second" }],
           metadata: {
-            timing: { streamStartTime: 2 },
+            timing: { streamStartTime: 2, totalChunks: 1, toolCallCount: 0 },
           },
         },
       ];
@@ -452,7 +587,11 @@ describe("convertExternalMessages", () => {
 
       expect(result).toHaveLength(1);
       const metadata = result[0]!.metadata as any;
-      expect(metadata.timing).toEqual({ streamStartTime: 2 });
+      expect(metadata.timing).toEqual({
+        streamStartTime: 2,
+        totalChunks: 1,
+        toolCallCount: 0,
+      });
     });
 
     it("merges custom across joined outputs, with later keys overwriting earlier ones", () => {

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/remote-thread-hook-on-tap.test.ts

@@ -7,11 +7,12 @@ import {
 } from "@assistant-ui/tap";
 import { useExternalStoreRuntime } from "./useExternalStoreRuntime";
 import type { AssistantRuntime } from "../../runtime/api/assistant-runtime";
+import type { ThreadMessage } from "../../types/message";
 
 const EMPTY_MESSAGES: readonly never[] = [];
 
 const useTestThreadRuntime = () =>
-  useExternalStoreRuntime({
+  useExternalStoreRuntime<ThreadMessage>({
     messages: EMPTY_MESSAGES,
     isRunning: false,
     onNew: async () => {},

package/src/react/runtimes/useExternalStoreRuntime.lifecycle.test.tsx

@@ -52,7 +52,10 @@ describe("useExternalStoreRuntime lifecycle", () => {
     const onNew = vi.fn(async () => {});
     const capture: { runtime: AssistantRuntime | null } = { runtime: null };
     const App = () => {
-      const runtime = useExternalStoreRuntime({ messages: [], onNew });
+      const runtime = useExternalStoreRuntime<ThreadMessage>({
+        messages: [],
+        onNew,
+      });
       capture.runtime = runtime;
       return null;
     };
@@ -81,7 +84,10 @@ describe("useExternalStoreRuntime lifecycle", () => {
     const onNew = vi.fn(async () => {});
     const capture: { runtime: AssistantRuntime | null } = { runtime: null };
     const App = () => {
-      const runtime = useExternalStoreRuntime({ messages: [], onNew });
+      const runtime = useExternalStoreRuntime<ThreadMessage>({
+        messages: [],
+        onNew,
+      });
       capture.runtime = runtime;
       return null;
     };

package/src/react/runtimes/useExternalStoreRuntime.suggestions.test.tsx

@@ -10,6 +10,7 @@ import {
   type AssistantClient,
 } from "@assistant-ui/store";
 import type { ThreadSuggestion } from "../../runtime/interfaces/thread-runtime-core";
+import type { ThreadMessage } from "../../types/message";
 import { Suggestions } from "../../store/clients/suggestions";
 import { AssistantRuntimeProvider } from "../AssistantRuntimeProvider";
 import { useExternalStoreRuntime } from "./useExternalStoreRuntime";
@@ -32,13 +33,16 @@ const App = ({
   config?: AuiConfig;
   strict?: boolean;
 }) => {
-  const runtime = useExternalStoreRuntime({
+  const runtime = useExternalStoreRuntime<ThreadMessage>({
     messages: [],
     onNew: async () => {},
     ...(suggestions && { suggestions }),
   });
   const tree = (
-    <AssistantRuntimeProvider runtime={runtime} config={config}>
+    <AssistantRuntimeProvider
+      runtime={runtime}
+      {...(config === undefined ? {} : { config })}
+    >
       <Consumer />
     </AssistantRuntimeProvider>
   );

package/src/react/runtimes/useRemoteThreadListRuntime.concurrent.test.tsx

@@ -1,10 +1,18 @@
 // @vitest-environment jsdom
 
 import { act, render } from "@testing-library/react";
-import { startTransition, Suspense, type ReactNode } from "react";
+import {
+  startTransition,
+  Suspense,
+  useInsertionEffect,
+  useLayoutEffect,
+  useRef,
+  type ReactNode,
+} from "react";
+import { createRoot } from "react-dom/client";
 import { describe, expect, it, vi } from "vitest";
 import type { AssistantRuntime } from "../../runtime/api/assistant-runtime";
-import type { AppendMessage } from "../../types/message";
+import type { AppendMessage, ThreadMessage } from "../../types/message";
 import { makeAdapter } from "../../tests/remote-thread-list-test-helpers";
 import { AssistantRuntimeProvider } from "../AssistantRuntimeProvider";
 import { useExternalStoreRuntime } from "./useExternalStoreRuntime";
@@ -51,7 +59,7 @@ const createHarness = ({ stableHook = false } = {}) => {
   // variant is the only way to exercise an unchanged published hook.
   const useHoistedThreadRuntime = () => {
     renderThreadRuntime();
-    return useExternalStoreRuntime({
+    return useExternalStoreRuntime<ThreadMessage>({
       messages: EMPTY_MESSAGES,
       onNew: onNewA,
     });
@@ -68,7 +76,10 @@ const createHarness = ({ stableHook = false } = {}) => {
       ? useHoistedThreadRuntime
       : () => {
           renderThreadRuntime();
-          return useExternalStoreRuntime({ messages: EMPTY_MESSAGES, onNew });
+          return useExternalStoreRuntime<ThreadMessage>({
+            messages: EMPTY_MESSAGES,
+            onNew,
+          });
         };
     const runtime = useRemoteThreadListRuntime({
       adapter,
@@ -214,3 +225,72 @@ describe("useRemoteThreadListRuntime concurrent options", () => {
     expect(onNewA).not.toHaveBeenCalled();
   });
 });
+
+describe("useRemoteThreadListRuntime commit phase", () => {
+  const LayoutProbe = ({
+    onLayout,
+  }: {
+    onLayout: (() => void) | undefined;
+  }) => {
+    useLayoutEffect(() => {
+      onLayout?.();
+    }, [onLayout]);
+    return null;
+  };
+
+  it("hands a hosted hook its committed options before a yielded commit settles work", async () => {
+    // A render past the scheduler's 5 ms frame budget yields before passive
+    // effects, and act would drain them first, so this renders outside act.
+    vi.stubGlobal("IS_REACT_ACT_ENVIRONMENT", false);
+    try {
+      const adapter = makeAdapter();
+      const settled: string[] = [];
+      let readLabel: (() => string) | undefined;
+      const App = ({
+        label,
+        renderMs,
+        onLayout,
+      }: {
+        label: string;
+        renderMs: number;
+        onLayout?: () => void;
+      }) => {
+        const runtime = useRemoteThreadListRuntime({
+          adapter,
+          runtimeHook: function useLabelThreadRuntime() {
+            const labelRef = useRef(label);
+            useInsertionEffect(() => {
+              labelRef.current = label;
+              readLabel = () => labelRef.current;
+            });
+            return useExternalStoreRuntime<ThreadMessage>({
+              messages: EMPTY_MESSAGES,
+              onNew: async () => {},
+            });
+          },
+        });
+        const renderEnd = performance.now() + renderMs;
+        while (performance.now() < renderEnd) {}
+        return (
+          <AssistantRuntimeProvider runtime={runtime}>
+            <LayoutProbe onLayout={onLayout} />
+          </AssistantRuntimeProvider>
+        );
+      };
+      const settleInMicrotask = () => {
+        queueMicrotask(() => settled.push(readLabel!()));
+      };
+
+      const root = createRoot(document.createElement("div"));
+      root.render(<App label="A" renderMs={0} />);
+      await vi.waitFor(() => expect(readLabel?.()).toBe("A"));
+      root.render(<App label="B" renderMs={30} onLayout={settleInMicrotask} />);
+      await vi.waitFor(() => expect(settled).toHaveLength(1));
+      root.unmount();
+
+      expect(settled).toEqual(["B"]);
+    } finally {
+      vi.unstubAllGlobals();
+    }
+  });
+});

package/src/react/runtimes/useRemoteThreadListRuntime.ts

@@ -2,6 +2,7 @@ import {
   useState,
   useEffect,
   useInsertionEffect,
+  useLayoutEffect,
   useMemo,
   useRef,
   useCallback,
@@ -86,7 +87,10 @@ export const useRemoteThreadListRuntime = (
   const [runtimeHookStore] = useState(
     () => new WritableSubscribable(options.runtimeHook),
   );
-  useEffect(() => {
+  // The layout phase re-renders hosted threads before this commit yields. An
+  // insertion effect cannot notify subscribers, so descendant layout effects
+  // of the same commit still see the previous hook.
+  useLayoutEffect(() => {
     runtimeHookStore.setState(options.runtimeHook);
   }, [runtimeHookStore, options.runtimeHook]);
 
@@ -109,7 +113,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/message-part-runtime.ts

@@ -92,8 +92,11 @@ export class MessagePartRuntimeImpl implements MessagePartRuntime {
       toolName,
       toolCallId,
       result: response.result,
-      artifact: response.artifact,
       isError: response.isError,
+      ...(response.artifact !== undefined && { artifact: response.artifact }),
+      ...(response.modelContent !== undefined && {
+        modelContent: response.modelContent,
+      }),
     });
   }