@assistant-ui/react 0.11.38 → 0.11.39

package/src/client/AssistantClient.ts

@@ -24,18 +24,14 @@ import {
   useAssistantApiImpl,
   extendApi,
 } from "../context/react/AssistantApiContext";
-import { ToolUIClient } from "./ToolUIClient";
 import { withEventsProvider } from "./EventContext";
 import { withModelContextProvider } from "./ModelContext";
-import { withToolUIProvider } from "./ToolUIContext";
-import { ToolUIApi, ToolUIState } from "./types/ToolUI";
 import { ToolsApi, ToolsState } from "./types/Tools";
 import { ModelContextApi, ModelContextState } from "./types/ModelContext";
 import { ModelContext as ModelContextResource } from "./ModelContextClient";
 
 type AssistantClientState = {
   readonly threads: ThreadListClientState;
-  readonly toolUIs: ToolUIState;
   readonly tools: ToolsState;
   readonly modelContext: ModelContextState;
 };
@@ -44,7 +40,6 @@ type AssistantClientApi = {
   getState(): AssistantClientState;
 
   readonly threads: ThreadListClientApi;
-  readonly toolUIs: ToolUIApi;
   readonly tools: ToolsApi;
   readonly modelContext: ModelContextApi;
 
@@ -61,44 +56,35 @@ const AssistantStore = resource(
     tools: toolsEl,
   }: AssistantClientProps) => {
     const events = tapInlineResource(EventManager());
-    const toolUIsResource = tapInlineResource(ToolUIClient());
 
-    const { threads, toolUIs, tools, modelContext } = withEventsProvider(
-      events,
-      () => {
-        const modelContextResource = tapResource(
-          modelContextEl ?? ModelContextResource(),
-          [modelContextEl],
-        );
+    const { threads, tools, modelContext } = withEventsProvider(events, () => {
+      const modelContextResource = tapResource(
+        modelContextEl ?? ModelContextResource(),
+        [modelContextEl],
+      );
 
-        return withModelContextProvider(modelContextResource.api, () => {
-          return withToolUIProvider(toolUIsResource.api, () => {
-            return {
-              toolUIs: toolUIsResource,
-              modelContext: modelContextResource,
-              tools: tapResource(toolsEl ?? Tools({}), [toolsEl]),
-              threads: tapResource(threadsEl, [threadsEl]),
-            };
-          });
-        });
-      },
-    );
+      return withModelContextProvider(modelContextResource.api, () => {
+        return {
+          modelContext: modelContextResource,
+          tools: tapResource(toolsEl ?? Tools({}), [toolsEl]),
+          threads: tapResource(threadsEl, [threadsEl]),
+        };
+      });
+    });
 
     const state = tapMemo<AssistantClientState>(
       () => ({
         threads: threads.state,
-        toolUIs: toolUIs.state,
         tools: tools.state,
         modelContext: modelContext.state,
       }),
-      [threads.state, toolUIs.state, tools.state, modelContext.state],
+      [threads.state, tools.state, modelContext.state],
     );
 
     return tapApi<AssistantClientApi>({
       getState: () => state,
 
       threads: threads.api,
-      toolUIs: toolUIs.api,
       tools: tools.api,
       modelContext: modelContext.api,
       on: events.on,
@@ -116,11 +102,6 @@ const getClientFromStore = (client: Store<{ api: AssistantClientApi }>) => {
       query: {},
       get: () => client.getState().api.threads,
     }),
-    toolUIs: createAssistantApiField({
-      source: "root",
-      query: {},
-      get: () => client.getState().api.toolUIs,
-    }),
     tools: createAssistantApiField({
       source: "root",
       query: {},

package/src/client/Tools.ts

@@ -1,41 +1,25 @@
 import { resource, tapState, tapEffect } from "@assistant-ui/tap";
 import { tapApi } from "../utils/tap-store";
 import { tapModelContext } from "./ModelContext";
-import { tapToolUI } from "./ToolUIContext";
 import { ToolsState, ToolsApi } from "./types/Tools";
 import type { Tool } from "assistant-stream";
-import {
-  type Toolkit,
-  FallbackSymbol,
-  LayoutSymbol,
-} from "../model-context/toolbox";
+import { type Toolkit } from "../model-context/toolbox";
 
 export const Tools = resource(({ toolkit }: { toolkit?: Toolkit }) => {
-  const [state] = tapState<ToolsState>(() => ({}));
+  const [state, setState] = tapState<ToolsState>(() => ({
+    tools: {},
+  }));
 
   const modelContext = tapModelContext();
-  const toolUI = tapToolUI();
 
   tapEffect(() => {
     if (!toolkit) return;
     const unsubscribes: (() => void)[] = [];
 
-    // Register fallback UI
-    const fallback = toolkit[FallbackSymbol];
-    if (fallback?.render) {
-      unsubscribes.push(toolUI.setFallbackToolUI(fallback.render));
-    }
-
-    // Register layout
-    const layout = toolkit[LayoutSymbol];
-    if (layout?.render) {
-      unsubscribes.push(toolUI.setToolUILayout(layout.render));
-    }
-
     // Register tool UIs (exclude symbols)
     for (const [toolName, tool] of Object.entries(toolkit)) {
       if (tool.render) {
-        unsubscribes.push(toolUI.setToolUI(toolName, tool.render));
+        unsubscribes.push(setToolUI(toolName, tool.render));
       }
     }
 
@@ -60,9 +44,34 @@ export const Tools = resource(({ toolkit }: { toolkit?: Toolkit }) => {
     return () => {
       unsubscribes.forEach((fn) => fn());
     };
-  }, [toolkit, modelContext, toolUI]);
+  }, [toolkit, modelContext]);
+
+  const setToolUI = (toolName: string, render: any) => {
+    setState((prev) => {
+      return {
+        ...prev,
+        tools: {
+          ...prev.tools,
+          [toolName]: [...(prev.tools[toolName] ?? []), render],
+        },
+      };
+    });
+
+    return () => {
+      setState((prev) => {
+        return {
+          ...prev,
+          tools: {
+            ...prev.tools,
+            [toolName]: prev.tools[toolName]?.filter((r) => r !== render) ?? [],
+          },
+        };
+      });
+    };
+  };
 
   return tapApi<ToolsApi>({
     getState: () => state,
+    setToolUI,
   });
 });

package/src/client/types/Tools.ts

@@ -1,8 +1,17 @@
-// eslint-disable-next-line @typescript-eslint/no-empty-object-type
-export type ToolsState = {};
+import { Unsubscribe } from "@assistant-ui/tap";
+import { ToolCallMessagePartComponent } from "../../types";
+
+export type ToolsState = {
+  tools: Record<string, ToolCallMessagePartComponent[]>;
+};
 
 export type ToolsApi = {
   getState(): ToolsState;
+
+  setToolUI(
+    toolName: string,
+    render: ToolCallMessagePartComponent,
+  ): Unsubscribe;
 };
 
 export type ToolsMeta = {

package/src/context/react/AssistantApiContext.tsx

@@ -10,7 +10,6 @@ import {
 } from "react";
 import { useResource } from "@assistant-ui/tap/react";
 
-import { ToolUIApi, ToolUIState, ToolUIMeta } from "../../client/types/ToolUI";
 import {
   MessageClientApi,
   MessageClientState,
@@ -49,7 +48,7 @@ import {
   AssistantClientProps,
   useAssistantClient,
 } from "../../client/AssistantClient";
-import { ToolsApi, ToolsMeta } from "../../client/types/Tools";
+import { ToolsApi, ToolsMeta, ToolsState } from "../../client/types/Tools";
 import {
   ModelContextApi,
   ModelContextMeta,
@@ -61,7 +60,7 @@ import {
 
 export type AssistantState = {
   readonly threads: ThreadListClientState;
-  readonly toolUIs: ToolUIState;
+  readonly tools: ToolsState;
 
   readonly threadListItem: ThreadListItemClientState;
   readonly thread: ThreadClientState;
@@ -122,7 +121,6 @@ type AttachmentMeta = {
 
 export type AssistantApi = {
   threads: AssistantApiField<ThreadListClientApi, ThreadsMeta>;
-  toolUIs: AssistantApiField<ToolUIApi, ToolUIMeta>;
   tools: AssistantApiField<ToolsApi, ToolsMeta>;
   modelContext: AssistantApiField<ModelContextApi, ModelContextMeta>;
   threadListItem: AssistantApiField<
@@ -168,13 +166,6 @@ const AssistantApiContext = createContext<AssistantApi>({
       throw new Error("Threads is only available inside <AssistantProvider />");
     },
   }),
-  toolUIs: createAssistantApiField({
-    source: null,
-    query: {},
-    get: (): never => {
-      throw new Error("ToolUIs is only available inside <AssistantProvider />");
-    },
-  }),
   tools: createAssistantApiField({
     source: null,
     query: {},

package/src/context/react/hooks/useAssistantState.tsx

@@ -15,8 +15,8 @@ class ProxiedAssistantState implements AssistantState {
     return this.#api.threads().getState();
   }
 
-  get toolUIs() {
-    return this.#api.toolUIs().getState();
+  get tools() {
+    return this.#api.tools().getState();
   }
 
   get threadListItem() {

package/src/legacy-runtime/cloud/auiV0.ts

@@ -112,7 +112,7 @@ export const auiV0Encode = (message: ThreadMessage): AuiV0Message => {
         }
 
         default: {
-          const unhandledType: "image" | "file" | "audio" = type;
+          const unhandledType: "image" | "file" | "audio" | "data" = type;
           throw new Error(
             `Message part type not supported by aui/v0: ${unhandledType}`,
           );

package/src/legacy-runtime/runtime-cores/external-store/ThreadMessageLike.tsx

@@ -12,6 +12,7 @@ import {
   ThreadSystemMessage,
   CompleteAttachment,
   FileMessagePart,
+  DataMessagePart,
   Unstable_AudioMessagePart,
 } from "../../../types";
 import {
@@ -31,6 +32,7 @@ export type ThreadMessageLike = {
         | SourceMessagePart
         | ImageMessagePart
         | FileMessagePart
+        | DataMessagePart
         | Unstable_AudioMessagePart
         | {
             readonly type: "tool-call";
@@ -123,6 +125,9 @@ export const fromThreadMessageLike = (
               case "image":
                 return sanitizeImageContent(part);
 
+              case "data":
+                return part;
+
               case "tool-call": {
                 const { parentId, messages, ...basePart } = part;
                 const commonProps = {
@@ -182,7 +187,11 @@ export const fromThreadMessageLike = (
               return part;
 
             default: {
-              const unhandledType: "tool-call" | "reasoning" | "source" = type;
+              const unhandledType:
+                | "tool-call"
+                | "reasoning"
+                | "source"
+                | "data" = type;
               throw new Error(
                 `Unsupported user message part type: ${unhandledType}`,
               );

package/src/legacy-runtime/runtime-cores/local/LocalThreadRuntimeCore.tsx

@@ -152,6 +152,7 @@ export class LocalThreadRuntimeCore
     this._options.adapters.history?.append({
       parentId: message.parentId,
       message: newMessage,
+      ...(message.runConfig !== undefined && { runConfig: message.runConfig }),
     });
 
     const startRun = message.startRun ?? message.role === "user";
@@ -183,8 +184,6 @@ export class LocalThreadRuntimeCore
   ): Promise<void> {
     this.ensureInitialized();
 
-    this.repository.resetHead(parentId);
-
     // add assistant message
     const id = generateId();
     let message: ThreadAssistantMessage = {
@@ -254,7 +253,7 @@ export class LocalThreadRuntimeCore
     runConfig: RunConfig | undefined,
     runCallback?: ChatModelAdapter["run"],
   ) {
-    const messages = this.repository.getMessages();
+    const messages = parentId ? this.repository.getMessages(parentId) : [];
 
     // abort existing run
     this.abortController?.abort();
@@ -327,6 +326,10 @@ export class LocalThreadRuntimeCore
           type: "running",
         },
       });
+
+      // Switch to the new message branch right after adding it for the first time
+      this.repository.resetHead(message.id);
+      this._notifySubscribers();
     }
 
     try {
@@ -401,6 +404,7 @@ export class LocalThreadRuntimeCore
         await this._options.adapters.history?.append({
           parentId,
           message: message,
+          runConfig: this._lastRunConfig,
         });
       }
     }

package/src/legacy-runtime/runtime-cores/utils/MessageRepository.tsx

@@ -1,4 +1,5 @@
 import type { ThreadMessage } from "../../../types";
+import { RunConfig } from "../../../types/AssistantTypes";
 import { generateId, generateOptimisticId } from "../../../utils/idUtils";
 import { ThreadMessageLike } from "../external-store";
 import { getAutoStatus } from "../external-store/auto-status";
@@ -34,6 +35,8 @@ export type ExportedMessageRepositoryItem = {
   message: ThreadMessage;
   /** ID of the parent message, or null for root messages */
   parentId: string | null;
+  /** The run configuration for the message */
+  runConfig?: RunConfig;
 };
 
 /**
@@ -46,6 +49,7 @@ export type ExportedMessageRepository = {
   messages: Array<{
     message: ThreadMessage;
     parentId: string | null;
+    runConfig?: RunConfig;
   }>;
 };
 
@@ -248,10 +252,30 @@ export class MessageRepository {
 
   /**
    * Gets all messages in the current active branch, from root to head.
-   * @returns Array of messages in the current branch
+   * @param headId - Optional ID of the head message to get messages for. If not provided, uses the current head.
+   * @returns Array of messages in the specified branch
    */
-  getMessages() {
-    return this._messages.value;
+  getMessages(headId?: string) {
+    if (headId === undefined || headId === this.head?.current.id) {
+      return this._messages.value;
+    }
+
+    const headMessage = this.messages.get(headId);
+    if (!headMessage) {
+      throw new Error(
+        "MessageRepository(getMessages): Head message not found. This is likely an internal bug in assistant-ui.",
+      );
+    }
+
+    const messages = new Array<ThreadMessage>(headMessage.level + 1);
+    for (
+      let current: RepositoryMessage | null = headMessage;
+      current;
+      current = current.prev
+    ) {
+      messages[current.level] = current.current;
+    }
+    return messages;
   }
 
   /**

package/src/model-context/index.ts

@@ -19,12 +19,7 @@ export { tool } from "./tool";
 
 export { makeAssistantVisible } from "./makeAssistantVisible";
 
-export {
-  Toolkit,
-  type ToolDefinition,
-  type ToolkitFallback,
-  type ToolkitLayout,
-} from "./toolbox";
+export type { Toolkit, ToolDefinition } from "./toolbox";
 
 export { Tools } from "../client/Tools";
 

package/src/model-context/makeAssistantVisible.tsx

@@ -11,18 +11,21 @@ import {
   createContext,
   useContext,
 } from "react";
-import { z } from "zod";
 import { useAssistantApi } from "../context/react/AssistantApiContext";
 import { useComposedRefs } from "@radix-ui/react-compose-refs";
 import { tool } from "./tool";
 
-// TODO replace zod with json-schema so we can drop the zod dep
-
 const click = tool({
-  parameters: z.object({
-    clickId: z.string(),
-  }),
-  execute: async ({ clickId }) => {
+  parameters: {
+    type: "object",
+    properties: {
+      clickId: {
+        type: "string",
+      },
+    },
+    required: ["clickId"],
+  },
+  execute: async ({ clickId }: { clickId: string }) => {
     const escapedClickId = CSS.escape(clickId);
     const el = document.querySelector(`[data-click-id='${escapedClickId}']`);
     if (el instanceof HTMLElement) {
@@ -38,11 +41,19 @@ const click = tool({
 });
 
 const edit = tool({
-  parameters: z.object({
-    editId: z.string(),
-    value: z.string(),
-  }),
-  execute: async ({ editId, value }) => {
+  parameters: {
+    type: "object",
+    properties: {
+      editId: {
+        type: "string",
+      },
+      value: {
+        type: "string",
+      },
+    },
+    required: ["editId", "value"],
+  },
+  execute: async ({ editId, value }: { editId: string; value: string }) => {
     const escapedEditId = CSS.escape(editId);
     const el = document.querySelector(`[data-edit-id='${escapedEditId}']`);
     if (el instanceof HTMLInputElement || el instanceof HTMLTextAreaElement) {

package/src/model-context/toolbox.tsx

@@ -1,6 +1,5 @@
 "use client";
 
-import type { ComponentType, ReactNode } from "react";
 import type { Tool } from "assistant-stream";
 import type { ToolCallMessagePartComponent } from "../types/MessagePartComponentTypes";
 
@@ -11,26 +10,7 @@ export type ToolDefinition<
   render?: ToolCallMessagePartComponent<TArgs, TResult> | undefined;
 };
 
-export const FallbackSymbol = Symbol("Toolkit.Fallback");
-export const LayoutSymbol = Symbol("Toolkit.Layout");
-
-export type ToolkitFallback = {
-  render: ToolCallMessagePartComponent<unknown, unknown>;
-};
-
-export type ToolkitLayout = {
-  render: ComponentType<{ children: ReactNode }>;
-};
-
-export type Toolkit = Record<string, ToolDefinition<any, any>> & {
-  [FallbackSymbol]?: ToolkitFallback;
-  [LayoutSymbol]?: ToolkitLayout;
-};
-
-export const Toolkit = {
-  Fallback: FallbackSymbol,
-  Layout: LayoutSymbol,
-} as const;
+export type Toolkit = Record<string, ToolDefinition<any, any>>;
 
 export type ToolsConfig = {
   toolkit: Toolkit;

package/src/model-context/useAssistantTool.tsx

@@ -23,7 +23,7 @@ export const useAssistantTool = <
 
   useEffect(() => {
     if (!tool.render) return undefined;
-    return api.toolUIs().setToolUI(tool.toolName, tool.render);
+    return api.tools().setToolUI(tool.toolName, tool.render);
   }, [api, tool.toolName, tool.render]);
 
   useEffect(() => {

package/src/model-context/useAssistantToolUI.tsx

@@ -15,6 +15,6 @@ export const useAssistantToolUI = (
   const api = useAssistantApi();
   useEffect(() => {
     if (!tool?.toolName || !tool?.render) return undefined;
-    return api.toolUIs().setToolUI(tool.toolName, tool.render);
+    return api.tools().setToolUI(tool.toolName, tool.render);
   }, [api, tool?.toolName, tool?.render]);
 };

package/src/primitives/index.ts

@@ -15,5 +15,6 @@ export { useMessagePartReasoning } from "./messagePart/useMessagePartReasoning";
 export { useMessagePartSource } from "./messagePart/useMessagePartSource";
 export { useMessagePartFile } from "./messagePart/useMessagePartFile";
 export { useMessagePartImage } from "./messagePart/useMessagePartImage";
+export { useMessagePartData } from "./messagePart/useMessagePartData";
 export { useThreadViewportAutoScroll } from "./thread/useThreadViewportAutoScroll";
 export { useScrollLock } from "./reasoning";

package/src/primitives/message/MessageParts.tsx

@@ -252,9 +252,8 @@ const ToolUIDisplay = ({
 }: {
   Fallback: ToolCallMessagePartComponent | undefined;
 } & ToolCallMessagePartProps) => {
-  const Render = useAssistantState(({ toolUIs }) => {
-    const Render =
-      toolUIs.tools[props.toolName] ?? toolUIs.fallback ?? Fallback;
+  const Render = useAssistantState(({ tools }) => {
+    const Render = tools.tools[props.toolName] ?? Fallback;
     if (Array.isArray(Render)) return Render[0] ?? Fallback;
     return Render;
   });
@@ -338,6 +337,9 @@ const MessagePartComponent: FC<MessagePartComponentProps> = ({
     case "audio":
       return <Audio {...part} />;
 
+    case "data":
+      return null;
+
     default:
       const unhandledType: never = type;
       throw new Error(`Unknown message part type: ${unhandledType}`);

package/src/primitives/message/MessagePartsGrouped.tsx

@@ -220,9 +220,8 @@ const ToolUIDisplay = ({
 }: {
   Fallback: ToolCallMessagePartComponent | undefined;
 } & ToolCallMessagePartProps) => {
-  const Render = useAssistantState(({ toolUIs }) => {
-    const Render =
-      toolUIs.tools[props.toolName] ?? toolUIs.fallback ?? Fallback;
+  const Render = useAssistantState(({ tools }) => {
+    const Render = tools.tools[props.toolName] ?? Fallback;
     if (Array.isArray(Render)) return Render[0] ?? Fallback;
     return Render;
   });
@@ -305,6 +304,9 @@ const MessagePartComponent: FC<MessagePartComponentProps> = ({
     case "audio":
       return <Audio {...part} />;
 
+    case "data":
+      return null;
+
     default:
       const unhandledType: never = type;
       throw new Error(`Unknown message part type: ${unhandledType}`);

package/src/primitives/messagePart/useMessagePartData.tsx

@@ -0,0 +1,24 @@
+"use client";
+
+import { useAssistantState } from "../../context";
+import { DataMessagePart } from "../../types";
+
+// eslint-disable-next-line @typescript-eslint/no-unused-vars
+export const useMessagePartData = <T = any,>(name?: string) => {
+  const part = useAssistantState(({ part }) => {
+    if (part.type !== "data") {
+      return null;
+    }
+    return part as DataMessagePart<T>;
+  });
+
+  if (!part) {
+    return null;
+  }
+
+  if (name && part.name !== name) {
+    return null;
+  }
+
+  return part;
+};

package/src/types/AssistantTypes.ts

@@ -6,6 +6,7 @@ import type {
   SourceMessagePart,
   ImageMessagePart,
   FileMessagePart,
+  DataMessagePart,
   Unstable_AudioMessagePart,
   ToolCallMessagePart,
   ThreadUserMessagePart,
@@ -19,6 +20,7 @@ export type {
   SourceMessagePart,
   ImageMessagePart,
   FileMessagePart,
+  DataMessagePart,
   Unstable_AudioMessagePart,
   ToolCallMessagePart,
   ThreadUserMessagePart,

package/src/types/MessagePartTypes.ts

@@ -43,6 +43,12 @@ export type Unstable_AudioMessagePart = {
   };
 };
 
+export type DataMessagePart<T = any> = {
+  readonly type: "data";
+  readonly name: string;
+  readonly data: T;
+};
+
 export type ToolCallMessagePart<
   TArgs = ReadonlyJSONObject,
   TResult = unknown,
@@ -72,4 +78,5 @@ export type ThreadAssistantMessagePart =
   | ToolCallMessagePart
   | SourceMessagePart
   | FileMessagePart
-  | ImageMessagePart;
+  | ImageMessagePart
+  | DataMessagePart;

package/src/types/index.ts

@@ -12,6 +12,7 @@ export type {
   SourceMessagePart,
   ImageMessagePart,
   FileMessagePart,
+  DataMessagePart,
   Unstable_AudioMessagePart,
   ToolCallMessagePart,
   MessageStatus,

package/dist/client/ToolUIClient.d.ts

@@ -1,7 +0,0 @@
-import { ToolUIApi } from "./types/ToolUI";
-export declare const ToolUIClient: import("@assistant-ui/tap").ResourceElementConstructor<{
-    key: string | undefined;
-    state: any;
-    api: ToolUIApi;
-}, undefined>;
-//# sourceMappingURL=ToolUIClient.d.ts.map

package/dist/client/ToolUIClient.d.ts.map

@@ -1 +0,0 @@
-{"version":3,"file":"ToolUIClient.d.ts","sourceRoot":"","sources":["../../src/client/ToolUIClient.ts"],"names":[],"mappings":"AAEA,OAAO,EAAe,SAAS,EAAE,MAAM,gBAAgB,CAAC;AAExD,eAAO,MAAM,YAAY;;;;aAuEvB,CAAC"}