@assistant-ui/core 0.3.18 → 0.3.19

package/src/react/primitives/message/MessageGroupedParts.tsx

@@ -8,7 +8,7 @@ import type {
   MessagePartStatus,
   ToolCallMessagePartStatus,
 } from "../../../types/message";
-import { getGroupStatus } from "../../../utils/getGroupStatus";
+import { getGroupSummary } from "../../../utils/getGroupStatus";
 import {
   buildGroupTree,
   GROUPBY_MEMO_KEY,
@@ -18,6 +18,14 @@ import {
 import { MessagePartChildren, type EnrichedPartState } from "./MessageParts";
 
 export namespace MessagePrimitiveGroupedParts {
+  /** Per status tallies over the group's `indices`; they sum to `indices.length`. */
+  export type GroupCounts = {
+    readonly running: number;
+    readonly complete: number;
+    readonly incomplete: number;
+    readonly requiresAction: number;
+  };
+
   /**
    * A coalesced group of adjacent parts. Surfaced through the same
    * `{ part }` channel as a leaf {@link EnrichedPartState} so consumers
@@ -28,6 +36,8 @@ export namespace MessagePrimitiveGroupedParts {
   export type GroupPart<TKey extends `group-${string}` = `group-${string}`> = {
     readonly type: TKey;
     readonly status: MessagePartStatus | ToolCallMessagePartStatus;
+    /** Per status tallies over `indices`. */
+    readonly counts: GroupCounts;
     readonly indices: readonly number[];
   };
 
@@ -85,10 +95,10 @@ export namespace MessagePrimitiveGroupedParts {
      * `switch (part.type)` can tell groups apart from real part types.
      *
      * **Prefer {@link groupPartByType}** for the common case of mapping by
-     * `part.type` — it ships a stable memo fingerprint so the tree
+     * `part.type` or tool name — it ships a stable memo fingerprint so the tree
      * survives unrelated re-renders. Use an inline function only when
      * the helper isn't expressive enough (e.g. branching on
-     * `part.toolName` or part metadata).
+     * `part.parentId` or part metadata).
      *
      * The second argument is a {@link GroupByContext} carrying the tool-UI
      * registry, for grouping that depends on it (e.g. standalone tool calls).
@@ -194,10 +204,11 @@ const renderNode = <TKey extends `group-${string}`>(
     );
   }
 
-  const status = getGroupStatus(parts, node.indices);
+  const { status, counts } = getGroupSummary(parts, node.indices);
   const groupPart: MessagePrimitiveGroupedParts.GroupPart<TKey> = {
     type: node.key as TKey,
     status,
+    counts,
     indices: node.indices,
   };
 

package/src/react/primitives/part/PartMessages.test.tsx

