@assistant-ui/core 0.3.16 → 0.3.17

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

@@ -0,0 +1,131 @@
+// @vitest-environment jsdom
+
+import { act, cleanup, render, waitFor } from "@testing-library/react";
+import { startTransition, Suspense } from "react";
+import { afterEach, describe, expect, it, vi } from "vitest";
+
+const { aui, getProvider, resetProvider } = vi.hoisted(() => {
+  type TestProvider = {
+    getModelContext: () => {
+      system?: string | undefined;
+      tools?:
+        | Record<string, { execute?: (() => unknown) | undefined }>
+        | undefined;
+    };
+  };
+  let provider: TestProvider | undefined;
+
+  return {
+    aui: {
+      modelContext: {
+        register: vi.fn((nextProvider: TestProvider) => {
+          provider = nextProvider;
+          return () => {
+            if (provider === nextProvider) provider = undefined;
+          };
+        }),
+      },
+    },
+    getProvider: () => provider,
+    resetProvider: () => {
+      provider = undefined;
+    },
+  };
+});
+
+vi.mock("@assistant-ui/store", async (importOriginal) => ({
+  ...(await importOriginal()),
+  useAui: () => aui,
+}));
+
+import { useAssistantContext } from "./useAssistantContext";
+import { useAuiToolOverrides } from "./useAuiToolOverrides";
+
+afterEach(() => {
+  cleanup();
+  resetProvider();
+});
+
+describe("model context hooks", () => {
+  it("keeps assistant context scoped to the committed render", async () => {
+    const pending = new Promise<never>(() => {});
+    const renderB = vi.fn();
+    const Probe = ({ value, suspend }: { value: string; suspend: boolean }) => {
+      useAssistantContext({ getContext: () => value });
+      if (value === "workspace-b") renderB();
+      if (suspend) throw pending;
+      return null;
+    };
+
+    const view = render(
+      <Suspense fallback={null}>
+        <Probe value="workspace-a" suspend={false} />
+      </Suspense>,
+    );
+    await waitFor(() => expect(getProvider()).toBeDefined());
+
+    act(() => {
+      startTransition(() => {
+        view.rerender(
+          <Suspense fallback={null}>
+            <Probe value="workspace-b" suspend />
+          </Suspense>,
+        );
+      });
+    });
+
+    expect(renderB).toHaveBeenCalled();
+    expect(getProvider()?.getModelContext().system).toBe("workspace-a");
+
+    view.rerender(
+      <Suspense fallback={null}>
+        <Probe value="workspace-b" suspend={false} />
+      </Suspense>,
+    );
+
+    expect(getProvider()?.getModelContext().system).toBe("workspace-b");
+  });
+
+  it("keeps tool overrides scoped to the committed render", async () => {
+    const pending = new Promise<never>(() => {});
+    const renderB = vi.fn();
+    const Probe = ({ value, suspend }: { value: string; suspend: boolean }) => {
+      useAuiToolOverrides({ search: { execute: () => value } });
+      if (value === "workspace-b") renderB();
+      if (suspend) throw pending;
+      return null;
+    };
+
+    const view = render(
+      <Suspense fallback={null}>
+        <Probe value="workspace-a" suspend={false} />
+      </Suspense>,
+    );
+    await waitFor(() => expect(getProvider()).toBeDefined());
+
+    act(() => {
+      startTransition(() => {
+        view.rerender(
+          <Suspense fallback={null}>
+            <Probe value="workspace-b" suspend />
+          </Suspense>,
+        );
+      });
+    });
+
+    expect(renderB).toHaveBeenCalled();
+    expect(getProvider()?.getModelContext().tools?.search?.execute?.()).toBe(
+      "workspace-a",
+    );
+
+    view.rerender(
+      <Suspense fallback={null}>
+        <Probe value="workspace-b" suspend={false} />
+      </Suspense>,
+    );
+
+    expect(getProvider()?.getModelContext().tools?.search?.execute?.()).toBe(
+      "workspace-b",
+    );
+  });
+});

package/src/react/model-context/toolbox.test.ts

@@ -19,7 +19,7 @@ describe("makeToolCallTextComponent", () => {
         status: { type: "running" },
         addResult: () => {},
         resume: () => {},
-        respondToApproval: () => {},
+        respondToApproval: async () => {},
       }),
     ).toBe("Searching...");
 
@@ -34,7 +34,7 @@ describe("makeToolCallTextComponent", () => {
         status: { type: "complete" },
         addResult: () => {},
         resume: () => {},
-        respondToApproval: () => {},
+        respondToApproval: async () => {},
       }),
     ).toBe("Done searching");
   });
