@assistant-ui/core 0.3.16 → 0.3.17

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

@@ -1,50 +1,120 @@
+/** @vitest-environment jsdom */
+import { render, screen } from "@testing-library/react";
+import { useState, type ReactNode } from "react";
 import { describe, expect, it, vi } from "vitest";
-import type { ReactElement } from "react";
-import { MessagePrimitiveAttachments } from "./MessageAttachments";
+import type { CompleteAttachment } from "../../../types/attachment";
 
-const mockUseAuiState = vi.fn();
-type UseAuiStateSelector = Parameters<
-  (typeof import("@assistant-ui/store"))["useAuiState"]
->[0];
-type AttachmentsElement = ReactElement<{ children: () => null }>;
-
-vi.mock("react", async (importOriginal) => {
-  const actual = await importOriginal<typeof import("react")>();
-  return {
-    ...actual,
-    useMemo: (factory: () => unknown) => factory(),
-  };
-});
+const mocks = vi.hoisted(() => ({
+  role: "user" as "user" | "assistant",
+  attachments: undefined as CompleteAttachment[] | undefined,
+  aui: {
+    message: {
+      attachment: ({ index }: { index: number }) => ({
+        getState: () => mocks.attachments?.[index],
+      }),
+    },
+  },
+}));
 
-vi.mock("@assistant-ui/store", async (importOriginal) => {
-  const actual = await importOriginal<typeof import("@assistant-ui/store")>();
-  return {
-    ...actual,
-    useAuiState: (selector: UseAuiStateSelector) => mockUseAuiState(selector),
-  };
-});
+vi.mock("@assistant-ui/store", async (importOriginal) => ({
+  ...(await importOriginal<typeof import("@assistant-ui/store")>()),
+  useAui: () => mocks.aui,
+  useAuiState: (selector: (state: unknown) => unknown) =>
+    selector({
+      message: {
+        role: mocks.role,
+        attachments: mocks.attachments,
+      },
+    }),
+  RenderChildrenWithAccessor: ({
+    getItemState,
+    children,
+  }: {
+    getItemState: (aui: typeof mocks.aui) => CompleteAttachment;
+    children: (getItem: () => CompleteAttachment) => ReactNode;
+  }) => children(() => getItemState(mocks.aui)),
+}));
 
