feat: useInlineRender hook (#1025)

.changeset/slimy-socks-unite.md

@@ -0,0 +1,5 @@
+---
+"@assistant-ui/react": patch
+---
+
+feat: useInlineRender hook

apps/docs/content/docs/advanced/ToolUI.mdx

@@ -92,19 +92,18 @@ const MyApp = () => {
 
 ### Inline Tool UI Hooks
 
-If you need access to component props, you can use the `useAssistantToolUI` hook.
+If you need access to component props, you can use the `useAssistantToolUI` hook. If you are passing a component inline, you should use the `useInlineRender` hook to prevent the component from being re-mounted on every render.
 
 ```tsx {5-11}
-import { useAssistantToolUI } from "@assistant-ui/react";
-import { useCallback } from "react";
+import { useAssistantToolUI, useInlineRender } from "@assistant-ui/react";
 
 const MyComponent = ({ product_id }) => {
   useAssistantToolUI({
     toolName: "current_product_info",
-    render: useCallback(({ args, status }) => {
+    render: useInlineRender(({ args, status }) => {
       // you can access component props here
       return <p>product_info({ product_id })</p>;
-    }, [product_id]),
+    }),
   });
 
   ...

packages/react/src/context/stores/AssistantToolUIs.ts

@@ -2,13 +2,14 @@
 
 import { create } from "zustand";
 import type { ToolCallContentPartComponent } from "../../types/ContentPartComponentTypes";
+import { Unsubscribe } from "../../types";
 
 export type AssistantToolUIsState = Readonly<{
   getToolUI: (toolName: string) => ToolCallContentPartComponent | null;
   setToolUI: (
     toolName: string,
     render: ToolCallContentPartComponent,
-  ) => () => void;
+  ) => Unsubscribe;
 }>;
 
 export const makeAssistantToolUIsStore = () =>

packages/react/src/model-config/index.ts

@@ -9,3 +9,4 @@ export {
   useAssistantToolUI,
   type AssistantToolUIProps,
 } from "./useAssistantToolUI";
+export { useInlineRender } from "./useInlineRender";

packages/react/src/model-config/useAssistantTool.tsx

@@ -24,22 +24,22 @@ export const useAssistantTool = <
 ) => {
   const assistantRuntime = useAssistantRuntime();
   const toolUIsStore = useToolUIsStore();
+
+  useEffect(() => {
+    return tool.render
+      ? toolUIsStore.getState().setToolUI(tool.toolName, tool.render)
+      : undefined;
+  }, [toolUIsStore, tool.toolName, tool.render]);
+
   useEffect(() => {
     const { toolName, render, ...rest } = tool;
     const config = {
       tools: {
-        [tool.toolName]: rest,
+        [toolName]: rest,
       },
     };
-    const unsub1 = assistantRuntime.registerModelConfigProvider({
+    return assistantRuntime.registerModelConfigProvider({
       getModelConfig: () => config,
     });
-    const unsub2 = render
-      ? toolUIsStore.getState().setToolUI(toolName, render)
-      : undefined;
-    return () => {
-      unsub1();
-      unsub2?.();
-    };
-  }, [assistantRuntime, toolUIsStore, tool]);
+  }, [assistantRuntime, tool]);
 };

packages/react/src/model-config/useAssistantToolUI.tsx

@@ -13,12 +13,12 @@ export type AssistantToolUIProps<
 };
 
 export const useAssistantToolUI = (
+  // TODO remove null option in 0.6
   tool: AssistantToolUIProps<any, any> | null,
 ) => {
   const toolUIsStore = useToolUIsStore();
   useEffect(() => {
     if (!tool) return;
-    const { toolName, render } = tool;
-    return toolUIsStore.getState().setToolUI(toolName, render);
-  }, [toolUIsStore, tool]);
+    return toolUIsStore.getState().setToolUI(tool.toolName, tool.render);
+  }, [toolUIsStore, tool?.toolName, tool?.render]);
 };

packages/react/src/model-config/useInlineRender.tsx

@@ -0,0 +1,18 @@
+import { FC, useCallback, useEffect, useState } from "react";
+import { ToolCallContentPartProps } from "../types";
+import { create } from "zustand";
+
+export const useInlineRender = <TArgs extends Record<string, unknown>, TResult>(
+  toolUI: FC<ToolCallContentPartProps<TArgs, TResult>>,
+): FC<ToolCallContentPartProps<TArgs, TResult>> => {
+  const [useToolUI] = useState(() => create(() => toolUI));
+
+  useEffect(() => {
+    useToolUI.setState(toolUI);
+  });
+
+  return useCallback((args) => {
+    const toolUI = useToolUI();
+    return toolUI(args);
+  }, []);
+};