@@ -59,7 +59,7 @@ describe("makeToolCallTextComponent", () => {
         status: { type: "running" },
         addResult: () => {},
         resume: () => {},
-        respondToApproval: () => {},
+        respondToApproval: async () => {},
       }),
     ).toBe("Searching docs...");
 
@@ -74,7 +74,7 @@ describe("makeToolCallTextComponent", () => {
         status: { type: "complete" },
         addResult: () => {},
         resume: () => {},
-        respondToApproval: () => {},
+        respondToApproval: async () => {},
       }),
     ).toBe("Found 3 results for docs");
   });
@@ -94,7 +94,7 @@ describe("makeToolCallTextComponent", () => {
         argsText: "{}",
         addResult: () => {},
         resume: () => {},
-        respondToApproval: () => {},
+        respondToApproval: async () => {},
       } as unknown as ToolCallMessagePartProps),
     ).toBe("Complete");
   });
@@ -115,7 +115,7 @@ describe("makeToolCallTextComponent", () => {
         status: { type: "incomplete", reason: "error" },
         addResult: () => {},
         resume: () => {},
-        respondToApproval: () => {},
+        respondToApproval: async () => {},
       }),
     ).toBe("Finished");
   });
@@ -136,7 +136,7 @@ describe("makeToolCallTextComponent", () => {
         status: { type: "requires-action", reason: "interrupt" },
         addResult: () => {},
         resume: () => {},
-        respondToApproval: () => {},
+        respondToApproval: async () => {},
       }),
     ).toBe("Waiting for approval...");
   });
@@ -156,7 +156,7 @@ describe("makeToolCallTextComponent", () => {
         status: { type: "complete" },
         addResult: () => {},
         resume: () => {},
-        respondToApproval: () => {},
+        respondToApproval: async () => {},
       }),
     ).toBeNull();
   });
@@ -176,7 +176,7 @@ describe("makeToolCallTextComponent", () => {
         status: { type: "running" },
         addResult: () => {},
         resume: () => {},
-        respondToApproval: () => {},
+        respondToApproval: async () => {},
       }),
     ).toBeNull();
   });

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

@@ -9,6 +9,15 @@ import type {
   ToolCallMessagePartComponent,
   ToolCallMessagePartProps,
 } from "../types/MessagePartComponentTypes";
