@assistant-ui/core 0.3.19 → 0.3.21

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

@@ -11,6 +11,7 @@ const mocks = vi.hoisted(() => {
       isDisabled: false,
       isRunning: false,
       capabilities: { queue: false },
+      voice: undefined as { canSendText: boolean } | undefined,
     },
   };
   const composerState = {
@@ -50,6 +51,7 @@ afterEach(() => {
   mocks.state.thread.isDisabled = false;
   mocks.state.thread.isRunning = false;
   mocks.state.thread.capabilities = { queue: false };
+  mocks.state.thread.voice = undefined;
   mocks.composerState.text = "";
 });
 
@@ -101,6 +103,37 @@ describe("useSuggestionTrigger", () => {
     expect(mocks.setText).not.toHaveBeenCalled();
   });
 
+  it("sends into a voice session that takes typed text while a spoken reply is running", () => {
+    mocks.state.thread.isRunning = true;
+    mocks.state.thread.voice = { canSendText: true };
+    mocks.composerState.text = "my draft";
+    const { result } = renderHook(() =>
+      useSuggestionTrigger({ prompt: "Hello", send: true }),
+    );
+
+    result.current.trigger();
+
+    expect(result.current.disabled).toBe(false);
+    expect(mocks.append).toHaveBeenCalledWith({
+      content: [{ type: "text", text: "Hello" }],
+      runConfig: { custom: { model: "gpt-test" } },
+    });
+    expect(mocks.setText).toHaveBeenCalledWith("");
+  });
+
+  it("disables and no-ops while a voice session cannot take typed text", () => {
+    mocks.state.thread.voice = { canSendText: false };
+    const { result } = renderHook(() =>
+      useSuggestionTrigger({ prompt: "Hello", send: true }),
+    );
+
+    result.current.trigger();
+
+    expect(result.current.disabled).toBe(true);
+    expect(mocks.append).not.toHaveBeenCalled();
+    expect(mocks.setText).not.toHaveBeenCalled();
+  });
+
   it("replaces the composer text when send is false, even while running", () => {
     mocks.state.thread.isRunning = true;
     mocks.composerState.text = "my draft";

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

@@ -1,6 +1,9 @@
 import { useCallback } from "react";
 import { useAui, useAuiState } from "@assistant-ui/store";
-import { suggestionTriggerDisabled } from "../../store/primitive-predicates";
+import {
+  suggestionSendMode,
+  suggestionTriggerDisabled,
+} from "../../store/primitive-predicates";
 
 export type UseSuggestionTriggerOptions = {
   prompt: string;
@@ -21,15 +24,15 @@ export const useSuggestionTrigger = ({
 
   const trigger = useCallback(() => {
     if (resolvedSend) {
-      const { isRunning, capabilities } = aui.thread.getState();
-      if (isRunning && !capabilities.queue) return;
+      const mode = suggestionSendMode(aui.thread.getState());
+      if (mode === "blocked") return;
 
       aui.thread.append({
         content: [{ type: "text", text: prompt }],
         runConfig: aui.composer.getState().runConfig,
       });
       // A queued send must not clear the draft the user is still composing.
-      if (clearComposer && !isRunning) {
+      if (clearComposer && mode === "now") {
         aui.composer.setText("");
       }
     } else {

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

@@ -0,0 +1,170 @@
+/** @vitest-environment jsdom */
+import { render } from "@testing-library/react";
+import { type ReactNode, useState } from "react";
+import { afterEach, describe, expect, it, vi } from "vitest";
+import type {
+  GenerativeUINode,
+  GenerativeUISpec,
+} from "../../../types/message";
+import type { GenerativeUIRenderProps } from "../../types/MessagePartComponentTypes";
+import { GenerativeUIRender, GenerativeUIRenderError } from "./GenerativeUI";
+
+const Card = ({ children }: { children?: ReactNode }) => (
+  <section>{children}</section>
+);
+
+const Fallback = ({ component }: { component: string }) => <i>{component}</i>;
+
+const inheritedNames = [
+  "constructor",
+  "hasOwnProperty",
+  "isPrototypeOf",
+  "toLocaleString",
+  "toString",
+  "valueOf",
+  "__proto__",
+];
+
+const renderRoot = (
+  root: unknown,
+  props: Partial<Pick<GenerativeUIRenderProps, "components" | "Fallback">> = {},
+) =>
+  render(
+    <GenerativeUIRender
+      spec={{ root } as GenerativeUISpec}
+      components={{ Card }}
+      {...props}
+    />,
+  ).container.innerHTML;
+
+const thrownBy = (fn: () => unknown): unknown => {
+  try {
+    fn();
+  } catch (error) {
+    return error;
+  }
+  return undefined;
+};
+
+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("renders number leaves and nested arrays recursively", () => {
+    const error = vi.spyOn(console, "error").mockImplementation(() => {});
+
+    expect(
+      renderRoot({
+        component: "Card",
+        children: ["Count: ", [42, { component: "Card", children: [7] }]],
+      }),
+    ).toBe("<section>Count: 42<section>7</section></section>");
+    expect(error).not.toHaveBeenCalled();
+  });
+
+  it("preserves keyed component identity when nested arrays reorder", () => {
+    let nextInstance = 0;
+    const StatefulCard = ({ label }: { label?: string }) => {
+      const [instance] = useState(() => ++nextInstance);
+      return <span>{`${label}:${instance}`}</span>;
+    };
+    const nodes = [
+      { component: "StatefulCard", key: "a", props: { label: "A" } },
+      { component: "StatefulCard", key: "b", props: { label: "B" } },
+    ] satisfies readonly GenerativeUINode[];
+    const view = render(
+      <GenerativeUIRender
+        spec={{
+          root: { component: "Card", children: [nodes] },
+        }}
+        components={{ Card, StatefulCard }}
+      />,
+    );
+
+    expect(view.container.textContent).toBe("A:1B:2");
+    view.rerender(
+      <GenerativeUIRender
+        spec={{
+          root: {
+            component: "Card",
+            children: [[nodes[1]!, nodes[0]!]],
+          },
+        }}
+        components={{ Card, StatefulCard }}
+      />,
+    );
+    expect(view.container.textContent).toBe("B:2A:1");
+  });
+
+  it("stops rendering arrays beyond the recursion limit", () => {
+    let root: GenerativeUINode = "too deep";
+    for (let depth = 0; depth < 66; depth += 1) root = [root];
+    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+
+    expect(renderRoot(root)).toBe("");
+    expect(warn).toHaveBeenCalledWith(
+      expect.stringContaining(
+        "[generative-ui] Skipping node nested past 64 levels at 0/",
+      ),
+    );
+  });
+
+  it("skips an array-like object children value as a malformed node", () => {
+    const children = { length: 1 };
+    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,
+    );
+  });
+
+  it.each(inheritedNames)(
+    "renders Fallback for the inherited name %s",
+    (component) => {
+      expect(renderRoot({ component }, { Fallback })).toBe(
+        `<i>${component}</i>`,
+      );
+    },
+  );
+
+  it.each(inheritedNames)(
+    "throws GenerativeUIRenderError for the inherited name %s",
+    (component) => {
+      vi.spyOn(console, "error").mockImplementation(() => {});
+
+      const error = thrownBy(() => renderRoot({ component }));
+
+      expect(error).toBeInstanceOf(GenerativeUIRenderError);
+      expect((error as GenerativeUIRenderError).componentName).toBe(component);
+    },
+  );
+
+  it("renders a component registered under an inherited name", () => {
+    expect(
+      renderRoot(
+        { component: "toString", children: "Sunny" },
+        { components: { toString: Card } },
+      ),
+    ).toBe("<section>Sunny</section>");
+  });
+});

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

@@ -38,36 +38,65 @@ export class GenerativeUIRenderError extends Error {
 
 const isObjectNode = (
   node: GenerativeUINode,
-): node is Exclude<GenerativeUINode, string> =>
-  typeof node === "object" && node !== null;
+): node is Extract<GenerativeUINode, { component: string }> =>
+  typeof node === "object" && node !== null && !Array.isArray(node);
+
+const toNodeList = (
+  value: GenerativeUINode | null | undefined,
+): readonly GenerativeUINode[] => {
+  if (value === undefined || value === null) return [];
+  return Array.isArray(value) ? value : [value];
+};
+
+const warn = (message: string, ...args: unknown[]) => {
+  if (
+    typeof process !== "undefined" &&
+    process.env?.NODE_ENV !== "production"
+  ) {
+    console.warn(message, ...args);
+  }
+};
+
+/** Bounds recursion so a runaway or adversarial model response cannot overflow the stack. */
+const MAX_DEPTH = 64;
 
 const renderNode = (
   node: GenerativeUINode | undefined,
   components: GenerativeUIComponentRegistry,
   Fallback: GenerativeUIRenderProps["Fallback"],
   path: string,
+  depth = 0,
 ): ReactNode => {
+  if (depth > MAX_DEPTH) {
+    warn(
+      `[generative-ui] Skipping node nested past ${MAX_DEPTH} levels at ${path}.`,
+    );
+    return null;
+  }
   if (node === undefined || node === null) return null;
 
-  if (typeof node === "string") return node;
+  if (typeof node === "string" || typeof node === "number") return node;
+
+  if (Array.isArray(node)) {
+    return node.map((child, i) =>
+      renderNode(child, components, Fallback, `${path}/${i}`, depth + 1),
+    );
+  }
 
   if (
     !isObjectNode(node) ||
     !("component" in node) ||
     typeof node.component !== "string"
   ) {
-    if (
-      typeof process !== "undefined" &&
-      process.env?.NODE_ENV !== "production"
-    ) {
-      console.warn(`[generative-ui] Skipping malformed node at ${path}:`, node);
-    }
+    warn(`[generative-ui] Skipping malformed node at ${path}:`, node);
     return null;
   }
 
   const { component, props, children, key } = node;
 
-  const Resolved = components[component];
+  const Resolved = Object.hasOwn(components, component)
+    ? components[component]
+    : undefined;
   if (!Resolved) {
     if (Fallback) {
       return <Fallback key={key ?? path} component={component} props={props} />;
@@ -75,29 +104,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}`, depth + 1),
+    ),
   );
 };
 
-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 +123,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/MessageAttachments.test.tsx

@@ -2,7 +2,7 @@
 import { render, screen } from "@testing-library/react";
 import { useState, type ReactNode } from "react";
 import { describe, expect, it, vi } from "vitest";
-import type { CompleteAttachment } from "../../../types/attachment";
+import type { Attachment, CompleteAttachment } from "../../../types/attachment";
 
 const mocks = vi.hoisted(() => ({
   role: "user" as "user" | "assistant",
@@ -56,11 +56,7 @@ const attachment = (id: string): CompleteAttachment => ({
   content: [],
 });
 
-const StatefulAttachment = ({
-  attachment,
-}: {
-  attachment: CompleteAttachment;
-}) => {
+const StatefulAttachment = ({ attachment }: { attachment: Attachment }) => {
   const [initialId] = useState(attachment.id);
   return <span>{initialId}</span>;
 };

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

@@ -1,4 +1,4 @@
-import type { CompleteAttachment } from "../../../types/attachment";
+import type { Attachment } from "../../../types/attachment";
 import {
   type ComponentType,
   type FC,
@@ -25,15 +25,19 @@ export namespace MessagePrimitiveAttachments {
         children?: never;
       }
     | {
-        /** Render function called for each attachment. Receives the attachment. */
-        children: (value: { attachment: CompleteAttachment }) => ReactNode;
+        /**
+         * Render function called for each attachment. Receives the attachment,
+         * which is still being prepared while the message is the composer's
+         * submission.
+         */
+        children: (value: { attachment: Attachment }) => ReactNode;
         components?: never;
       };
 }
 
 const getComponent = (
   components: MessageAttachmentsComponentConfig | undefined,
-  attachment: CompleteAttachment,
+  attachment: Attachment,
 ) => {
   const type = attachment.type;
   switch (type) {
@@ -54,7 +58,7 @@ const AttachmentComponent: FC<{
   const attachment = useAuiState((s) => s.attachment);
   if (!attachment) return null;
 
-  const Component = getComponent(components, attachment as CompleteAttachment);
+  const Component = getComponent(components, attachment);
   if (!Component) return null;
   return <Component />;
 };
@@ -90,12 +94,14 @@ MessagePrimitiveAttachmentByIndex.displayName =
   "MessagePrimitive.AttachmentByIndex";
 
 const MessagePrimitiveAttachmentsInner: FC<{
-  children: (value: { attachment: CompleteAttachment }) => ReactNode;
+  children: (value: { attachment: Attachment }) => ReactNode;
 }> = ({ children }) => {
   const attachmentIds = useAuiState(
     useShallowSelector((s) => {
       if (s.message.role !== "user") return [];
-      return (s.message.attachments ?? []).map((attachment) => attachment.id);
+      const attachments =
+        s.message.submission?.attachments ?? s.message.attachments;
+      return (attachments ?? []).map((attachment) => attachment.id);
     }),
   );
 
@@ -109,7 +115,7 @@ const MessagePrimitiveAttachmentsInner: FC<{
             {(getItem) =>
               children({
                 get attachment() {
-                  return getItem() as CompleteAttachment;
+                  return getItem();
                 },
               })
             }

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

@@ -1,25 +1,54 @@
 // @vitest-environment jsdom
 import { cleanup, render } from "@testing-library/react";
+import { useEffect } from "react";
 import { afterEach, describe, expect, it } from "vitest";
 import type { ThreadMessageLike } from "../../../runtime/utils/thread-message-like";
+import { useAui } from "@assistant-ui/store";
 import { AssistantRuntimeProvider } from "../../AssistantRuntimeProvider";
 import { ThreadPrimitiveMessages } from "../thread/ThreadMessages";
 import { useExternalStoreRuntime } from "../../runtimes/useExternalStoreRuntime";
 import { groupPartByType } from "../../utils/groupParts";
+import { useAssistantDataUI } from "../../model-context/useAssistantDataUI";
+import { useAssistantToolUI } from "../../model-context/useAssistantToolUI";
 import { MessagePrimitiveGroupedParts } from "./MessageGroupedParts";
 
+const NamedTool = () => <b>named</b>;
+const RegisterNamedTool = () => {
+  useAssistantToolUI({ toolName: "task", render: NamedTool });
+  return null;
+};
+
+const NamedData = ({ data }: { data: { value: string } }) => (
+  <b>{data.value}</b>
+);
+const RegisterNamedData = () => {
+  useAssistantDataUI({ name: "status", render: NamedData });
+  return null;
+};
+const RegisterFallbackData = () => {
+  const aui = useAui();
+  useEffect(() => aui.dataRenderers.setFallbackDataUI(NamedData), [aui]);
+  return null;
+};
+
 type Msg = {
   id: string;
-  content: readonly {
-    type: "tool-call";
-    toolCallId: string;
-    toolName: string;
-    args: {};
-    result?: { ok: true };
-  }[];
+  content: readonly (
+    | {
+        type: "tool-call";
+        toolCallId: string;
+        toolName: string;
+        args: {};
+        result?: { ok: true };
+      }
+    | { type: "data"; name: string; data: { value: string } }
+  )[];
 };
 
-const task = (toolCallId: string, hasResult: boolean) => ({
+const task = (
+  toolCallId: string,
+  hasResult: boolean,
+): Msg["content"][number] => ({
   type: "tool-call" as const,
   toolCallId,
   toolName: "task",
@@ -93,4 +122,122 @@ describe("MessagePrimitive.GroupedParts", () => {
     });
     expect(group?.indices).toHaveLength(3);
   });
+
+  it("renders registered tool UIs when the render function returns null", () => {
+    const GroupedParts = () => (
+      <>
+        <RegisterNamedTool />
+        <MessagePrimitiveGroupedParts groupBy={groupPartByType({})}>
+          {() => null}
+        </MessagePrimitiveGroupedParts>
+      </>
+    );
+
+    const App = () => {
+      const runtime = useExternalStoreRuntime<Msg>({
+        messages: [{ id: "assistant-1", content: [task("task-1", true)] }],
+        convertMessage,
+        onNew: async () => {},
+      });
+      return (
+        <AssistantRuntimeProvider runtime={runtime}>
+          <ThreadPrimitiveMessages components={{ Message: GroupedParts }} />
+        </AssistantRuntimeProvider>
+      );
+    };
+
+    expect(render(<App />).container.innerHTML).toContain("named");
+  });
+
+  it("renders a registered data UI when the render function returns null", () => {
+    const GroupedParts = () => (
+      <>
+        <RegisterNamedData />
+        <MessagePrimitiveGroupedParts groupBy={groupPartByType({})}>
+          {() => null}
+        </MessagePrimitiveGroupedParts>
+      </>
+    );
+
+    const App = () => {
+      const runtime = useExternalStoreRuntime<Msg>({
+        messages: [
+          {
+            id: "assistant-1",
+            content: [
+              { type: "data", name: "status", data: { value: "named" } },
+            ],
+          },
+        ],
+        convertMessage,
+        onNew: async () => {},
+      });
+      return (
+        <AssistantRuntimeProvider runtime={runtime}>
+          <ThreadPrimitiveMessages components={{ Message: GroupedParts }} />
+        </AssistantRuntimeProvider>
+      );
+    };
+
+    expect(render(<App />).container.innerHTML).toContain("named");
+  });
+
+  it("does not render registered UIs when a leaf returns an empty fragment", () => {
+    const GroupedParts = () => (
+      <>
+        <RegisterNamedTool />
+        <MessagePrimitiveGroupedParts groupBy={groupPartByType({})}>
+          {() => <></>}
+        </MessagePrimitiveGroupedParts>
+      </>
+    );
+
+    const App = () => {
+      const runtime = useExternalStoreRuntime<Msg>({
+        messages: [{ id: "assistant-1", content: [task("task-1", true)] }],
+        convertMessage,
+        onNew: async () => {},
+      });
+      return (
+        <AssistantRuntimeProvider runtime={runtime}>
+          <ThreadPrimitiveMessages components={{ Message: GroupedParts }} />
+        </AssistantRuntimeProvider>
+      );
+    };
+
+    expect(render(<App />).container.innerHTML).not.toContain("named");
+  });
+
+  it("renders the fallback data UI for an unregistered data part", () => {
+    const GroupedParts = () => (
+      <>
+        <RegisterFallbackData />
+        <MessagePrimitiveGroupedParts groupBy={groupPartByType({})}>
+          {() => null}
+        </MessagePrimitiveGroupedParts>
+      </>
+    );
+
+    const App = () => {
+      const runtime = useExternalStoreRuntime<Msg>({
+        messages: [
+          {
+            id: "assistant-1",
+            content: [
+              { type: "data", name: "unknown", data: { value: "fallback" } },
+            ],
+          },
+        ],
+        convertMessage,
+        onNew: async () => {},
+      });
+      return (
+        <AssistantRuntimeProvider runtime={runtime}>
+          <ThreadPrimitiveMessages components={{ Message: GroupedParts }} />
+        </AssistantRuntimeProvider>
+      );
+    };
+
+    expect(render(<App />).container.innerHTML).toContain("fallback");
+  });
 });