@assistant-ui/core 0.3.18 → 0.3.20

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

@@ -0,0 +1,204 @@
+// @vitest-environment jsdom
+
+import { cleanup, renderHook, waitFor } from "@testing-library/react";
+import { afterEach, describe, expect, it, vi } from "vitest";
+import type { Unstable_InteractableStateSchema } from "../types/scopes/interactables";
+
+type JsonSchema = Exclude<
+  Unstable_InteractableStateSchema,
+  { "~standard": unknown }
+>;
+type StandardSchemaWithJsonSchema = Extract<
+  Unstable_InteractableStateSchema,
+  { "~standard": unknown }
+> & {
+  "~standard": {
+    jsonSchema: { input: () => JsonSchema; output: () => JsonSchema };
+  };
+};
+
+const mocks = vi.hoisted(() => {
+  const unregister = vi.fn();
+  const register = vi.fn(() => unregister);
+  return {
+    register,
+    unregister,
+    aui: { unstable_interactables: { register } },
+    state: { optional: { part: undefined }, thread: { messages: [] } },
+    methods: {
+      setState: vi.fn(),
+      isPending: false,
+      error: undefined,
+      flush: vi.fn(async () => undefined),
+    },
+  };
+});
+
+vi.mock("@assistant-ui/store", async (importOriginal) => ({
+  ...(await importOriginal<typeof import("@assistant-ui/store")>()),
+  useAui: () => mocks.aui,
+  useAuiState: (selector: (state: typeof mocks.state) => unknown) =>
+    selector(mocks.state),
+}));
+
+vi.mock("./useInteractableState", async (importOriginal) => ({
+  ...(await importOriginal<typeof import("./useInteractableState")>()),
+  unstable_useInteractableState: () => [undefined, mocks.methods],
+}));
+
+import { unstable_useInteractable } from "./useInteractable";
+
+afterEach(() => {
+  cleanup();
+  vi.clearAllMocks();
+});
+
+describe("unstable_useInteractable", () => {
+  it("refreshes the registration when its JSON schema changes", async () => {
+    const schemaA = {
+      type: "object",
+      properties: { first: { type: "string" } },
+    } satisfies Unstable_InteractableStateSchema;
+    const schemaB = {
+      type: "object",
+      properties: { second: { type: "number" } },
+    } satisfies Unstable_InteractableStateSchema;
+    const initialA = { first: "one" };
+    const initialB = { second: 2 };
+
+    const hook = renderHook<
+      unknown,
+      { stateSchema: Unstable_InteractableStateSchema; initialState: unknown }
+    >(
+      ({ stateSchema, initialState }) =>
+        unstable_useInteractable("panel", {
+          id: "panel-1",
+          description: "A panel",
+          stateSchema,
+          initialState,
+        }),
+      { initialProps: { stateSchema: schemaA, initialState: initialA } },
+    );
+    await waitFor(() => expect(mocks.register).toHaveBeenCalledTimes(1));
+
+    hook.rerender({ stateSchema: schemaB, initialState: initialB });
+
+    await waitFor(() => expect(mocks.register).toHaveBeenCalledTimes(2));
+    expect(mocks.unregister).toHaveBeenCalledTimes(1);
+    expect(mocks.register).toHaveBeenLastCalledWith(
+      expect.objectContaining({ stateSchema: schemaB, initialState: initialB }),
+    );
+
+    hook.rerender({
+      stateSchema: {
+        type: "object",
+        properties: { second: { type: "number" } },
+      },
+      initialState: { second: 3 },
+    });
+    expect(mocks.register).toHaveBeenCalledTimes(2);
+  });
+
+  it("stabilizes equivalent rebuilt standard schemas", async () => {
+    const createSchema = (property: string): StandardSchemaWithJsonSchema => ({
+      "~standard": {
+        version: 1,
+        vendor: "test",
+        validate: () => ({ value: {} }),
+        jsonSchema: {
+          input: () => ({
+            type: "object",
+            properties: { [property]: { type: "string" } },
+          }),
+          output: () => ({ type: "object" }),
+        },
+      },
+    });
+    const firstSchema = createSchema("value");
+
+    const hook = renderHook<
+      unknown,
+      { stateSchema: Unstable_InteractableStateSchema }
+    >(
+      ({ stateSchema }) =>
+        unstable_useInteractable("panel", {
+          id: "panel-1",
+          description: "A panel",
+          stateSchema,
+          initialState: {},
+        }),
+      { initialProps: { stateSchema: firstSchema } },
+    );
+    await waitFor(() => expect(mocks.register).toHaveBeenCalledTimes(1));
+
+    hook.rerender({ stateSchema: createSchema("value") });
+    expect(mocks.register).toHaveBeenCalledTimes(1);
+
+    const changedSchema = createSchema("other");
+    hook.rerender({ stateSchema: changedSchema });
+    await waitFor(() => expect(mocks.register).toHaveBeenCalledTimes(2));
+    expect(mocks.register).toHaveBeenLastCalledWith(
+      expect.objectContaining({ stateSchema: changedSchema }),
+    );
+  });
+
+  it("keeps unsupported standard schemas from failing during render", async () => {
+    const createSchema = () =>
+      ({
+        "~standard": {
+          version: 1 as const,
+          vendor: "test",
+          validate: () => ({ value: {} }),
+        },
+      }) satisfies Unstable_InteractableStateSchema;
+
+    const hook = renderHook<
+      unknown,
+      { stateSchema: Unstable_InteractableStateSchema }
+    >(
+      ({ stateSchema }) =>
+        unstable_useInteractable("panel", {
+          id: "panel-1",
+          description: "A panel",
+          stateSchema,
+          initialState: {},
+        }),
+      { initialProps: { stateSchema: createSchema() } },
+    );
+    await waitFor(() => expect(mocks.register).toHaveBeenCalledTimes(1));
+
+    hook.rerender({ stateSchema: createSchema() });
+    expect(mocks.register).toHaveBeenCalledTimes(1);
+  });
+
+  it("converts a referentially stable schema only once", async () => {
+    const convert = vi.fn(() => ({
+      type: "object" as const,
+      properties: { value: { type: "string" as const } },
+    }));
+    const stateSchema: StandardSchemaWithJsonSchema = {
+      "~standard": {
+        version: 1,
+        vendor: "test",
+        validate: () => ({ value: {} }),
+        jsonSchema: { input: convert, output: convert },
+      },
+    };
+
+    const hook = renderHook<unknown, { initialState: Record<string, unknown> }>(
+      ({ initialState }) =>
+        unstable_useInteractable("panel", {
+          id: "panel-1",
+          description: "A panel",
+          stateSchema,
+          initialState,
+        }),
+      { initialProps: { initialState: { value: "first" } } },
+    );
+    await waitFor(() => expect(mocks.register).toHaveBeenCalledTimes(1));
+
+    hook.rerender({ initialState: { value: "second" } });
+    expect(convert).toHaveBeenCalledTimes(1);
+    expect(mocks.register).toHaveBeenCalledTimes(1);
+  });
+});

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

