@assistant-ui/core 0.2.9 → 0.2.11

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

@@ -0,0 +1,182 @@
+import { describe, expect, it } from "vitest";
+import { makeToolCallTextComponent } from "./toolbox";
+
+describe("makeToolCallTextComponent", () => {
+  it("renders static running and complete text", () => {
+    const Render = makeToolCallTextComponent({
+      running: "Searching...",
+      complete: "Done searching",
+    });
+
+    expect(
+      Render({
+        type: "tool-call",
+        toolCallId: "call-1",
+        toolName: "search",
+        args: {},
+        argsText: "{}",
+        status: { type: "running" },
+        addResult: () => {},
+        resume: () => {},
+        respondToApproval: () => {},
+      }),
+    ).toBe("Searching...");
+
+    expect(
+      Render({
+        type: "tool-call",
+        toolCallId: "call-1",
+        toolName: "search",
+        args: {},
+        argsText: "{}",
+        result: "ok",
+        status: { type: "complete" },
+        addResult: () => {},
+        resume: () => {},
+        respondToApproval: () => {},
+      }),
+    ).toBe("Done searching");
+  });
+
+  it("passes args and result to dynamic text functions", () => {
+    const Render = makeToolCallTextComponent<
+      { query: string },
+      { count: number }
+    >({
+      running: ({ args }) => `Searching ${args.query}...`,
+      complete: ({ args, result }) =>
+        `Found ${result?.count ?? 0} results for ${args.query}`,
+    });
+
+    expect(
+      Render({
+        type: "tool-call",
+        toolCallId: "call-1",
+        toolName: "search",
+        args: { query: "docs" },
+        argsText: '{"query":"docs"}',
+        status: { type: "running" },
+        addResult: () => {},
+        resume: () => {},
+        respondToApproval: () => {},
+      }),
+    ).toBe("Searching docs...");
+
+    expect(
+      Render({
+        type: "tool-call",
+        toolCallId: "call-1",
+        toolName: "search",
+        args: { query: "docs" },
+        argsText: '{"query":"docs"}',
+        result: { count: 3 },
+        status: { type: "complete" },
+        addResult: () => {},
+        resume: () => {},
+        respondToApproval: () => {},
+      }),
+    ).toBe("Found 3 results for docs");
+  });
+
+  it("treats missing status as complete", () => {
+    const Render = makeToolCallTextComponent({
+      running: "Running",
+      complete: "Complete",
+    });
+
+    expect(
+      Render({
+        type: "tool-call",
+        toolCallId: "call-1",
+        toolName: "search",
+        args: {},
+        argsText: "{}",
+        addResult: () => {},
+        resume: () => {},
+        respondToApproval: () => {},
+      }),
+    ).toBe("Complete");
+  });
+
+  it("treats incomplete status as a terminal state", () => {
+    const Render = makeToolCallTextComponent({
+      running: "Searching...",
+      complete: "Finished",
+    });
+
+    expect(
+      Render({
+        type: "tool-call",
+        toolCallId: "call-1",
+        toolName: "search",
+        args: {},
+        argsText: "{}",
+        status: { type: "incomplete", reason: "error" },
+        addResult: () => {},
+        resume: () => {},
+        respondToApproval: () => {},
+      }),
+    ).toBe("Finished");
+  });
+
+  it("treats requires-action status as running", () => {
+    const Render = makeToolCallTextComponent({
+      running: "Waiting for approval...",
+      complete: "Done",
+    });
+
+    expect(
+      Render({
+        type: "tool-call",
+        toolCallId: "call-1",
+        toolName: "confirm",
+        args: {},
+        argsText: "{}",
+        status: { type: "requires-action", reason: "interrupt" },
+        addResult: () => {},
+        resume: () => {},
+        respondToApproval: () => {},
+      }),
+    ).toBe("Waiting for approval...");
+  });
+
+  it("renders null for terminal status when only running text is provided", () => {
+    const Render = makeToolCallTextComponent({
+      running: "Searching...",
+    });
+
+    expect(
+      Render({
+        type: "tool-call",
+        toolCallId: "call-1",
+        toolName: "search",
+        args: {},
+        argsText: "{}",
+        status: { type: "complete" },
+        addResult: () => {},
+        resume: () => {},
+        respondToApproval: () => {},
+      }),
+    ).toBeNull();
+  });
+
+  it("renders null for running status when only complete text is provided", () => {
+    const Render = makeToolCallTextComponent({
+      complete: "Done",
+    });
+
+    expect(
+      Render({
+        type: "tool-call",
+        toolCallId: "call-1",
+        toolName: "search",
+        args: {},
+        argsText: "{}",
+        status: { type: "running" },
+        addResult: () => {},
+        resume: () => {},
+        respondToApproval: () => {},
+      }),
+    ).toBeNull();
+  });
+});

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

