Merge pull request #327 from ryanhopperlowe/ui/fix/remove-flash-on-cr…

…eating/switching-threads

fix: remove flash on creating new thread and switching threads

ui/admin/app/components/chat/ChatContext.tsx

@@ -35,6 +35,7 @@ interface ChatContextType {
     invoke: (prompt?: string) => void;
     readOnly?: boolean;
     isRunning: boolean;
+    isLoading: boolean;
 }
 
 const ChatContext = createContext<ChatContextType | undefined>(undefined);
@@ -58,8 +59,15 @@ export function ChatProvider({
     const [generatingMessage, setGeneratingMessage] = useState<string | null>(
         null
     );
-    const [isRunning, setIsRunning] = useState(false);
+    const [isRunning, _setIsRunning] = useState(false);
+    const isRunningRef = useRef(false);
     const isRunningToolCall = useRef(false);
+
+    const setIsRunning = (value: boolean) => {
+        isRunningRef.current = value;
+        _setIsRunning(value);
+    };
+
     // todo(tylerslaton): this is a huge hack to get the generating message and runId to be
     // interactable during workflow invokes. take a look at invokeWorkflow to see why this is
     // currently needed.
@@ -83,19 +91,22 @@ export function ChatProvider({
         ({ threadId }) => ThreadsService.getThreadEvents(threadId),
         {
             onSuccess: () => setInsertedMessages([]),
-            revalidateIfStale: false,
             revalidateOnFocus: false,
             revalidateOnReconnect: false,
         }
     );
 
     const messages = useMemo(
-        () => chatEventsToMessages(getThreadEvents.data || []),
+        () =>
+            chatEventsToMessages(combineChatEvents(getThreadEvents.data || [])),
         [getThreadEvents.data]
     );
 
     // clear out inserted messages when the threadId changes
-    useEffect(() => setInsertedMessages([]), [threadId]);
+    useEffect(() => {
+        if (isRunningRef.current) return;
+        setInsertedMessages([]);
+    }, [threadId]);
 
     /** inserts message optimistically */
     const insertMessage = (message: Message) => {
@@ -149,6 +160,14 @@ export function ChatProvider({
 
             setIsRunning(true);
 
+            if (responseThreadId && !threadId) {
+                // persist the threadId
+                onCreateThreadId?.(responseThreadId);
+
+                // revalidate threads
+                mutate(ThreadsService.getThreads.key());
+            }
+
             readStream<ChatEvent>({
                 reader,
                 onChunk: (chunk) =>
@@ -184,29 +203,8 @@ export function ChatProvider({
                         }
                     }),
                 onComplete: async (chunks) => {
-                    const compactEvents = combineChatEvents(chunks);
-
-                    if (responseThreadId && !threadId) {
-                        // if this is a new thread, persist it by
-                        // prepopulating the cache with the events before setting the threadId
-                        // to avoid a flash of no messages
-                        await mutate(
-                            ThreadsService.getThreadEvents.key(
-                                responseThreadId
-                            ),
-                            compactEvents,
-                            { revalidate: false }
-                        );
-
-                        // persist the threadId
-                        onCreateThreadId?.(responseThreadId);
-
-                        // revalidate threads
-                        mutate(ThreadsService.getThreads.key());
-                        clearGeneratingMessage();
-                    } else {
-                        insertGeneratingMessage(chunks[0]?.runID);
-                    }
+                    insertGeneratingMessage(chunks[0]?.runID);
+                    clearGeneratingMessage();
 
                     invokeAgent.clear();
                     generatingRunIdRef.current = null;
@@ -255,6 +253,7 @@ export function ChatProvider({
                 generatingMessage: outGeneratingMessage,
                 invoke,
                 isRunning,
+                isLoading: getThreadEvents.isLoading,
                 readOnly,
             }}
         >

ui/admin/app/components/chat/MessagePane.tsx

@@ -6,6 +6,7 @@ import { Message } from "~/components/chat/Message";
 import { NoMessages } from "~/components/chat/NoMessages";
 import { ScrollArea } from "~/components/ui/scroll-area";
 
+import { LoadingSpinner } from "../ui/LoadingSpinner";
 import { useChat } from "./ChatContext";
 
 interface MessagePaneProps {
@@ -25,27 +26,23 @@ export function MessagePane({
     classNames = {},
     generatingMessage,
 }: MessagePaneProps) {
-    const { mode, readOnly } = useChat();
+    const { readOnly, isLoading } = useChat();
+
+    const isEmpty = messages.length === 0 && !generatingMessage && !readOnly;
 
     return (
         <div className={cn("flex flex-col h-full", className, classNames.root)}>
-            {messages.length === 0 &&
-            !generatingMessage &&
-            mode === "agent" &&
-            !readOnly ? (
-                <div className="flex-grow flex items-center justify-center">
+            <ScrollArea
+                startScrollAt="bottom"
+                enableScrollTo="bottom"
+                enableScrollStick="bottom"
+                className={cn("h-full w-full relative", classNames.messageList)}
+            >
+                {isLoading && isEmpty ? (
+                    <LoadingSpinner fillContainer />
+                ) : isEmpty ? (
                     <NoMessages />
-                </div>
-            ) : (
-                <ScrollArea
-                    startScrollAt="bottom"
-                    enableScrollTo="bottom"
-                    enableScrollStick="bottom"
-                    className={cn(
-                        "h-full w-full relative",
-                        classNames.messageList
-                    )}
-                >
+                ) : (
                     <div className="p-4 space-y-6 w-full">
                         {messages.map((message, i) => (
                             <Message key={i} message={message} />
@@ -54,8 +51,8 @@ export function MessagePane({
                             <Message message={generatingMessage} />
                         )}
                     </div>
-                </ScrollArea>
-            )}
+                )}
+            </ScrollArea>
         </div>
     );
 }

ui/admin/app/components/ui/LoadingSpinner.tsx

@@ -1,30 +1,38 @@
+import { Loader2 } from "lucide-react";
+import { ComponentProps } from "react";
+
 import { cn } from "~/lib/utils";
 
-export interface ISVGProps extends React.SVGProps<SVGSVGElement> {
+export interface LoadingSpinnerProps extends ComponentProps<typeof Loader2> {
     size?: number;
     className?: string;
+    fillContainer?: boolean;
+    classNames?: {
+        root?: string;
+        container?: string;
+    };
 }
 
 export const LoadingSpinner = ({
-    size = 24,
     className,
+    fillContainer,
+    classNames = {},
     ...props
-}: ISVGProps) => {
-    return (
-        <svg
-            xmlns="http://www.w3.org/2000/svg"
-            width={size}
-            height={size}
-            {...props}
-            viewBox="0 0 24 24"
-            fill="none"
-            stroke="currentColor"
-            strokeWidth="2"
-            strokeLinecap="round"
-            strokeLinejoin="round"
-            className={cn("animate-spin", className)}
+}: LoadingSpinnerProps) => {
+    const content = (
+        <Loader2 className={cn("animate-spin", className)} {...props} />
+    );
+
+    return fillContainer ? (
+        <div
+            className={cn(
+                "min-w-fit h-full flex-auto flex items-center justify-center",
+                classNames.container
+            )}
         >
-            <path d="M21 12a9 9 0 1 1-6.219-8.56" />
-        </svg>
+            {content}
+        </div>
+    ) : (
+        content
     );
 };

ui/admin/app/hooks/useRefState.ts

@@ -0,0 +1,13 @@
+import { useCallback, useRef, useState } from "react";
+
+export function useRefState<T>(initialValue: T) {
+    const ref = useRef(initialValue);
+    const [state, _setState] = useState(initialValue);
+
+    const setState = useCallback((value: T) => {
+        ref.current = value;
+        _setState(value);
+    }, []);
+
+    return [ref, state, setState] as const;
+}