fix: Stable state in tool components

frontend/app/components/GeneratingQuestionsToolUI.tsx

@@ -11,6 +11,7 @@ import {
 import { Sheet, SheetContent, SheetTrigger } from "./ui/sheet";
 import { TooltipIconButton } from "./ui/assistant-ui/tooltip-icon-button";
 import { DocumentCard, Document } from "./DocumentCard";
+import { useCallback } from "react";
 
 type Question = {
   question: string;
@@ -103,7 +104,9 @@ const QuestionCard = ({ question }: { question: Question }) => {
 export const useGeneratingQuestionsUI = () =>
   useAssistantToolUI({
     toolName: "generating_questions",
-    render: (input) => {
+    // Wrap the component in a useCallback to keep the identity stable.
+    // Allows the component to be interactable and not be re-rendered on every state change.
+    render: useCallback((input) => {
       if (!input.args?.questions || input.args.questions.length === 0) {
         return null;
       }
@@ -128,5 +131,5 @@ export const useGeneratingQuestionsUI = () =>
           </div>
         </div>
       );
-    },
+    }, []),
   });

frontend/app/components/ProgressToolUI.tsx

@@ -1,6 +1,7 @@
 import { useAssistantToolUI } from "@assistant-ui/react";
 import { Progress } from "./ui/progress";
 import { cn } from "../utils/cn";
+import { useCallback } from "react";
 
 export const stepToProgressFields = (step: number) => {
   switch (step) {
@@ -40,7 +41,9 @@ export const stepToProgressFields = (step: number) => {
 export const useProgressToolUI = () =>
   useAssistantToolUI({
     toolName: "progress",
-    render: (input) => {
+    // Wrap the component in a useCallback to keep the identity stable.
+    // Allows the component to be interactable and not be re-rendered on every state change.
+    render: useCallback((input) => {
       const { text, progress } = stepToProgressFields(input.args.step);
 
       return (
@@ -60,5 +63,5 @@ export const useProgressToolUI = () =>
           </p>
         </div>
       );
-    },
+    }, []),
   });

frontend/app/components/SelectedDocumentsToolUI.tsx

@@ -3,11 +3,14 @@ import { BookOpenText, Plus } from "lucide-react";
 import { Sheet, SheetContent, SheetTrigger } from "./ui/sheet";
 import { DocumentDialog } from "./DocumentDialog";
 import { DocumentCard, Document } from "./DocumentCard";
+import { useCallback } from "react";
 
 export const useSelectedDocumentsUI = () =>
   useAssistantToolUI({
     toolName: "selected_documents",
-    render: (input) => {
+    // Wrap the component in a useCallback to keep the identity stable.
+    // Allows the component to be interactable and not be re-rendered on every state change.
+    render: useCallback((input) => {
       if (!input.args?.documents || input.args.documents.length === 0) {
         return null;
       }
@@ -75,5 +78,5 @@ export const useSelectedDocumentsUI = () =>
           </div>
         </div>
       );
-    },
+    }, []),
   });