-const renderAttachmentsInner = () => {
-  const element = MessagePrimitiveAttachments({
-    children: () => null,
-  }) as AttachmentsElement;
+vi.mock(
+  "../../providers/AttachmentByIndexProvider",
+  async (importOriginal) => ({
+    ...(await importOriginal<
+      typeof import("../../providers/AttachmentByIndexProvider")
+    >()),
+    MessageAttachmentByIndexProvider: ({ children }: { children: ReactNode }) =>
+      children,
+  }),
+);
 
-  const Inner = element.type as (props: typeof element.props) => unknown;
-  return Inner(element.props);
+import { MessagePrimitiveAttachments } from "./MessageAttachments";
+
+const attachment = (id: string): CompleteAttachment => ({
+  id,
+  type: "file",
+  name: `${id}.txt`,
+  status: { type: "complete" },
+  content: [],
+});
+
+const StatefulAttachment = ({
+  attachment,
+}: {
+  attachment: CompleteAttachment;
+}) => {
+  const [initialId] = useState(attachment.id);
+  return <span>{initialId}</span>;
 };
 
 describe("MessagePrimitiveAttachments", () => {
   it("treats missing user message attachments as empty", () => {
-    mockUseAuiState.mockImplementation((selector: UseAuiStateSelector) =>
-      selector({
-        message: {
-          role: "user",
-          attachments: undefined,
-        },
-      } as never),
+    mocks.role = "user";
+    mocks.attachments = undefined;
+
+    const view = render(
+      <MessagePrimitiveAttachments>{() => null}</MessagePrimitiveAttachments>,
+    );
+
+    expect(view.container.childElementCount).toBe(0);
+  });
+
+  it("does not reuse component state for a different attachment", () => {
+    mocks.role = "user";
+    mocks.attachments = [attachment("first-file"), attachment("second-file")];
+
+    const view = render(
+      <MessagePrimitiveAttachments>
+        {({ attachment }) => <StatefulAttachment attachment={attachment} />}
+      </MessagePrimitiveAttachments>,
+    );
+
+    mocks.attachments = [mocks.attachments[1]!];
+    view.rerender(
+      <MessagePrimitiveAttachments>
+        {({ attachment }) => <StatefulAttachment attachment={attachment} />}
+      </MessagePrimitiveAttachments>,
+    );
+
+    expect(screen.queryByText("second-file")).not.toBeNull();
+    expect(screen.queryByText("first-file")).toBeNull();
+  });
+
+  it("keeps component state with attachments when they are reordered", () => {
+    mocks.role = "user";
+    mocks.attachments = [attachment("first-file"), attachment("second-file")];
+
+    const view = render(
+      <MessagePrimitiveAttachments>
+        {({ attachment }) => <StatefulAttachment attachment={attachment} />}
+      </MessagePrimitiveAttachments>,
+    );
+
+    mocks.attachments = [mocks.attachments[1]!, mocks.attachments[0]!];
+    view.rerender(
+      <MessagePrimitiveAttachments>
+        {({ attachment }) => <StatefulAttachment attachment={attachment} />}
+      </MessagePrimitiveAttachments>,
     );
 
-    expect(() => renderAttachmentsInner()).not.toThrow();
-    expect(renderAttachmentsInner()).toEqual([]);
+    expect(view.container.textContent).toBe("second-filefirst-file");
   });
 });

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

@@ -7,6 +7,7 @@ import {
   useMemo,
 } from "react";
 import { RenderChildrenWithAccessor, useAuiState } from "@assistant-ui/store";
+import { useShallowSelector } from "@assistant-ui/store/internal";
 import { MessageAttachmentByIndexProvider } from "../../providers/AttachmentByIndexProvider";
 
 type MessageAttachmentsComponentConfig = {
@@ -91,15 +92,17 @@ MessagePrimitiveAttachmentByIndex.displayName =
 const MessagePrimitiveAttachmentsInner: FC<{
   children: (value: { attachment: CompleteAttachment }) => ReactNode;
 }> = ({ children }) => {
-  const attachmentsCount = useAuiState((s) => {
-    if (s.message.role !== "user") return 0;
-    return (s.message.attachments ?? []).length;
-  });
+  const attachmentIds = useAuiState(
+    useShallowSelector((s) => {
+      if (s.message.role !== "user") return [];
+      return (s.message.attachments ?? []).map((attachment) => attachment.id);
+    }),
+  );
 
   return useMemo(
     () =>
-      Array.from({ length: attachmentsCount }, (_, index) => (
-        <MessageAttachmentByIndexProvider key={index} index={index}>
+      attachmentIds.map((attachmentId, index) => (
+        <MessageAttachmentByIndexProvider key={attachmentId} index={index}>
           <RenderChildrenWithAccessor
             getItemState={(aui) => aui.message.attachment({ index }).getState()}
           >
@@ -113,7 +116,7 @@ const MessagePrimitiveAttachmentsInner: FC<{
           </RenderChildrenWithAccessor>
         </MessageAttachmentByIndexProvider>
       )),
-    [attachmentsCount, children],
+    [attachmentIds, children],
   );
 };
 

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

@@ -2,7 +2,7 @@
 
 import { Fragment, type FC, type ReactNode, useMemo } from "react";
 import { useAuiState } from "@assistant-ui/store";
-import { useShallow } from "zustand/shallow";
+import { useShallowSelector } from "@assistant-ui/store/internal";
 import type { PartState } from "../../../store/scopes/part";
 import type {
   MessagePartStatus,
@@ -251,7 +251,7 @@ export const MessagePrimitiveGroupedParts = <TKey extends `group-${string}`>({
   indicator = "no-text",
   children,
 }: MessagePrimitiveGroupedParts.Props<TKey>): ReactNode => {
-  const parts = useAuiState(useShallow((s) => s.message.parts));
+  const parts = useAuiState(useShallowSelector((s) => s.message.parts));
   // Handed to `groupBy` as its `context` argument (see GroupByContext).
   const toolUIs = useAuiState((s) => s.tools.toolUIs);
   // Subscribe to a boolean, not the status object: the tree only needs to

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

@@ -40,7 +40,7 @@ import {
 } from "../../../types/message";
 import type { DataRenderersState } from "../../types/scopes/dataRenderers";
 import type { ToolsState } from "../../types/scopes/tools";
-import { useShallow } from "zustand/shallow";
+import { useShallowSelector } from "@assistant-ui/store/internal";
 
 type MessagePartRange =
   | { type: "single"; index: number }
@@ -174,10 +174,10 @@ const useMessagePartsGroups = (
   useChainOfThought: boolean,
 ): { ranges: MessagePartRange[]; partIds: (string | undefined)[] } => {
   const messageTypes = useAuiState(
-    useShallow((s) => s.message.parts.map((c: any) => c.type)),
+    useShallowSelector((s) => s.message.parts.map((c: any) => c.type)),
   );
   const partIds = useAuiState(
-    useShallow((s) =>
+    useShallowSelector((s) =>
       s.message.parts.map((c: any) =>
         c.type === "tool-call" ? c.toolCallId : undefined,
       ),

package/src/react/primitives/threadList/ThreadListItems.test.tsx

@@ -0,0 +1,99 @@
+/** @vitest-environment jsdom */
+import { render, screen } from "@testing-library/react";
+import { useState, type ReactNode } from "react";
+import { describe, expect, it, vi } from "vitest";
+import type { ThreadListItemState } from "../../../store/scopes/thread-list-item";
+
+const mocks = vi.hoisted(() => ({
+  threadItems: [] as ThreadListItemState[],
+  archivedThreadItems: [] as ThreadListItemState[],
+  aui: {
+    threads: {
+      item: ({ index, archived }: { index: number; archived?: boolean }) => ({
+        getState: () =>
+          archived
+            ? mocks.archivedThreadItems[index]
+            : mocks.threadItems[index],
+      }),
+    },
+  },
+}));
+
+vi.mock("@assistant-ui/store", async (importOriginal) => ({
+  ...(await importOriginal<typeof import("@assistant-ui/store")>()),
+  useAui: () => mocks.aui,
+  useAuiState: (selector: (state: unknown) => unknown) =>
+    selector({
+      threads: {
+        threadIds: mocks.threadItems.map((item) => item.id),
+        archivedThreadIds: mocks.archivedThreadItems.map((item) => item.id),
+      },
+    }),
+  RenderChildrenWithAccessor: ({
+    getItemState,
+    children,
+  }: {
+    getItemState: (aui: typeof mocks.aui) => ThreadListItemState;
+    children: (getItem: () => ThreadListItemState) => ReactNode;
+  }) => children(() => getItemState(mocks.aui)),
+}));
+
+vi.mock(
+  "../../providers/ThreadListItemByIndexProvider",
+  async (importOriginal) => ({
+    ...(await importOriginal<
+      typeof import("../../providers/ThreadListItemByIndexProvider")
+    >()),
+    ThreadListItemByIndexProvider: ({ children }: { children: ReactNode }) =>
+      children,
+  }),
+);
+
+import { ThreadListPrimitiveItems } from "./ThreadListItems";
+
+const threadItem = (id: string, title: string): ThreadListItemState => ({
+  id,
+  remoteId: id,
+  externalId: undefined,
+  title,
+  status: "regular",
+  isRunning: false,
+});
+
+const StatefulThreadItem = ({
+  threadListItem,
+}: {
+  threadListItem: ThreadListItemState;
+}) => {
+  const [initialTitle] = useState(threadListItem.title);
+  return <span>{initialTitle}</span>;
+};
+
+describe("ThreadListPrimitiveItems", () => {
+  it("does not reuse component state for a different thread", () => {
+    mocks.threadItems = [
+      threadItem("first", "First thread"),
+      threadItem("second", "Second thread"),
+    ];
+
+    const view = render(
+      <ThreadListPrimitiveItems>
+        {({ threadListItem }) => (
+          <StatefulThreadItem threadListItem={threadListItem} />
+        )}
+      </ThreadListPrimitiveItems>,
+    );
+
+    mocks.threadItems = [mocks.threadItems[1]!];
+    view.rerender(
+      <ThreadListPrimitiveItems>
+        {({ threadListItem }) => (
+          <StatefulThreadItem threadListItem={threadListItem} />
+        )}
+      </ThreadListPrimitiveItems>,
+    );
+
+    expect(screen.queryByText("Second thread")).not.toBeNull();
+    expect(screen.queryByText("First thread")).toBeNull();
+  });
+});

package/src/react/primitives/threadList/ThreadListItems.tsx

@@ -6,6 +6,7 @@ import {
   useMemo,
 } from "react";
 import { RenderChildrenWithAccessor, useAuiState } from "@assistant-ui/store";
+import { useShallowSelector } from "@assistant-ui/store/internal";
 import type { ThreadListItemState } from "../../../store/scopes/thread-list-item";
 import { ThreadListItemByIndexProvider } from "../../providers/ThreadListItemByIndexProvider";
 
@@ -64,15 +65,17 @@ const ThreadListPrimitiveItemsInner: FC<{
   archived: boolean;
   children: (value: { threadListItem: ThreadListItemState }) => ReactNode;
 }> = ({ archived, children }) => {
-  const contentLength = useAuiState((s) =>
-    archived ? s.threads.archivedThreadIds.length : s.threads.threadIds.length,
+  const threadIds = useAuiState(
+    useShallowSelector((s) =>
+      archived ? s.threads.archivedThreadIds : s.threads.threadIds,
+    ),
   );
 
   return useMemo(
     () =>
-      Array.from({ length: contentLength }, (_, index) => (
+      threadIds.map((threadId, index) => (
         <ThreadListItemByIndexProvider
-          key={index}
+          key={threadId}
           index={index}
           archived={archived}
         >
@@ -91,7 +94,7 @@ const ThreadListPrimitiveItemsInner: FC<{
           </RenderChildrenWithAccessor>
         </ThreadListItemByIndexProvider>
       )),
-    [contentLength, archived, children],
+    [threadIds, archived, children],
   );
 };
 

package/src/react/runtimes/RemoteThreadListHookInstanceManager.running.test.tsx

@@ -5,9 +5,12 @@ import type { ThreadMessage } from "../../types/message";
 import { RemoteThreadListHookInstanceManager } from "./RemoteThreadListHookInstanceManager";
 
 const makeRuntime = (
-  initial: Partial<Pick<ThreadRuntimeCore, "isRunning" | "messages">> = {},
+  initial: Partial<Pick<ThreadRuntimeCore, "isRunning" | "messages">> & {
+    initialized?: boolean;
+  } = {},
 ) => {
   const subscribers = new Set<() => void>();
+  const eventListeners = new Map<string, Set<() => void>>();
   const runtime = {
     isRunning: initial.isRunning,
     messages: initial.messages ?? [],
@@ -15,6 +18,23 @@ const makeRuntime = (
       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);
+      // mirrors the real core: initialize latches, so a subscriber that
+      // attaches after initialization gets a deferred replay
+      if (event === "initialize" && initial.initialized) {
+        const latched = listeners;
+        queueMicrotask(() => {
+          if (latched.has(callback)) callback();
+        });
+      }
+      return () => listeners.delete(callback);
+    },
   } as unknown as ThreadRuntimeCore & {
     isRunning: boolean | undefined;
     messages: readonly ThreadMessage[];
@@ -23,6 +43,14 @@ const makeRuntime = (
   return {
     runtime,
     subscriberCount: () => subscribers.size,
+    eventListenerCount: () =>
+      [...eventListeners.values()].reduce(
+        (total, listeners) => total + listeners.size,
+        0,
+      ),
+    emit: (event: string) => {
+      for (const callback of eventListeners.get(event) ?? []) callback();
+    },
     setRunning: (isRunning: boolean | undefined) => {
       runtime.isRunning = isRunning;
       for (const callback of subscribers) callback();
@@ -152,6 +180,71 @@ describe("RemoteThreadListHookInstanceManager run tracking", () => {
     manager.stopThreadRuntime("thread-1");
 
     expect(thread.subscriberCount()).toBe(0);
+    expect(thread.eventListenerCount()).toBe(0);
     expect(manager.__internal_isThreadRunning("thread-1")).toBe(false);
   });
+
+  it("forwards every lifecycle event from every tracked thread with its thread id", () => {
+    const manager = makeManager();
+    const events: unknown[] = [];
+    manager.__internal_subscribeThreadEvents((event) => events.push(event));
+
+    start(manager, "thread-1");
+    start(manager, "thread-2");
+    const first = makeRuntime();
+    const second = makeRuntime();
+    publish(manager, "thread-1", first.runtime);
+    publish(manager, "thread-2", second.runtime);
+
+    second.emit("runStart");
+    first.emit("runEnd");
+    first.emit("initialize");
+    second.emit("modelContextUpdate");
+
+    expect(events).toEqual([
+      { threadId: "thread-2", type: "runStart" },
+      { threadId: "thread-1", type: "runEnd" },
+      { threadId: "thread-1", type: "initialize" },
+      { threadId: "thread-2", type: "modelContextUpdate" },
+    ]);
+  });
+
+  it("replays a latched initialize from the surviving runtime after a restart", async () => {
+    const manager = makeManager();
+    const events: unknown[] = [];
+    manager.__internal_subscribeThreadEvents((event) => events.push(event));
+
+    start(manager, "thread-1");
+    const before = makeRuntime({ initialized: true });
+    publish(manager, "thread-1", before.runtime);
+    const after = makeRuntime({ initialized: true });
+    publish(manager, "thread-1", after.runtime);
+    await new Promise<void>((resolve) => queueMicrotask(resolve));
+
+    expect(before.eventListenerCount()).toBe(0);
+    expect(events).toEqual([{ threadId: "thread-1", type: "initialize" }]);
+  });
+
+  it("stops forwarding lifecycle events from a runtime a restart replaced", () => {
+    const manager = makeManager();
+    const events: unknown[] = [];
+    manager.__internal_subscribeThreadEvents((event) => events.push(event));
+
+    start(manager, "thread-1");
+    const before = makeRuntime();
+    publish(manager, "thread-1", before.runtime);
+    const after = makeRuntime();
+    publish(manager, "thread-1", after.runtime);
+
+    before.emit("runEnd");
+    before.emit("initialize");
+    after.emit("runEnd");
+    after.emit("modelContextUpdate");
+
+    expect(before.eventListenerCount()).toBe(0);
+    expect(events).toEqual([
+      { threadId: "thread-1", type: "runEnd" },
+      { threadId: "thread-1", type: "modelContextUpdate" },
+    ]);
+  });
 });