@@ -1,4 +1,13 @@
-import { expectTypeOf } from "vitest";
+// @vitest-environment jsdom
+
+import { cleanup, render, screen } from "@testing-library/react";
+import { useAuiState } from "@assistant-ui/store";
+import { afterEach, describe, expect, expectTypeOf, it } from "vitest";
+import type { ThreadMessage } from "../../../types/message";
+import { AssistantRuntimeProvider } from "../../AssistantRuntimeProvider";
+import { useExternalStoreRuntime } from "../../runtimes/useExternalStoreRuntime";
+import { MessagePrimitiveParts } from "../message/MessageParts";
+import { ThreadPrimitiveMessages } from "../thread/ThreadMessages";
 import { PartPrimitiveMessages } from "./PartMessages";
 
 const checkPartMessagesProps = () => {
@@ -15,3 +24,97 @@ const checkPartMessagesProps = () => {
   </PartPrimitiveMessages>;
 };
 expectTypeOf(checkPartMessagesProps).toEqualTypeOf<() => void>();
+
+const nestedAssistant: ThreadMessage = {
+  id: "nested-assistant",
+  createdAt: new Date(0),
+  role: "assistant",
+  content: [],
+  status: { type: "running" },
+  metadata: {
+    unstable_state: {},
+    unstable_annotations: [],
+    unstable_data: [],
+    steps: [],
+    custom: {},
+  },
+};
+
+const NestedThreadStatus = () => {
+  const isRunning = useAuiState((s) => s.thread.isRunning);
+  return (
+    <output data-testid="nested-thread-running">{String(isRunning)}</output>
+  );
+};
+
+const OuterMessageStatus = () => {
+  const status = useAuiState((s) => s.message.status?.type);
+  return <output data-testid="outer-message-status">{status}</output>;
+};
+
+afterEach(cleanup);
+
+describe("PartPrimitiveMessages", () => {
+  it("keeps the outer message running and reports the nested thread as running after the outer run ends", () => {
+    const App = () => {
+      const runtime = useExternalStoreRuntime({
+        messages: [
+          {
+            id: "outer-assistant",
+            createdAt: new Date(0),
+            role: "assistant" as const,
+            content: [
+              {
+                type: "tool-call" as const,
+                toolCallId: "delegate-1",
+                toolName: "delegate",
+                args: {},
+                argsText: "",
+                messages: [nestedAssistant],
+              },
+            ],
+            metadata: {
+              unstable_state: {},
+              unstable_annotations: [],
+              unstable_data: [],
+              steps: [],
+              custom: {},
+            },
+          },
+        ],
+        isRunning: false,
+        convertMessage: (message) => message,
+        onNew: async () => {},
+      });
+      return (
+        <AssistantRuntimeProvider runtime={runtime}>
+          <ThreadPrimitiveMessages>
+            {() => (
+              <>
+                <OuterMessageStatus />
+                <MessagePrimitiveParts>
+                  {({ part }) =>
+                    part.type === "tool-call" ? (
+                      <PartPrimitiveMessages>
+                        {() => <NestedThreadStatus />}
+                      </PartPrimitiveMessages>
+                    ) : null
+                  }
+                </MessagePrimitiveParts>
+              </>
+            )}
+          </ThreadPrimitiveMessages>
+        </AssistantRuntimeProvider>
+      );
+    };
+
+    render(<App />);
+
+    expect(screen.getByTestId("outer-message-status").textContent).toBe(
+      "running",
+    );
+    expect(screen.getByTestId("nested-thread-running").textContent).toBe(
+      "true",
+    );
+  });
+});

package/src/react/primitives/part/PartMessages.tsx

@@ -36,6 +36,7 @@ const usePartMessages = (): readonly ThreadMessage[] | undefined => {
  * This primitive reads `messages` from the current tool call part in the PartScope
  * and renders them using a readonly thread context. All existing message and part
  * primitives work inside, and parent tool UI registrations are inherited.
+ * The nested thread reports `isRunning` while its trailing assistant message is running.
  *
  * @example
  * ```tsx

package/src/react/primitives/thread/ThreadMessages.test.tsx

@@ -0,0 +1,142 @@
+// @vitest-environment jsdom
+
+import { cleanup, render, screen, waitFor } from "@testing-library/react";
+import { AuiProvider, useAui } from "@assistant-ui/store";
+import { type FC, useState } from "react";
+import { afterEach, describe, expect, it } from "vitest";
+import {
+  ExternalThread,
+  type ExternalThreadProps,
+} from "../../../store/clients/external-thread";
+import type { MessageState } from "../../../store";
+import { AssistantRuntimeProvider } from "../../AssistantRuntimeProvider";
+import { useExternalStoreRuntime } from "../../runtimes/useExternalStoreRuntime";
+import { ThreadPrimitiveMessages } from "./ThreadMessages";
+
+const message = (id: string, role: "user" | "assistant") =>
+  ({
+    id,
+    createdAt: new Date(0),
+    role,
+    content: [{ type: "text", text: id }],
+    attachments: [],
+    metadata: {
+      custom: {},
+    },
+    ...(role === "assistant" ? { status: { type: "complete" } } : undefined),
+  }) as ExternalThreadProps["messages"][number];
+
+const StatefulMessage = ({ message }: { message: MessageState }) => {
+  const [initialId] = useState(message.id);
+  return <span data-role={message.role}>{`${initialId}:${message.id}`}</span>;
+};
+
+const renderMessages = (messages: ExternalThreadProps["messages"]) => {
+  const App: FC<{ messages: ExternalThreadProps["messages"] }> = ({
+    messages,
+  }) => {
+    const aui = useAui({ thread: ExternalThread({ messages }) });
+    return (
+      <AuiProvider value={aui}>
+        <ThreadPrimitiveMessages>
+          {({ message }) => <StatefulMessage message={message} />}
+        </ThreadPrimitiveMessages>
+      </AuiProvider>
+    );
+  };
+  const result = render(<App messages={messages} />);
+  return {
+    rerender: (nextMessages: ExternalThreadProps["messages"]) =>
+      result.rerender(<App messages={nextMessages} />),
+  };
+};
+
+const renderExternalStoreMessages = (
+  messages: ExternalThreadProps["messages"],
+  isRunning: boolean,
+) => {
+  const App: FC<{
+    messages: ExternalThreadProps["messages"];
+    isRunning: boolean;
+  }> = ({ messages, isRunning }) => {
+    const runtime = useExternalStoreRuntime({
+      messages,
+      isRunning,
+      convertMessage: (message) => message,
+      onNew: async () => {},
+    });
+    return (
+      <AssistantRuntimeProvider runtime={runtime}>
+        <ThreadPrimitiveMessages>
+          {({ message }) => <StatefulMessage message={message} />}
+        </ThreadPrimitiveMessages>
+      </AssistantRuntimeProvider>
+    );
+  };
+  const result = render(<App messages={messages} isRunning={isRunning} />);
+  return {
+    container: result.container,
+    rerender: (
+      nextMessages: ExternalThreadProps["messages"],
+      nextIsRunning: boolean,
+    ) =>
+      result.rerender(
+        <App messages={nextMessages} isRunning={nextIsRunning} />,
+      ),
+  };
+};
+
+describe("ThreadPrimitiveMessages", () => {
+  afterEach(cleanup);
+
+  it("keeps component state with the surviving message after removal", () => {
+    const first = message("first", "user");
+    const second = message("second", "assistant");
+    const thread = renderMessages([first, second]);
+
+    thread.rerender([second]);
+
+    expect(screen.queryByText("second:second")).not.toBeNull();
+    expect(screen.queryByText("first:second")).toBeNull();
+  });
+
+  it("keeps component state with messages after prepending history", () => {
+    const first = message("first", "user");
+    const second = message("second", "assistant");
+    const thread = renderMessages([first, second]);
+
+    thread.rerender([message("zeroth", "assistant"), first, second]);
+
+    expect(screen.queryByText("first:first")).not.toBeNull();
+    expect(screen.queryByText("second:second")).not.toBeNull();
+    expect(screen.queryByText("first:zeroth")).toBeNull();
+    expect(screen.queryByText("second:first")).toBeNull();
+  });
+
+  it("remounts a message when its runtime ID changes", () => {
+    const thread = renderMessages([message("client-id", "assistant")]);
+
+    thread.rerender([message("server-id", "assistant")]);
+
+    expect(screen.queryByText("server-id:server-id")).not.toBeNull();
+    expect(screen.queryByText("client-id:server-id")).toBeNull();
+  });
+
+  it("remounts an optimistic placeholder as the real assistant message", async () => {
+    const userMessage = message("user-id", "user");
+    const thread = renderExternalStoreMessages([userMessage], true);
+    const optimisticRow = thread.container.querySelector(
+      '[data-role="assistant"]',
+    );
+
+    expect(optimisticRow).not.toBeNull();
+    expect(optimisticRow?.textContent).not.toBe("assistant-id:assistant-id");
+
+    thread.rerender([userMessage, message("assistant-id", "assistant")], false);
+
+    await waitFor(() => {
+      expect(screen.queryByText("assistant-id:assistant-id")).not.toBeNull();
+    });
+    expect(optimisticRow?.isConnected).toBe(false);
+  });
+});

package/src/react/primitives/thread/ThreadMessages.tsx

@@ -6,6 +6,7 @@ import {
   useMemo,
 } from "react";
 import { RenderChildrenWithAccessor, useAuiState } from "@assistant-ui/store";
+import { useShallowSelector } from "@assistant-ui/store/internal";
 import { MessageByIndexProvider } from "../../providers/MessageByIndexProvider";
 import { MessageByIdProvider } from "../../providers/MessageByIdProvider";
 import type { MessageState } from "../../../store";
@@ -241,12 +242,14 @@ ThreadPrimitiveUnstable_MessageById.displayName =
 const ThreadPrimitiveMessagesInner: FC<{
   children: (value: { message: MessageState }) => ReactNode;
 }> = ({ children }) => {
-  const messagesLength = useAuiState((s) => s.thread.messages.length);
+  const messageIds = useAuiState(
+    useShallowSelector((s) => s.thread.messages.map((message) => message.id)),
+  );
 
   return useMemo(() => {
-    if (messagesLength === 0) return null;
-    return Array.from({ length: messagesLength }, (_, index) => (
-      <MessageByIndexProvider key={index} index={index}>
+    if (messageIds.length === 0) return null;
+    return messageIds.map((messageId, index) => (
+      <MessageByIndexProvider key={messageId} index={index}>
         <RenderChildrenWithAccessor
           getItemState={(aui) => aui.thread.message({ index }).getState()}
         >
@@ -260,7 +263,7 @@ const ThreadPrimitiveMessagesInner: FC<{
         </RenderChildrenWithAccessor>
       </MessageByIndexProvider>
     ));
-  }, [messagesLength, children]);
+  }, [messageIds, children]);
 };
 
 /**

package/src/react/providers/ReadonlyThreadProvider.test.tsx

@@ -0,0 +1,56 @@
+// @vitest-environment jsdom
+
+import { cleanup, render, screen, waitFor } from "@testing-library/react";
+import { useAuiState } from "@assistant-ui/store";
+import { afterEach, describe, expect, it } from "vitest";
+import type { ThreadMessage } from "../../types/message";
+import { ReadonlyThreadProvider } from "./ReadonlyThreadProvider";
+
+const messages = (status: "running" | "complete"): readonly ThreadMessage[] => [
+  {
+    id: "nested-assistant",
+    createdAt: new Date(0),
+    role: "assistant",
+    content: [],
+    status:
+      status === "running"
+        ? { type: "running" }
+        : { type: "complete", reason: "unknown" },
+    metadata: {
+      unstable_state: {},
+      unstable_annotations: [],
+      unstable_data: [],
+      steps: [],
+      custom: {},
+    },
+  },
+];
+
+const ThreadStatus = () => {
+  const isRunning = useAuiState((s) => s.thread.isRunning);
+  return <output data-testid="thread-running">{String(isRunning)}</output>;
+};
+
+afterEach(cleanup);
+
+describe("ReadonlyThreadProvider", () => {
+  it("reports isRunning from the trailing assistant message", async () => {
+    const result = render(
+      <ReadonlyThreadProvider messages={messages("running")}>
+        <ThreadStatus />
+      </ReadonlyThreadProvider>,
+    );
+
+    expect(screen.getByTestId("thread-running").textContent).toBe("true");
+
+    result.rerender(
+      <ReadonlyThreadProvider messages={messages("complete")}>
+        <ThreadStatus />
+      </ReadonlyThreadProvider>,
+    );
+
+    await waitFor(() => {
+      expect(screen.getByTestId("thread-running").textContent).toBe("false");
+    });
+  });
+});

package/src/react/runtimes/RemoteThreadListThreadListRuntimeCore.tsx

@@ -58,6 +58,7 @@ import { useAui } from "@assistant-ui/store";
 import type { ModelContextProvider } from "../../model-context/types";
 import { RuntimeAdapterProvider } from "./RuntimeAdapterProvider";
 import { useStableRuntimeAdapters } from "./useRuntimeAdapters";
+import { invokeUserCallback } from "../../utils/invoke-user-callback";
 
 const threadNotFoundError = (threadIdOrRemoteId: string, action: string) =>
   new Error(`Thread "${threadIdOrRemoteId}" not found while ${action}.`);
@@ -573,7 +574,12 @@ export class RemoteThreadListThreadListRuntimeCore
     if (this._lastNotifiedThreadId === threadId) return;
     this._lastNotifiedThreadId = threadId;
     if (emit) {
-      this._options.onThreadIdChange?.(threadId);
+      invokeUserCallback(
+        "assistant-ui",
+        "onThreadIdChange",
+        this._options.onThreadIdChange,
+        threadId,
+      );
     }
   }
 

package/src/react/runtimes/cloud/AssistantCloudThreadHistoryAdapter.test.ts

@@ -28,8 +28,8 @@ describe("extractAuiV0", () => {
   it("reports a terminal message once with its usage", () => {
     const result = extractAuiV0(auiV0Message({ type: "complete" }));
     expect(result?.status).toBe("completed");
-    expect(result?.inputTokens).toBe(1);
-    expect(result?.outputTokens).toBe(2);
+    expect(result?.usage?.inputTokens).toBe(1);
+    expect(result?.usage?.outputTokens).toBe(2);
   });
 
   it("sums step usage reported under the AI SDK v7 token details", () => {
@@ -57,22 +57,28 @@ describe("extractAuiV0", () => {
       },
     });
 
-    expect(result?.inputTokens).toBe(15);
-    expect(result?.outputTokens).toBe(6);
-    expect(result?.cachedInputTokens).toBe(9);
-    expect(result?.reasoningTokens).toBe(2);
+    expect(result?.usage).toEqual({
+      inputTokens: 15,
+      outputTokens: 6,
+      cachedInputTokens: 9,
+      reasoningTokens: 2,
+    });
     expect(result?.steps).toEqual([
       {
-        input_tokens: 10,
-        output_tokens: 4,
-        reasoning_tokens: 1,
-        cached_input_tokens: 6,
+        usage: {
+          inputTokens: 10,
+          outputTokens: 4,
+          inputTokenDetails: { cacheReadTokens: 6 },
+          outputTokenDetails: { reasoningTokens: 1 },
+        },
       },
       {
-        input_tokens: 5,
-        output_tokens: 2,
-        reasoning_tokens: 1,
-        cached_input_tokens: 3,
+        usage: {
+          inputTokens: 5,
+          outputTokens: 2,
+          inputTokenDetails: { cacheReadTokens: 3 },
+          outputTokenDetails: { reasoningTokens: 1 },
+        },
       },
     ]);
   });