@@ -10,6 +10,7 @@ import {
 } from "../../model-context/interactable-composer-metadata";
 import { unstable_useInteractableState as useInteractableState } from "./useInteractableState";
 import { useJSONEqual } from "../utils/useJSONEqual";
+import { useJSONSchemaDependency } from "../utils/useJSONSchemaDependency";
 
 /**
  * The state type described by an interactable's `stateSchema`. Resolves the
@@ -116,6 +117,7 @@ const useInteractable = <TSchema extends Unstable_InteractableStateSchema>(
 
   const stateSchemaRef = useRef(config.stateSchema);
   stateSchemaRef.current = config.stateSchema;
+  const normalizedStateSchema = useJSONSchemaDependency(config.stateSchema);
   const initialStateRef = useRef(config.initialState);
   initialStateRef.current = config.initialState;
 
@@ -138,6 +140,7 @@ const useInteractable = <TSchema extends Unstable_InteractableStateSchema>(
     id,
     name,
     config.description,
+    normalizedStateSchema,
     internalScope,
     config.updateRender,
   ]);

package/src/react/primitive-hooks/useActionBarCopy.test.ts

@@ -21,13 +21,13 @@ const mocks = vi.hoisted(() => {
     },
     aui: {
       message: {
-        getCopyText: () => "Hello",
+        getCopyText: (): string => "Hello",
         setIsCopied,
       },
     },
     currentAui: {
       message: {
-        getCopyText: () => "Hello",
+        getCopyText: (): string => "Hello",
         setIsCopied,
       },
     },
@@ -37,15 +37,14 @@ const mocks = vi.hoisted(() => {
 vi.mock("@assistant-ui/store", async (importOriginal) => ({
   ...(await importOriginal<typeof import("@assistant-ui/store")>()),
   useAui: () => mocks.currentAui,
-  useAuiState: ((selector: (state: typeof mocks.state) => unknown) =>
-    selector(mocks.state)) as typeof import("@assistant-ui/store").useAuiState,
+  useAuiState: (selector: (state: typeof mocks.state) => unknown) =>
+    selector(mocks.state),
 }));
 
 import { useActionBarCopy } from "./useActionBarCopy";
 
 afterEach(() => {
   cleanup();
-  vi.clearAllMocks();
   vi.useRealTimers();
   mocks.currentAui = mocks.aui;
 });
@@ -176,7 +175,7 @@ describe("useActionBarCopy", () => {
 
     mocks.currentAui = {
       message: {
-        getCopyText: () => "Next message",
+        getCopyText: (): string => "Next message",
         setIsCopied: nextSetIsCopied,
       },
     };

package/src/react/primitive-hooks/useSuggestionTrigger.test.ts

@@ -47,7 +47,6 @@ import { useSuggestionTrigger } from "./useSuggestionTrigger";
 
 afterEach(() => {
   cleanup();
-  vi.clearAllMocks();
   mocks.state.thread.isDisabled = false;
   mocks.state.thread.isRunning = false;
   mocks.state.thread.capabilities = { queue: false };

package/src/react/primitives/generativeUI/GenerativeUI.test.tsx

@@ -0,0 +1,54 @@
+/** @vitest-environment jsdom */
+import { render } from "@testing-library/react";
+import type { ReactNode } from "react";
+import { afterEach, describe, expect, it, vi } from "vitest";
+import type { GenerativeUISpec } from "../../../types/message";
+import { GenerativeUIRender } from "./GenerativeUI";
+
+const Card = ({ children }: { children?: ReactNode }) => (
+  <section>{children}</section>
+);
+
+const renderRoot = (root: unknown) =>
+  render(
+    <GenerativeUIRender
+      spec={{ root } as GenerativeUISpec}
+      components={{ Card }}
+    />,
+  ).container.innerHTML;
+
+describe("GenerativeUIRender", () => {
+  afterEach(() => {
+    vi.restoreAllMocks();
+  });
+
+  it("renders a string children value as the only child", () => {
+    expect(renderRoot({ component: "Card", children: "Sunny" })).toBe(
+      "<section>Sunny</section>",
+    );
+  });
+
+  it("renders a node children value as the only child", () => {
+    expect(
+      renderRoot({
+        component: "Card",
+        children: { component: "Card", children: ["Sunny"] },
+      }),
+    ).toBe("<section><section>Sunny</section></section>");
+  });
+
+  it.each([
+    ["a number", 42],
+    ["an array-like object", { length: 1 }],
+  ])("skips %s children value as a malformed node", (_label, children) => {
+    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+
+    expect(renderRoot({ component: "Card", children })).toBe(
+      "<section></section>",
+    );
+    expect(warn).toHaveBeenCalledWith(
+      "[generative-ui] Skipping malformed node at 0/0:",
+      children,
+    );
+  });
+});