@@ -1,5 +1,14 @@
-import type { Tool, ToolDeclaration } from "assistant-stream";
-import type { ToolCallMessagePartComponent } from "../types/MessagePartComponentTypes";
+import type {
+  Tool,
+  ToolCallReader,
+  ToolDeclaration,
+  ToolModelOutputFunction,
+} from "assistant-stream";
+import type { ReactNode } from "react";
+import type {
+  ToolCallMessagePartComponent,
+  ToolCallMessagePartProps,
+} from "../types/MessagePartComponentTypes";
 
 /**
  * Resolves whether a tool's UI should be presented standalone (outside the
@@ -20,23 +29,96 @@ export const isStandaloneToolDisplay = (
 type WithRender<T, TArgs extends Record<string, unknown>, TResult> = T extends {
   type: "frontend" | "human";
 }
-  ? T & { render: ToolCallMessagePartComponent<TArgs, TResult> }
+  ? T &
+      (T extends { type: "frontend" }
+        ?
+            | { render: ToolCallMessagePartComponent<TArgs, TResult> }
+            | {
+                render?: ToolCallMessagePartComponent<TArgs, TResult>;
+                renderText: ToolCallText<TArgs, TResult>;
+              }
+        : { render: ToolCallMessagePartComponent<TArgs, TResult> })
   : T & {
       render?: ToolCallMessagePartComponent<TArgs, TResult> | undefined;
+      renderText?: ToolCallText<TArgs, TResult> | undefined;
     };
 
+type ToolParameters<TArgs extends Record<string, unknown>> =
+  ToolDeclaration<TArgs>["parameters"];
+
+// ToolExecutionContext is not re-exported from assistant-stream's public entry.
+type ToolExecuteContext = Parameters<
+  NonNullable<ToolDeclaration["execute"]>
+>[1];
+
+type ToolExecute<TArgs extends Record<string, unknown>, TResult> = (
+  args: TArgs,
+  context: ToolExecuteContext,
+) => TResult | Promise<TResult>;
+
+type ToolStreamCall<TArgs extends Record<string, unknown>, TResult> = (
+  reader: ToolCallReader<TArgs, 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,
+>(
+  text: ToolCallText<TArgs, TResult>,
+): ToolCallMessagePartComponent<TArgs, TResult> => {
+  return function ToolCallTextComponent(part) {
+    return resolveToolCallText(text, part);
+  };
+};
+
 /**
  * Tool definition accepted by the React tool registry.
  *
- * Extends the core tool contract with a render component. Human tools rely on
- * the renderer to collect input from the user. Frontend tools execute in the
- * browser and require a UI surface for their progress and result. Backend
- * tools execute server-side and may omit a renderer. The `render` component is
- * required for frontend and human tools and optional for backend tools.
+ * Extends the core tool contract with tool-call display options. Human tools
+ * rely on `render` to collect input from the user. Frontend tools execute in
+ * the browser and require either `render` or `renderText` for their progress
+ * and result. Backend tools execute server-side and may omit a renderer.
  */
 export type ToolDefinition<
-  TArgs extends Record<string, unknown>,
-  TResult,
+  TArgs extends Record<string, unknown> = Record<string, unknown>,
+  TResult = unknown,
 > = WithRender<Tool<TArgs, TResult>, TArgs, TResult>;
 
 /**
@@ -46,7 +128,7 @@ export type ToolDefinition<
  *
  * @example
  * ```tsx
- * const toolkit = {
+ * const toolkit = defineToolkit({
  *   get_weather: {
  *     type: "frontend",
  *     description: "Get the weather for a city.",
@@ -54,7 +136,7 @@ export type ToolDefinition<
  *     execute: async ({ city }: { city: string }) => fetchWeather(city),
  *     render: WeatherToolUI,
  *   },
- * } satisfies Toolkit;
+ * });
  * ```
  */
 export type Toolkit = Record<string, ToolDefinition<any, any>>;
