@assistant-ui/core 0.3.16 → 0.3.18

package/src/react/model-context/useToolArgsStatus.test.tsx

@@ -0,0 +1,30 @@
+// @vitest-environment jsdom
+
+import { renderHook } from "@testing-library/react";
+import { describe, expect, it, vi } from "vitest";
+
+const part = vi.hoisted(() => ({
+  type: "tool-call" as const,
+  status: { type: "complete" as const },
+  args: JSON.parse('{"__proto__":"value"}') as Record<string, unknown>,
+}));
+
+vi.mock("@assistant-ui/store", async (importOriginal) => {
+  const actual = await importOriginal<typeof import("@assistant-ui/store")>();
+  return {
+    ...actual,
+    useAuiState: (selector: (state: { part: typeof part }) => unknown) =>
+      selector({ part }),
+  };
+});
+
+import { useToolArgsStatus } from "./useToolArgsStatus";
+
+describe("useToolArgsStatus", () => {
+  it("reports status for an argument named __proto__", () => {
+    const { result } = renderHook(() => useToolArgsStatus());
+
+    expect(Object.hasOwn(result.current.propStatus, "__proto__")).toBe(true);
+    expect(result.current.propStatus.__proto__).toBe("complete");
+  });
+});

package/src/react/model-context/useToolArgsStatus.ts

@@ -4,6 +4,7 @@ import {
   getPartialJsonObjectFieldState,
   getPartialJsonObjectMeta,
 } from "assistant-stream/utils";
+import { nullProtoRecord } from "../../utils/record";
 
 type PropFieldStatus = "streaming" | "complete";
 
@@ -60,7 +61,7 @@ export const useToolArgsStatus = <
     const isStreaming = statusType === "running";
     const args = part.args as Record<string, unknown>;
     const meta = getPartialJsonObjectMeta(args as Record<symbol, unknown>);