package/src/react/primitives/generativeUI/GenerativeUI.tsx

@@ -41,6 +41,15 @@ const isObjectNode = (
 ): node is Exclude<GenerativeUINode, string> =>
   typeof node === "object" && node !== null;
 
+const toNodeList = (
+  value: GenerativeUINode | readonly GenerativeUINode[] | null | undefined,
+): readonly GenerativeUINode[] => {
+  if (value === undefined || value === null) return [];
+  return Array.isArray(value)
+    ? (value as readonly GenerativeUINode[])
+    : [value as GenerativeUINode];
+};
+
 const renderNode = (
   node: GenerativeUINode | undefined,
   components: GenerativeUIComponentRegistry,
@@ -75,29 +84,15 @@ const renderNode = (
     throw new GenerativeUIRenderError(component);
   }
 
-  const renderedChildren = children?.length
-    ? children.map((child, i) =>
-        renderNode(child, components, Fallback, `${path}/${i}`),
-      )
-    : undefined;
-
   return createElement(
     Resolved,
     { ...(props ?? {}), key: key ?? path },
-    ...(renderedChildren ?? []),
+    ...toNodeList(children).map((child, i) =>
+      renderNode(child, components, Fallback, `${path}/${i}`),
+    ),
   );
 };
 
-const normalizeRoot = (
-  spec: GenerativeUISpec | undefined,
-): readonly GenerativeUINode[] => {
-  if (!spec || spec.root === undefined || spec.root === null) return [];
-  const root = spec.root;
-  return Array.isArray(root)
-    ? (root as readonly GenerativeUINode[])
-    : [root as GenerativeUINode];
-};
-
 /**
  * Internal renderer. Resolves a {@link GenerativeUISpec} against the consumer
  * allowlist. Used by `MessagePrimitive.GenerativeUI` and by
@@ -108,7 +103,7 @@ export const GenerativeUIRender: FC<GenerativeUIRenderProps> = ({
   components,
   Fallback,
 }) => {
-  const nodes = useMemo(() => normalizeRoot(spec), [spec]);
+  const nodes = useMemo(() => toNodeList(spec?.root), [spec]);
 
   return (
     <>

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

@@ -0,0 +1,99 @@
+// @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,
+): Msg["content"][number] => ({
+  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