@@ -63,19 +145,96 @@ export type Toolkit = Record<string, ToolDefinition<any, any>>;
  * A tool as authored, before the build splits it: like {@link ToolDefinition}
  * but it may declare `description`, `parameters`, and a server-side `execute`
  * alongside its `render`. The `type` field is **not** authored — the
- * `"use generative"` compiler infers it (`execute: hitl()` → human; `execute`
- * with a `"use client"` directive → frontend; otherwise backend) and writes it
- * back — so declaring it here is a type error.
+ * `"use generative"` compiler infers it (`execute: humanTool()` → human;
+ * `execute: providerTool(...)` → provider; `execute` with a `"use client"`
+ * directive → frontend; otherwise backend) and writes it back — so declaring it
+ * here is a type error.
  */
-export type ToolkitDeclarationDefinition<
+type OverrideOptionalField<
+  T,
+  TKey extends keyof T,
+  TValue,
+> = undefined extends T[TKey]
+  ? // Preserve `?: undefined` fields (for variants that explicitly disallow a
+    // callback) instead of widening them to accept the override value.
+    Exclude<T[TKey], undefined> extends never
+    ? { [K in TKey]?: undefined }
+    : { [K in TKey]?: TValue | undefined }
+  : { [K in TKey]: TValue };
+
+type OverrideToolDeclarationCallbacks<
+  T extends { streamCall?: unknown },
+  TArgs extends Record<string, unknown>,
+  TResult,
+> = Omit<
+  T,
+  | "type"
+  | "execute"
+  | "toModelOutput"
+  | "experimental_onSchemaValidationError"
+  | "streamCall"
+> & {
+  type?: never;
+} & ("execute" extends keyof T
+    ? OverrideOptionalField<T, "execute", ToolExecute<NoInfer<TArgs>, TResult>>
+    : {}) &
+  ("toModelOutput" extends keyof T
+    ? OverrideOptionalField<
+        T,
+        "toModelOutput",
+        ToolModelOutputFunction<NoInfer<TArgs>, NoInfer<TResult>>
+      >
+    : {}) &
+  ("experimental_onSchemaValidationError" extends keyof T
+    ? OverrideOptionalField<
+        T,
+        "experimental_onSchemaValidationError",
+        (
+          args: unknown,
+          context: ToolExecuteContext,
+        ) => NoInfer<TResult> | Promise<NoInfer<TResult>>
+      >
+    : {}) &
+  OverrideOptionalField<
+    T,
+    "streamCall",
+    ToolStreamCall<TArgs, NoInfer<TResult>>
+  >;
+
+// Keep the authored shape tied to ToolDeclaration's union variants while
+// overriding callback fields to avoid inference pollution.
+type ToolkitDefinitionInput<
   TArgs extends Record<string, unknown>,
   TResult,
 > = WithRender<
-  Omit<ToolDeclaration<TArgs, TResult>, "type">,
+  ToolDeclaration<TArgs, TResult> extends infer T
+    ? T extends { streamCall?: unknown }
+      ? OverrideToolDeclarationCallbacks<T, TArgs, TResult>
+      : never
+    : never,
   TArgs,
   TResult
