@assistant-ui/core 0.3.17 → 0.3.19

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

@@ -0,0 +1,96 @@
+// @vitest-environment jsdom
+import { cleanup, render } from "@testing-library/react";
+import { afterEach, describe, expect, it } from "vitest";
+import type { ThreadMessageLike } from "../../../runtime/utils/thread-message-like";
+import { AssistantRuntimeProvider } from "../../AssistantRuntimeProvider";
+import { ThreadPrimitiveMessages } from "../thread/ThreadMessages";
+import { useExternalStoreRuntime } from "../../runtimes/useExternalStoreRuntime";
+import { groupPartByType } from "../../utils/groupParts";
+import { MessagePrimitiveGroupedParts } from "./MessageGroupedParts";
+
+type Msg = {
+  id: string;
+  content: readonly {
+    type: "tool-call";
+    toolCallId: string;
+    toolName: string;
+    args: {};
+    result?: { ok: true };
+  }[];
+};
+
+const task = (toolCallId: string, hasResult: boolean) => ({
+  type: "tool-call" as const,
+  toolCallId,
+  toolName: "task",
+  args: {},
+  ...(hasResult ? { result: { ok: true } } : {}),
+});
+
+const convertMessage = (message: Msg): ThreadMessageLike => ({
+  id: message.id,
+  role: "assistant",
+  content: message.content,
+});
+
+afterEach(cleanup);
+
+describe("MessagePrimitive.GroupedParts", () => {
+  it("passes status counts to a tool-name group", () => {
+    let group:
+      | {
+          counts: MessagePrimitiveGroupedParts.GroupCounts;
+          indices: readonly number[];
+        }
+      | undefined;
+
+    const GroupedParts = () => (
+      <MessagePrimitiveGroupedParts
+        groupBy={groupPartByType({
+          "tool-call:task": ["group-subagents"],
+        })}
+      >
+        {({ part, children }) => {
+          if (part.type === "group-subagents") {
+            group = { counts: part.counts, indices: part.indices };
+            return children;
+          }
+          return null;
+        }}
+      </MessagePrimitiveGroupedParts>
+    );
+
+    const App = () => {
+      const runtime = useExternalStoreRuntime<Msg>({
+        messages: [
+          {
+            id: "assistant-1",
+            content: [
+              task("task-1", true),
+              task("task-2", true),
+              task("task-3", false),
+            ],
+          },
+        ],
+        isRunning: true,
+        convertMessage,
+        onNew: async () => {},
+      });
+      return (
+        <AssistantRuntimeProvider runtime={runtime}>
+          <ThreadPrimitiveMessages components={{ Message: GroupedParts }} />
+        </AssistantRuntimeProvider>
+      );
+    };
+
+    render(<App />);
+
+    expect(group?.counts).toEqual({
+      running: 1,
+      complete: 2,
+      incomplete: 0,
+      requiresAction: 0,
+    });
+    expect(group?.indices).toHaveLength(3);
+  });
+});

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/RemoteThreadListHookInstanceManager.destroy.test.tsx