+import {
+  resolveToolCallText,
+  type ToolCallText as NeutralToolCallText,
+} from "../../model-context/tool-call-text";
+
+export type ToolCallText<
+  TArgs extends Record<string, unknown>,
+  TResult,
+> = NeutralToolCallText<TArgs, TResult, ReactNode>;
 
 /**
  * Resolves whether a tool's UI should be presented standalone (outside the
@@ -61,42 +70,6 @@ type ToolStreamCall<TArgs extends Record<string, unknown>, TResult> = (
   context: ToolExecuteContext,
 ) => void;
 
-type ToolCallRunningText<TArgs extends Record<string, unknown>> =
-  | ReactNode
-  | ((options: { args: TArgs }) => ReactNode);
-
-type ToolCallCompleteText<TArgs extends Record<string, unknown>, TResult> =
-  | ReactNode
-  | ((options: { args: TArgs; result: TResult | undefined }) => ReactNode);
-
-export type ToolCallText<TArgs extends Record<string, unknown>, TResult> =
-  | {
-      running: ToolCallRunningText<TArgs>;
-      complete?: ToolCallCompleteText<TArgs, TResult> | undefined;
-    }
-  | {
-      running?: ToolCallRunningText<TArgs> | undefined;
-      complete: ToolCallCompleteText<TArgs, TResult>;
-    };
-
-const resolveToolCallText = <TArgs extends Record<string, unknown>, TResult>(
-  text: ToolCallText<TArgs, TResult>,
-  part: ToolCallMessagePartProps<TArgs, TResult>,
-): ReactNode => {
-  const isRunning =
-    part.status?.type === "running" || part.status?.type === "requires-action";
-
-  if (!isRunning) {
-    const value = text.complete;
-    if (typeof value !== "function") return value ?? null;
-    return value({ args: part.args, result: part.result });
-  }
-
-  const value = text.running;
-  if (typeof value !== "function") return value ?? null;
-  return value({ args: part.args });
-};
-
 export const makeToolCallTextComponent = <
   TArgs extends Record<string, unknown>,
   TResult,

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

@@ -1,4 +1,4 @@
-import { useEffect, useRef } from "react";
+import { useEffect, useInsertionEffect, useRef } from "react";
 import { useAui } from "@assistant-ui/store";
 import type { AssistantContextConfig } from "../..";
 
@@ -8,7 +8,9 @@ export const useAssistantContext = (config: AssistantContextConfig) => {
   const { getContext, disabled = false } = config;
   const aui = useAui();
   const getContextRef = useRef(getContext);
-  getContextRef.current = getContext;
+  useInsertionEffect(() => {
+    getContextRef.current = getContext;
+  }, [getContext]);
 
   useEffect(() => {
     if (disabled) return;

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

@@ -1,4 +1,4 @@
-import { useEffect, useRef } from "react";
+import { useEffect, useInsertionEffect, useRef } from "react";
 import { useAui } from "@assistant-ui/store";
 import type { Tool } from "assistant-stream";
 
@@ -25,7 +25,9 @@ type AuiToolOverrides = Record<string, AuiToolOverride<any, any>>;
 export function useAuiToolOverrides(overrides: AuiToolOverrides): void {
   const aui = useAui();
   const overridesRef = useRef(overrides);
-  overridesRef.current = overrides;
+  useInsertionEffect(() => {
+    overridesRef.current = overrides;
+  }, [overrides]);
 
   useEffect(() => {
     return aui.modelContext.register({

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

@@ -0,0 +1,48 @@
+// @vitest-environment jsdom
+import { useState, type FC } from "react";
+import { act, cleanup, render, screen } from "@testing-library/react";
+import { afterEach, describe, expect, it } from "vitest";
+import { useInlineRender } from "./useInlineRender";
+import type { ToolCallMessagePartProps } from "../types/MessagePartComponentTypes";
+
+afterEach(() => cleanup());
+
+type Props = ToolCallMessagePartProps<any, any>;
+
+// The point of useInlineRender is a component whose identity survives a
+// changing toolUI, so inner state must not be lost when toolUI swaps.
+const makeToolUI = (label: string): FC<Props> =>
+  function Inner() {
+    const [n, setN] = useState(0);
+    return (
+      <button onClick={() => setN((v) => v + 1)}>
+        {label}:{n}
+      </button>
+    );
+  };
+
+describe("useInlineRender", () => {
+  it("follows toolUI changes without remounting the inner tree", async () => {
+    let swap!: (fc: FC<Props>) => void;
+    const identities = new Set<unknown>();
+
+    const Host = () => {
+      const [toolUI, setToolUI] = useState(() => makeToolUI("a"));
+      swap = (fc) => setToolUI(() => fc);
+      const ToolUI = useInlineRender(toolUI);
+      identities.add(ToolUI);
+      return <ToolUI {...({} as Props)} />;
+    };
+
+    render(<Host />);
+    expect(screen.getByRole("button").textContent).toBe("a:0");
+
+    await act(async () => screen.getByRole("button").click());
+    expect(screen.getByRole("button").textContent).toBe("a:1");
+
+    await act(async () => swap(makeToolUI("b")));
+    // label follows the new toolUI, counter survives => no remount
+    expect(screen.getByRole("button").textContent).toBe("b:1");
+    expect(identities.size).toBe(1);
+  });
+});

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

@@ -1,25 +1,29 @@
 import { type FC, useCallback, useEffect, useState } from "react";
 import type { ToolCallMessagePartProps } from "../types/MessagePartComponentTypes";
-import { create } from "zustand";
+import { WritableSubscribable } from "../../subscribable/subscribable";
+import { useSubscribable } from "../../store/runtime-clients/useSubscribable";
 
 export const useInlineRender = <TArgs, TResult>(
   toolUI: FC<ToolCallMessagePartProps<TArgs, TResult>>,
 ): FC<ToolCallMessagePartProps<TArgs, TResult>> => {
-  const [useToolUIStore] = useState(() =>
-    create(() => ({
-      toolUI,
-    })),
+  const [toolUIStore] = useState(
+    () =>
+      new WritableSubscribable<FC<ToolCallMessagePartProps<TArgs, TResult>>>(
+        toolUI,
+      ),
   );
 
   useEffect(() => {
-    useToolUIStore.setState({ toolUI });
-  }, [toolUI, useToolUIStore]);
+    toolUIStore.setState(toolUI);
+  }, [toolUI, toolUIStore]);
 
   return useCallback(
     function ToolUI(args) {
-      const store = useToolUIStore();
-      return store.toolUI(args);
+      // Invoked as a plain function so its hooks belong to this component,
+      // which keeps its identity while `toolUI` changes underneath.
+      const currentToolUI = useSubscribable(toolUIStore);
+      return currentToolUI(args);
     },
-    [useToolUIStore],
+    [toolUIStore],
   );
 };

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],
   );
 };