-    const propStatus: Partial<Record<string, PropFieldStatus>> = {};
+    const propStatus = nullProtoRecord<PropFieldStatus>();
 
     for (const key of Object.keys(args)) {
       if (meta) {

package/src/react/primitive-hooks/useMessageError.ts

@@ -1,23 +1,6 @@
 import { useAuiState } from "@assistant-ui/store";
+import { messageErrorText } from "../../store/primitive-predicates";
 
 export const useMessageError = () => {
-  return useAuiState((s) => {
-    if (
-      s.message.status?.type !== "incomplete" ||
-      s.message.status.reason !== "error"
-    ) {
-      return undefined;
-    }
-    const error = s.message.status.error;
-    if (typeof error === "string") return error;
-    if (
-      typeof error === "object" &&
-      error !== null &&
-      "message" in error &&
-      typeof error.message === "string"
-    ) {
-      return error.message;
-    }
-    return error ?? "An error occurred";
-  });
+  return useAuiState(messageErrorText);
 };

package/src/react/primitive-hooks/useThreadListLoadMore.ts

@@ -1,11 +1,10 @@
 import { useCallback } from "react";
 import { useAui, useAuiState } from "@assistant-ui/store";
+import { threadListLoadMoreDisabled } from "../../store/primitive-predicates";
 
 export const useThreadListLoadMore = () => {
   const aui = useAui();
-  const disabled = useAuiState(
-    (s) => !s.threads.hasMore || s.threads.isLoading || s.threads.isLoadingMore,
-  );
+  const disabled = useAuiState(threadListLoadMoreDisabled);
 
   const loadMore = useCallback(() => {
     aui.threads.loadMore();

package/src/react/primitives/composer/ComposerAttachments.test.tsx

@@ -0,0 +1,76 @@
+/** @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 { Attachment } from "../../../types/attachment";
+
+const mocks = vi.hoisted(() => ({
+  attachments: [] as Attachment[],
+  aui: {
+    composer: {
+      attachment: ({ index }: { index: number }) => ({
+        getState: () => mocks.attachments[index],
+      }),
+    },
+  },
+}));
+
+vi.mock("@assistant-ui/store", async (importOriginal) => ({
+  ...(await importOriginal<typeof import("@assistant-ui/store")>()),
+  useAui: () => mocks.aui,
+  useAuiState: (selector: (state: unknown) => unknown) =>
+    selector({ composer: { attachments: mocks.attachments } }),
+  RenderChildrenWithAccessor: ({
+    getItemState,
+    children,
+  }: {
+    getItemState: (aui: typeof mocks.aui) => Attachment;
+    children: (getItem: () => Attachment) => ReactNode;
+  }) => children(() => getItemState(mocks.aui)),
+}));
+
+vi.mock("../../providers/AttachmentByIndexProvider", () => ({
+  ComposerAttachmentByIndexProvider: ({ children }: { children: ReactNode }) =>
+    children,
+}));
+
+import { ComposerPrimitiveAttachments } from "./ComposerAttachments";
+
+const attachment = (id: string, name: string): Attachment => ({
+  id,
+  name,
+  type: "file",
+  contentType: "text/plain",
+  status: { type: "complete" },
+  content: [],
+});
+
+const StatefulAttachment = ({ attachment }: { attachment: Attachment }) => {
+  const [initialName] = useState(attachment.name);
+  return <span>{initialName}</span>;
+};
+
+describe("ComposerPrimitiveAttachments", () => {
+  it("does not reuse component state for a different attachment", () => {
+    mocks.attachments = [
+      attachment("first", "first.txt"),
+      attachment("second", "second.txt"),
+    ];
+
+    const view = render(
+      <ComposerPrimitiveAttachments>
+        {({ attachment }) => <StatefulAttachment attachment={attachment} />}
+      </ComposerPrimitiveAttachments>,
+    );
+
+    mocks.attachments = [mocks.attachments[1]!];
+    view.rerender(
+      <ComposerPrimitiveAttachments>
+        {({ attachment }) => <StatefulAttachment attachment={attachment} />}
+      </ComposerPrimitiveAttachments>,
+    );
+
+    expect(screen.queryByText("second.txt")).not.toBeNull();
+    expect(screen.queryByText("first.txt")).toBeNull();
+  });
+});

package/src/react/primitives/composer/ComposerAttachments.tsx

@@ -7,6 +7,7 @@ import {
   useMemo,
 } from "react";
 import { RenderChildrenWithAccessor, useAuiState } from "@assistant-ui/store";
+import { useShallowSelector } from "@assistant-ui/store/internal";
 import { ComposerAttachmentByIndexProvider } from "../../providers/AttachmentByIndexProvider";
 
 type ComposerAttachmentsComponentConfig = {
@@ -91,12 +92,16 @@ ComposerPrimitiveAttachmentByIndex.displayName =
 const ComposerPrimitiveAttachmentsInner: FC<{
   children: (value: { attachment: Attachment }) => ReactNode;
 }> = ({ children }) => {
-  const attachmentsCount = useAuiState((s) => s.composer.attachments.length);
+  const attachmentIds = useAuiState(
+    useShallowSelector((s) =>
+      s.composer.attachments.map((attachment) => attachment.id),
+    ),
+  );
 
   return useMemo(
     () =>
-      Array.from({ length: attachmentsCount }, (_, index) => (
-        <ComposerAttachmentByIndexProvider key={index} index={index}>
+      attachmentIds.map((attachmentId, index) => (
+        <ComposerAttachmentByIndexProvider key={attachmentId} index={index}>
           <RenderChildrenWithAccessor
             getItemState={(aui) =>
               aui.composer.attachment({ index }).getState()
@@ -112,7 +117,7 @@ const ComposerPrimitiveAttachmentsInner: FC<{
           </RenderChildrenWithAccessor>
         </ComposerAttachmentByIndexProvider>
       )),
-    [attachmentsCount, children],
+    [attachmentIds, children],
   );
 };
 

package/src/react/primitives/composer/ComposerQueue.test.tsx

@@ -0,0 +1,72 @@
+/** @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 { QueueItemState } from "../../../runtime/queue/queue-item";
+
+const mocks = vi.hoisted(() => ({
+  queue: [] as QueueItemState[],
+  aui: {
+    composer: {
+      queueItem: ({ index }: { index: number }) => ({
+        getState: () => mocks.queue[index],
+      }),
+    },
+  },
+}));
+
+vi.mock("@assistant-ui/store", async (importOriginal) => ({
+  ...(await importOriginal<typeof import("@assistant-ui/store")>()),
+  useAui: () => mocks.aui,
+  useAuiState: (selector: (state: unknown) => unknown) =>
+    selector({ composer: { queue: mocks.queue } }),
+  RenderChildrenWithAccessor: ({
+    getItemState,
+    children,
+  }: {
+    getItemState: (aui: typeof mocks.aui) => QueueItemState;
+    children: (getItem: () => QueueItemState) => ReactNode;
+  }) => children(() => getItemState(mocks.aui)),
+}));
+
+vi.mock("../../providers/QueueItemByIndexProvider", () => ({
+  QueueItemByIndexProvider: ({ children }: { children: ReactNode }) => children,
+}));
+
+import { ComposerPrimitiveQueue } from "./ComposerQueue";
+
+const queueItem = (id: string, prompt: string): QueueItemState => ({
+  id,
+  prompt,
+  parts: [{ type: "text", text: prompt }],
+});
+
+const StatefulQueueItem = ({ queueItem }: { queueItem: QueueItemState }) => {
+  const [initialPrompt] = useState(queueItem.prompt);
+  return <span>{initialPrompt}</span>;
+};
+
+describe("ComposerPrimitiveQueue", () => {
+  it("does not reuse component state for a different queue item", () => {
+    mocks.queue = [
+      queueItem("first", "first message"),
+      queueItem("second", "second message"),
+    ];
+
+    const view = render(
+      <ComposerPrimitiveQueue>
+        {({ queueItem }) => <StatefulQueueItem queueItem={queueItem} />}
+      </ComposerPrimitiveQueue>,
+    );
+
+    mocks.queue = [mocks.queue[1]!];
+    view.rerender(
+      <ComposerPrimitiveQueue>
+        {({ queueItem }) => <StatefulQueueItem queueItem={queueItem} />}
+      </ComposerPrimitiveQueue>,
+    );
+
+    expect(screen.queryByText("second message")).not.toBeNull();
+    expect(screen.queryByText("first message")).toBeNull();
+  });
+});

package/src/react/primitives/composer/ComposerQueue.tsx

@@ -1,5 +1,6 @@
 import { type FC, type ReactNode, memo, useMemo } from "react";
 import { RenderChildrenWithAccessor, useAuiState } from "@assistant-ui/store";
+import { useShallowSelector } from "@assistant-ui/store/internal";
 import type { QueueItemState } from "../../../store/scopes/queue-item";
 import { QueueItemByIndexProvider } from "../../providers/QueueItemByIndexProvider";
 
@@ -13,12 +14,14 @@ export namespace ComposerPrimitiveQueue {
 const ComposerPrimitiveQueueInner: FC<{
   children: (value: { queueItem: QueueItemState }) => ReactNode;
 }> = ({ children }) => {
-  const queue = useAuiState((s) => s.composer.queue.length);
+  const queueItemIds = useAuiState(
+    useShallowSelector((s) => s.composer.queue.map((item) => item.id)),
+  );
 
   return useMemo(
     () =>
-      Array.from({ length: queue }, (_, index) => (
-        <QueueItemByIndexProvider key={index} index={index}>
+      queueItemIds.map((queueItemId, index) => (
+        <QueueItemByIndexProvider key={queueItemId} index={index}>
           <RenderChildrenWithAccessor
             getItemState={(aui) => aui.composer.queueItem({ index }).getState()}
           >
@@ -32,7 +35,7 @@ const ComposerPrimitiveQueueInner: FC<{
           </RenderChildrenWithAccessor>
         </QueueItemByIndexProvider>
       )),
-    [queue, children],
+    [queueItemIds, children],
   );
 };
 

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();
+  });
+});