-> & {
-  type?: never;
+>;
+
+/**
+ * A single entry in a {@link ToolkitDefinition}.
+ *
+ * Either authored inline (whose `type` the compiler infers) or an already-formed
+ * {@link ToolDefinition} produced by a factory whose own build splits it across
+ * targets — e.g. `new JSONGenerativeUI({ library }).present()`. The factory case
+ * carries a `type`, so it can only match the {@link ToolDefinition} arm of this
+ * union.
+ */
+export type ToolkitDefinitionEntry<
+  TArgs extends Record<string, unknown> = Record<string, unknown>,
+  TResult = unknown,
+> = ToolkitDefinitionInput<TArgs, TResult> | ToolDefinition<any, any>;
+
+export type ToolkitDefinitionEntryWithParameters<
+  TArgs extends Record<string, unknown> = Record<string, unknown>,
+  TResult = unknown,
+> = ToolkitDefinitionInput<TArgs, TResult> & {
+  parameters: NonNullable<ToolParameters<TArgs>>;
 };
 
 /**
@@ -83,10 +242,19 @@ export type ToolkitDeclarationDefinition<
  * {@link defineToolkit}. Backend entries may carry their server `execute` here;
  * the canonical {@link Toolkit} keeps those fields `undefined`.
  */
-export type ToolkitDeclaration = Record<
-  string,
-  ToolkitDeclarationDefinition<any, any>
->;
+export type ToolkitDefinition<
+  TArgsByName extends {
+    [K in keyof TArgsByName]: Record<string, unknown>;
+  } = Record<string, any>,
+  TResultByName extends { [K in keyof TArgsByName]: unknown } = {
+    [K in keyof TArgsByName]: any;
+  },
+> = {
+  [K in keyof TArgsByName]: ToolkitDefinitionEntry<
+    TArgsByName[K],
+    TResultByName[K]
+  >;
+};
 
 /** Configuration for the {@link Tools} resource. */
 export type ToolsConfig = {

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

@@ -1,11 +1,19 @@
-import { useEffect } from "react";
+import { useEffect, useMemo } from "react";
 import { useAui } from "@assistant-ui/store";
 import type { ToolCallMessagePartComponent } from "../types/MessagePartComponentTypes";
 import type { AssistantToolProps as CoreAssistantToolProps } from "../..";
-import { isStandaloneToolDisplay } from "./toolbox";
+import {
+  isStandaloneToolDisplay,
+  makeToolCallTextComponent,
+  type ToolCallText,
+} from "./toolbox";
 
 /**
  * Props used to register a tool from React.
+ *
+ * @deprecated Use a toolkit with `Tools({ toolkit })` and register it via
+ * `useAui({ tools: Tools({ toolkit }) })` instead. See
+ * https://assistant-ui.com/docs/migrations/toolkit-tools.
  */
 export type AssistantToolProps<
   TArgs extends Record<string, unknown>,
@@ -13,6 +21,8 @@ export type AssistantToolProps<
 > = CoreAssistantToolProps<TArgs, TResult> & {
   /** Component used to render calls to this tool in assistant messages. */
   render?: ToolCallMessagePartComponent<TArgs, TResult> | undefined;
+  /** Lightweight text rendered while a tool call is running or complete. */
+  renderText?: ToolCallText<TArgs, TResult> | undefined;
 };
 
 /**
@@ -28,6 +38,10 @@ export type AssistantToolProps<
  *
  * @param tool - Tool definition and name to register.
  *
+ * @deprecated Use a toolkit with `Tools({ toolkit })` and register it via
+ * `useAui({ tools: Tools({ toolkit }) })` instead. See
+ * https://assistant-ui.com/docs/migrations/toolkit-tools.
+ *
  * @example
  * ```tsx
  * const weatherTool = {
@@ -54,16 +68,22 @@ export const useAssistantTool = <
   const aui = useAui();
 
   const standalone = isStandaloneToolDisplay(tool);
+  const renderTextComponent = useMemo(
+    () =>
+      tool.renderText ? makeToolCallTextComponent(tool.renderText) : undefined,
+    [tool.renderText],
+  );
+  const render = tool.render ?? renderTextComponent;
 
   useEffect(() => {
-    if (!tool.render) return undefined;
-    return aui.tools().setToolUI(tool.toolName, tool.render, { standalone });
-  }, [aui, tool.toolName, tool.render, standalone]);
+    if (!render) return undefined;
+    return aui.tools().setToolUI(tool.toolName, render, { standalone });
+  }, [aui, tool.toolName, render, standalone]);
 
   useEffect(() => {
-    // `render` and `display` are client-only presentation concerns and never
-    // reach the model.
-    const { toolName, render, display, ...rest } = tool;
+    // `render`, `renderText`, and `display` are client-only presentation
+    // concerns and never reach the model.
+    const { toolName, render, renderText, display, ...rest } = tool;
     const context = {
       tools: {
         [toolName]: rest,

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

@@ -2,7 +2,13 @@ import { useEffect } from "react";
 import { useAui } from "@assistant-ui/store";
 import type { ToolCallMessagePartComponent } from "../types/MessagePartComponentTypes";
 
-/** Props used to register a renderer for tool-call message parts. */
+/**
+ * Props used to register a renderer for tool-call message parts.
+ *
+ * @deprecated Put `render`/`renderText` on the matching toolkit entry, or use
+ * `MessagePrimitive.Parts` inline tool render overrides for per-message UI.
+ * See https://assistant-ui.com/docs/migrations/toolkit-tools.
+ */
 export type AssistantToolUIProps<TArgs, TResult> = {
   /** Name of the tool whose calls should use this renderer. */
   toolName: string;
@@ -19,11 +25,14 @@ export type AssistantToolUIProps<TArgs, TResult> = {
 /**
  * Registers a tool-call renderer while the component is mounted.
  *
- * This only affects rendering. Pair it with {@link useAssistantTool},
- * {@link Tools}, or a backend tool registry to expose the actual tool
- * definition to the model.
+ * This only affects rendering. Pair it with {@link Tools} or a backend tool
+ * registry to expose the actual tool definition to the model.
  *
  * @param tool - Tool renderer registration, or `null` to skip registration.
+ *
+ * @deprecated Put `render`/`renderText` on the matching toolkit entry, or use
+ * `MessagePrimitive.Parts` inline tool render overrides for per-message UI.
+ * See https://assistant-ui.com/docs/migrations/toolkit-tools.
  */
 export const useAssistantToolUI = (
   tool: AssistantToolUIProps<any, any> | null,

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

@@ -0,0 +1,38 @@
+import { useEffect, useRef } from "react";
+import { useAui } from "@assistant-ui/store";
+import type { Tool } from "assistant-stream";
+
+type AuiToolOverride<
+  TArgs extends Record<string, unknown> = Record<string, unknown>,
+  TResult = unknown,
+> = Partial<Tool<TArgs, TResult>>;
+
+type AuiToolOverrides = Record<string, AuiToolOverride<any, any>>;
+
+/**
+ * Overrides toolkit entries for the current assistant scope.
+ *
+ * This is intended for dynamic local-state tools whose model-facing contract is
+ * declared in a `"use generative"` toolkit file with `execute: stubTool()`, but
+ * whose actual executor must close over React state in the mounted component.
+ * Keep the override keys stable after mount; dynamic key addition/removal is not
+ * currently observed.
+ * Overrides are registered at priority 1000, above toolkit defaults. Only one
+ * mounted override provider may define a given tool name at a time.
+ *
+ * @deprecated Experimental, API may change.
+ */
+export function useAuiToolOverrides(overrides: AuiToolOverrides): void {
+  const aui = useAui();
+  const overridesRef = useRef(overrides);
+  overridesRef.current = overrides;
+
+  useEffect(() => {
+    return aui.modelContext().register({
+      getModelContext: () => ({
+        priority: 1000,
+        tools: overridesRef.current as Record<string, Tool<any, any>>,
+      }),
+    });
+  }, [aui]);
+}

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

@@ -265,7 +265,7 @@ export const MessagePrimitiveGroupedParts = <TKey extends `group-${string}`>({
   const tree = useMemo(() => {
     const context: GroupByContext = { toolUIs };
     return buildGroupTree(parts.map((part) => groupBy(part, context) ?? []));
-    // oxlint-disable-next-line tap-hooks/exhaustive-deps -- groupBy is captured via memoDep (either its identity or the helper's memoKey fingerprint); listing it directly would defeat the helper-tagged memo path
+    // oxlint-disable-next-line react/exhaustive-deps -- groupBy is captured via memoDep (either its identity or the helper's memoKey fingerprint); listing it directly would defeat the helper-tagged memo path
   }, [parts, memoDep, toolUIs]);
 
   return (

package/src/react/primitives/part/PartMessages.tsx

@@ -34,16 +34,18 @@ const usePartMessages = (): readonly ThreadMessage[] | undefined => {
  *
  * @example
  * ```tsx
- * const SubAgentToolUI = makeAssistantToolUI({
- *   toolName: "invoke_sub_agent",
- *   render: () => (
- *     <PartPrimitive.Messages>
- *       {({ message }) => {
- *         if (message.role === "user") return <MyUserMessage />;
- *         return <MyAssistantMessage />;
- *       }}
- *     </PartPrimitive.Messages>
- *   ),
+ * const toolkit = defineToolkit({
+ *   invoke_sub_agent: {
+ *     type: "backend",
+ *     render: () => (
+ *       <PartPrimitive.Messages>
+ *         {({ message }) => {
+ *           if (message.role === "user") return <MyUserMessage />;
+ *           return <MyAssistantMessage />;
+ *         }}
+ *       </PartPrimitive.Messages>
+ *     ),
+ *   },
  * });
  * ```
  */