@@ -0,0 +1,172 @@
+import { describe, expect, it } from "vitest";
+import type { ThreadListRuntimeCore } from "../../runtime/interfaces/thread-list-runtime-core";
+import type { ThreadRuntimeCore } from "../../runtime/interfaces/thread-runtime-core";
+import { RemoteThreadListHookInstanceManager } from "./RemoteThreadListHookInstanceManager";
+
+const makeManager = () =>
+  new RemoteThreadListHookInstanceManager(
+    () => ({}) as never,
+    {} as ThreadListRuntimeCore,
+  );
+
+// No AdapterSink attaches a runtime here, so the start and restart promises
+// stay pending or reject on stop; neither is what these tests are about.
+const start = (manager: RemoteThreadListHookInstanceManager, id: string) => {
+  manager.startThreadRuntime(id).catch(() => {});
+};
+const restart = (manager: RemoteThreadListHookInstanceManager, id: string) => {
+  manager.__internal_restartThreadRuntime(id).catch(() => {});
+};
+
+const publishedSignal = (
+  manager: RemoteThreadListHookInstanceManager,
+  id: string,
+) => {
+  const { hostStore } = manager as unknown as {
+    hostStore: {
+      getState: () => {
+        threads: readonly { id: string; destroySignal: AbortSignal }[];
+      };
+    };
+  };
+  return hostStore.getState().threads.find((thread) => thread.id === id)
+    ?.destroySignal;
+};
+
+describe("RemoteThreadListHookInstanceManager destroy signal", () => {
+  it("publishes a live signal for each started thread", () => {
+    const manager = makeManager();
+    start(manager, "thread-1");
+
+    expect(publishedSignal(manager, "thread-1")?.aborted).toBe(false);
+  });
+
+  it("aborts the thread's signal when its runtime stops", () => {
+    const manager = makeManager();
+    start(manager, "thread-1");
+    const signal = publishedSignal(manager, "thread-1")!;
+
+    manager.stopThreadRuntime("thread-1");
+
+    expect(signal.aborted).toBe(true);
+  });
+
+  it("leaves sibling threads untouched when one stops", () => {
+    const manager = makeManager();
+    start(manager, "thread-1");
+    start(manager, "thread-2");
+    const sibling = publishedSignal(manager, "thread-2")!;
+
+    manager.stopThreadRuntime("thread-1");
+
+    expect(sibling.aborted).toBe(false);
+  });
+
+  it("hands the next generation a fresh signal on restart", () => {
+    const manager = makeManager();
+    start(manager, "thread-1");
+    const first = publishedSignal(manager, "thread-1")!;
+
+    restart(manager, "thread-1");
+    const second = publishedSignal(manager, "thread-1")!;
+
+    expect(first.aborted).toBe(true);
+    expect(second).not.toBe(first);
+    expect(second.aborted).toBe(false);
+  });
+
+  it("aborts every thread when the manager is disposed", () => {
+    const manager = makeManager();
+    start(manager, "thread-1");
+    start(manager, "thread-2");
+    const signals = [
+      publishedSignal(manager, "thread-1")!,
+      publishedSignal(manager, "thread-2")!,
+    ];
+
+    manager.__internal_dispose();
+
+    expect(signals.map((signal) => signal.aborted)).toEqual([true, true]);
+  });
+});
+
+const makeRunningRuntime = () => {
+  const subscribers = new Set<() => void>();
+  const eventListeners = new Map<string, Set<() => void>>();
+  const runtime = {
+    isRunning: true,
+    messages: [],
+    subscribe: (callback: () => void) => {
+      subscribers.add(callback);
+      return () => subscribers.delete(callback);
+    },
+    unstable_on: (event: string, callback: () => void) => {
+      let listeners = eventListeners.get(event);
+      if (!listeners) {
+        listeners = new Set();
+        eventListeners.set(event, listeners);
+      }
+      listeners.add(callback);
+      return () => listeners.delete(callback);
+    },
+  } as unknown as ThreadRuntimeCore & { isRunning: boolean };
+
+  return {
+    runtime,
+    // what a stopOnClientDestroy consumer does when the destroy signal fires
+    stop: () => {
+      runtime.isRunning = false;
+      for (const callback of eventListeners.get("runEnd") ?? []) callback();
+      for (const callback of subscribers) callback();
+    },
+  };
+};
+
+const publish = (
+  manager: RemoteThreadListHookInstanceManager,
+  threadId: string,
+  runtime: ThreadRuntimeCore,
+) => {
+  const internals = manager as unknown as {
+    instances: Map<string, { generation: number }>;
+    _publishThreadRuntime: (
+      threadId: string,
+      runtime: ThreadRuntimeCore,
+      generation: number,
+    ) => void;
+  };
+  internals._publishThreadRuntime(
+    threadId,
+    runtime,
+    internals.instances.get(threadId)!.generation,
+  );
+};
+
+describe("RemoteThreadListHookInstanceManager restart teardown", () => {
+  it("keeps the outgoing generation's terminal events off the thread's subscribers", () => {
+    const manager = makeManager();
+    start(manager, "thread-1");
+    const { runtime, stop } = makeRunningRuntime();
+    publish(manager, "thread-1", runtime);
+
+    const signal = publishedSignal(manager, "thread-1")!;
+    signal.addEventListener("abort", stop, { once: true });
+
+    const events: string[] = [];
+    manager.__internal_subscribeThreadEvents((event) =>
+      events.push(event.type),
+    );
+    const runningChanges: boolean[] = [];
+    manager.__internal_subscribeRunningChanged(() =>
+      runningChanges.push(manager.__internal_isThreadRunning("thread-1")),
+    );
+    expect(manager.__internal_isThreadRunning("thread-1")).toBe(true);
+
+    restart(manager, "thread-1");
+
+    expect(signal.aborted).toBe(true);
+    expect(events).toEqual([]);
+    expect(runningChanges).toEqual([]);
+    expect(manager.__internal_isThreadRunning("thread-1")).toBe(true);
+  });
+});