@marimo-team/frontend 0.15.1-dev14 → 0.15.1-dev16

package/src/__tests__/chat-utils.test.ts

@@ -1,14 +1,26 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
 import { describe, expect, it } from "vitest";
+import { Maps } from "@/utils/maps";
 import { addMessageToChat } from "../core/ai/chat-utils";
-import type { ChatState } from "../core/ai/state";
+import type { Chat, ChatId, ChatState } from "../core/ai/state";
+
+const CHAT_1 = "chat-1" as ChatId;
+const CHAT_2 = "chat-2" as ChatId;
+
+function first(map: Map<ChatId, Chat>) {
+  return [...map.values()][0];
+}
+
+function asMap(list: Iterable<Chat>) {
+  return Maps.keyBy(list, (c) => c.id);
+}
 
 describe("addMessageToChat", () => {
   const mockChatState: ChatState = {
-    chats: [
+    chats: asMap([
       {
-        id: "chat-1",
+        id: CHAT_1,
         title: "Test Chat 1",
         messages: [
           {
@@ -28,7 +40,7 @@ describe("addMessageToChat", () => {
         updatedAt: 2000,
       },
       {
-        id: "chat-2",
+        id: CHAT_2,
         title: "Test Chat 2",
         messages: [
           {
@@ -41,21 +53,21 @@ describe("addMessageToChat", () => {
         createdAt: 3000,
         updatedAt: 3000,
       },
-    ],
-    activeChatId: "chat-1",
+    ]),
+    activeChatId: CHAT_1,
   };
 
   it("should add a new message to an existing chat", () => {
     const result = addMessageToChat(
       mockChatState,
-      "chat-1",
+      CHAT_1,
       "msg-4",
       "user",
       "New message",
     );
 
     expect(result.chats).toHaveLength(2);
-    const updatedChat = result.chats.find((chat) => chat.id === "chat-1");
+    const updatedChat = result.chats.get(CHAT_1);
     expect(updatedChat?.messages).toHaveLength(3);
     expect(updatedChat?.messages[2]).toEqual({
       id: "msg-4",
@@ -64,21 +76,21 @@ describe("addMessageToChat", () => {
       timestamp: expect.any(Number),
     });
     expect(updatedChat?.updatedAt).toBeGreaterThan(
-      mockChatState.chats[0].updatedAt,
+      first(mockChatState.chats).updatedAt,
     );
   });
 
   it("should update an existing message", () => {
     const result = addMessageToChat(
       mockChatState,
-      "chat-1",
+      CHAT_1,
       "msg-1",
       "user",
       "Updated content",
     );
 
     expect(result.chats).toHaveLength(2);
-    const updatedChat = result.chats.find((chat) => chat.id === "chat-1");
+    const updatedChat = result.chats.get(CHAT_1);
     expect(updatedChat?.messages).toHaveLength(2);
     expect(updatedChat?.messages[0]).toEqual({
       id: "msg-1",
@@ -87,7 +99,7 @@ describe("addMessageToChat", () => {
       timestamp: 1000,
     });
     expect(updatedChat?.updatedAt).toBeGreaterThan(
-      mockChatState.chats[0].updatedAt,
+      first(mockChatState.chats).updatedAt,
     );
   });
 
@@ -95,48 +107,49 @@ describe("addMessageToChat", () => {
     const parts = [{ type: "text" as const, text: "Part content" }];
     const result = addMessageToChat(
       mockChatState,
-      "chat-1",
+      CHAT_1,
       "msg-5",
       "assistant",
       "Message with parts",
       parts,
     );
 
-    const updatedChat = result.chats.find((chat) => chat.id === "chat-1");
+    const updatedChat = result.chats.get(CHAT_1);
     expect(updatedChat?.messages[2].parts).toEqual(parts);
   });
 
   it("should update message parts", () => {
     const originalParts = [{ type: "text" as const, text: "Original" }];
     const updatedParts = [{ type: "text" as const, text: "Updated" }];
+    const chats = [...mockChatState.chats.values()];
 
     const stateWithParts: ChatState = {
       ...mockChatState,
-      chats: [
+      chats: asMap([
         {
-          ...mockChatState.chats[0],
+          ...chats[0],
           messages: [
             {
-              ...mockChatState.chats[0].messages[0],
+              ...chats[0].messages[0],
               parts: originalParts,
             },
-            mockChatState.chats[0].messages[1],
+            chats[0].messages[1],
           ],
         },
-        mockChatState.chats[1],
-      ],
+        chats[1],
+      ]),
     };
 
     const result = addMessageToChat(
       stateWithParts,
-      "chat-1",
+      CHAT_1,
       "msg-1",
       "user",
       "Updated content",
       updatedParts,
     );
 
-    const updatedChat = result.chats.find((chat) => chat.id === "chat-1");
+    const updatedChat = result.chats.get(CHAT_1);
     expect(updatedChat?.messages[0].parts).toEqual(updatedParts);
   });
 
@@ -155,7 +168,7 @@ describe("addMessageToChat", () => {
   it("should return unchanged state when chatId does not exist", () => {
     const result = addMessageToChat(
       mockChatState,
-      "non-existent-chat",
+      "non-existent-chat" as ChatId,
       "msg-4",
       "user",
       "New message",
@@ -167,54 +180,55 @@ describe("addMessageToChat", () => {
   it("should not modify other chats when updating a specific chat", () => {
     const result = addMessageToChat(
       mockChatState,
-      "chat-1",
+      CHAT_1,
       "msg-4",
       "user",
       "New message",
     );
 
-    const unchangedChat = result.chats.find((chat) => chat.id === "chat-2");
-    expect(unchangedChat).toEqual(mockChatState.chats[1]);
+    const unchangedChat = result.chats.get(CHAT_2);
+    expect(unchangedChat).toEqual([...mockChatState.chats.values()][1]);
   });
 
   it("should preserve message order when adding new messages", () => {
     const result = addMessageToChat(
       mockChatState,
-      "chat-1",
+      CHAT_1,
       "msg-4",
       "user",
       "New message",
     );
 
-    const updatedChat = result.chats.find((chat) => chat.id === "chat-1");
+    const updatedChat = result.chats.get(CHAT_1);
     expect(updatedChat?.messages[0].id).toBe("msg-1");
     expect(updatedChat?.messages[1].id).toBe("msg-2");
     expect(updatedChat?.messages[2].id).toBe("msg-4");
   });
 
   it("should handle empty chat messages array", () => {
+    const chatId = "empty-chat" as ChatId;
     const emptyChatState: ChatState = {
-      chats: [
+      chats: asMap([
         {
-          id: "empty-chat",
+          id: chatId,
           title: "Empty Chat",
           messages: [],
           createdAt: 1000,
           updatedAt: 1000,
         },
-      ],
-      activeChatId: "empty-chat",
+      ]),
+      activeChatId: chatId,
     };
 
     const result = addMessageToChat(
       emptyChatState,
-      "empty-chat",
+      chatId,
       "msg-1",
       "user",
       "First message",
     );
 
-    const updatedChat = result.chats.find((chat) => chat.id === "empty-chat");
+    const updatedChat = result.chats.get(chatId);
     expect(updatedChat?.messages).toHaveLength(1);
     expect(updatedChat?.messages[0].content).toBe("First message");
   });
@@ -222,13 +236,13 @@ describe("addMessageToChat", () => {
   it("should handle different message roles", () => {
     const result = addMessageToChat(
       mockChatState,
-      "chat-1",
+      CHAT_1,
       "msg-4",
       "assistant",
       "Assistant response",
     );
 
-    const updatedChat = result.chats.find((chat) => chat.id === "chat-1");
+    const updatedChat = result.chats.get(CHAT_1);
     expect(updatedChat?.messages[2].role).toBe("assistant");
   });
 });

package/src/components/chat/chat-panel.tsx

@@ -18,9 +18,11 @@ import {
   memo,
   type SetStateAction,
   useEffect,
+  useMemo,
   useRef,
   useState,
 } from "react";
+import useEvent from "react-use-event-hook";
 import { Button } from "@/components/ui/button";
 import {
   Popover,
@@ -41,6 +43,7 @@ import { useModelChange } from "@/core/ai/config";
 import {
   activeChatAtom,
   type Chat,
+  type ChatId,
   type ChatState,
   chatStateAtom,
 } from "@/core/ai/state";
@@ -68,8 +71,8 @@ import { ToolCallAccordion } from "./tool-call-accordion";
 
 interface ChatHeaderProps {
   onNewChat: () => void;
-  activeChatId: string | undefined;
-  setActiveChat: (id: string | null) => void;
+  activeChatId: ChatId | undefined;
+  setActiveChat: (id: ChatId | null) => void;
   chats: Chat[];
 }
 
@@ -151,11 +154,11 @@ interface ChatMessageProps {
   setChatState: Dispatch<SetStateAction<ChatState>>;
   chatState: ChatState;
   isStreamingReasoning: boolean;
-  totalMessages: number;
+  isLast: boolean;
 }
 
 const ChatMessage: React.FC<ChatMessageProps> = memo(
-  ({ message, index, onEdit, isStreamingReasoning, totalMessages }) => (
+  ({ message, index, onEdit, isStreamingReasoning, isLast }) => (
     <div
       className={cn(
         "flex group relative",
@@ -199,7 +202,7 @@ const ChatMessage: React.FC<ChatMessageProps> = memo(
                     key={i}
                     index={i}
                     isStreaming={
-                      index === totalMessages - 1 &&
+                      isLast &&
                       isStreamingReasoning &&
                       // If there are multiple reasoning parts, only show the last one
                       i === (message.parts?.length || 0) - 1
@@ -235,7 +238,7 @@ const ChatMessage: React.FC<ChatMessageProps> = memo(
 ChatMessage.displayName = "ChatMessage";
 
 interface ChatInputFooterProps {
-  input: string;
+  isEmpty: boolean;
   onSendClick: () => void;
   isLoading: boolean;
   onStop: () => void;
@@ -244,7 +247,7 @@ interface ChatInputFooterProps {
 const DEFAULT_MODE = "manual";
 
 const ChatInputFooter: React.FC<ChatInputFooterProps> = memo(
-  ({ input, onSendClick, isLoading, onStop }) => {
+  ({ isEmpty, onSendClick, isLoading, onStop }) => {
     const ai = useAtomValue(aiAtom);
     const currentMode = ai?.mode || DEFAULT_MODE;
     const currentModel = ai?.models?.chat_model || DEFAULT_AI_MODEL;
@@ -308,7 +311,7 @@ const ChatInputFooter: React.FC<ChatInputFooterProps> = memo(
           size="sm"
           className="h-6 w-6 p-0 hover:bg-muted/30"
           onClick={isLoading ? onStop : onSendClick}
-          disabled={isLoading ? false : !input.trim()}
+          disabled={isLoading ? false : isEmpty}
         >
           {isLoading ? (
             <SquareIcon className="h-3 w-3 fill-current" />
@@ -334,11 +337,11 @@ interface ChatInputProps {
 
 const ChatInput: React.FC<ChatInputProps> = memo(
   ({ input, setInput, onSubmit, inputRef, isLoading, onStop }) => {
-    const handleSendClick = () => {
+    const handleSendClick = useEvent(() => {
       if (input.trim()) {
         onSubmit(undefined, input);
       }
-    };
+    });
 
     return (
       <div className="border-t relative shrink-0 min-h-[80px] flex flex-col">
@@ -352,7 +355,7 @@ const ChatInput: React.FC<ChatInputProps> = memo(
           />
         </div>
         <ChatInputFooter
-          input={input}
+          isEmpty={!input.trim()}
           onSendClick={handleSendClick}
           isLoading={isLoading}
           onStop={onStop}
@@ -463,8 +466,9 @@ const ChatPanelBody = () => {
   const isLoading = status === "submitted" || status === "streaming";
 
   // Sync user messages from useChat to storage when they become available
+  // Only when we are done loading, for performance.
   useEffect(() => {
-    if (!chatState.activeChatId || messages.length === 0) {
+    if (!chatState.activeChatId || messages.length === 0 || isLoading) {
       return;
     }
 
@@ -474,9 +478,7 @@ const ChatPanelBody = () => {
       return;
     }
 
-    const currentChat = chatState.chats.find(
-      (c) => c.id === chatState.activeChatId,
-    );
+    const currentChat = chatState.chats.get(chatState.activeChatId);
     if (!currentChat) {
       return;
     }
@@ -505,31 +507,13 @@ const ChatPanelBody = () => {
         return result;
       });
     }
-  }, [messages, chatState.activeChatId, chatState.chats, setChatState]);
-
-  const isLastMessageReasoning = (messages: Message[]): boolean => {
-    if (messages.length === 0) {
-      return false;
-    }
-
-    const lastMessage = messages.at(-1);
-    if (!lastMessage) {
-      return false;
-    }
-
-    if (lastMessage.role !== "assistant" || !lastMessage.parts) {
-      return false;
-    }
-
-    const parts = lastMessage.parts;
-    if (parts.length === 0) {
-      return false;
-    }
-
-    // Check if the last part is reasoning
-    const lastPart = parts[parts.length - 1];
-    return lastPart.type === "reasoning";
-  };
+  }, [
+    messages,
+    chatState.activeChatId,
+    chatState.chats,
+    setChatState,
+    isLoading,
+  ]);
 
   // Check if we're currently streaming reasoning in the latest message
   const isStreamingReasoning =
@@ -550,7 +534,7 @@ const ChatPanelBody = () => {
   const createNewThread = (initialMessage: string) => {
     const CURRENT_TIME = Date.now();
     const newChat: Chat = {
-      id: generateUUID(),
+      id: generateUUID() as ChatId,
       title:
         initialMessage.length > 50
           ? `${initialMessage.slice(0, 50)}...`
@@ -562,53 +546,59 @@ const ChatPanelBody = () => {
 
     // Create new chat and set as active
     setChatState((prev) => {
+      const newChats = new Map(prev.chats);
+      newChats.set(newChat.id, newChat);
       const newState = {
         ...prev,
-        chats: [...prev.chats, newChat],
+        chats: newChats,
         activeChatId: newChat.id,
       };
       return newState;
     });
 
     // Trigger AI conversation with append
-    const MESSAGE_ID = generateUUID();
     append({
-      id: MESSAGE_ID,
+      id: generateUUID(),
       role: "user",
       content: initialMessage,
     });
     setInput("");
   };
 
-  const handleNewChat = () => {
+  const handleNewChat = useEvent(() => {
     setActiveChat(null);
     setInput("");
     setNewThreadInput("");
-  };
+  });
 
-  const handleMessageEdit = (index: number, newValue: string) => {
+  const handleMessageEdit = useEvent((index: number, newValue: string) => {
     // Truncate both useChat and storage
     setMessages((messages) => messages.slice(0, index));
-    if (chatState.activeChatId) {
-      setChatState((prev) => ({
-        ...prev,
-        chats: prev.chats.map((chat) =>
-          chat.id === chatState.activeChatId
-            ? {
-                ...chat,
-                messages: chat.messages.slice(0, index),
-                updatedAt: Date.now(),
-              }
-            : chat,
-        ),
-      }));
+    const activeChatId = chatState.activeChatId;
+    if (activeChatId) {
+      setChatState((prev) => {
+        const nextChats = new Map(prev.chats);
+        const activeChat = chatState.chats.get(activeChatId);
+        if (activeChat) {
+          nextChats.set(activeChat.id, {
+            ...activeChat,
+            messages: activeChat.messages.slice(0, index),
+            updatedAt: Date.now(),
+          });
+        }
+
+        return {
+          ...prev,
+          chats: nextChats,
+        };
+      });
     }
 
     append({
       role: "user",
       content: newValue,
     });
-  };
+  });
 
   const handleChatInputSubmit = (
     e: KeyboardEvent | undefined,
@@ -624,12 +614,21 @@ const ChatPanelBody = () => {
     reload();
   };
 
-  const handleNewThreadSubmit = () => {
-    newThreadInput.trim() && createNewThread(newThreadInput.trim());
-  };
+  const handleNewThreadSubmit = useEvent(() => {
+    if (!newThreadInput.trim()) {
+      return;
+    }
+    createNewThread(newThreadInput.trim());
+  });
 
   const handleOnCloseThread = () => newThreadInputRef.current?.editor?.blur();
 
+  const sortedChats = useMemo(() => {
+    return [...chatState.chats.values()].sort(
+      (a, b) => b.updatedAt - a.updatedAt,
+    );
+  }, [chatState.chats]);
+
   return (
     <div className="flex flex-col h-[calc(100%-53px)]">
       <TooltipProvider>
@@ -637,7 +636,7 @@ const ChatPanelBody = () => {
           onNewChat={handleNewChat}
           activeChatId={activeChat?.id}
           setActiveChat={setActiveChat}
-          chats={[...chatState.chats].sort((a, b) => b.updatedAt - a.updatedAt)}
+          chats={sortedChats}
         />
       </TooltipProvider>
 
@@ -658,7 +657,7 @@ const ChatPanelBody = () => {
               />
             </div>
             <ChatInputFooter
-              input={newThreadInput}
+              isEmpty={!newThreadInput.trim()}
               onSendClick={handleNewThreadSubmit}
               isLoading={isLoading}
               onStop={stop}
@@ -668,14 +667,14 @@ const ChatPanelBody = () => {
 
         {messages.map((message, idx) => (
           <ChatMessage
-            key={idx}
+            key={message.id}
             message={message}
             index={idx}
             onEdit={handleMessageEdit}
             setChatState={setChatState}
             chatState={chatState}
             isStreamingReasoning={isStreamingReasoning}
-            totalMessages={messages.length}
+            isLast={idx === messages.length - 1}
           />
         ))}
 
@@ -718,3 +717,27 @@ const ChatPanelBody = () => {
     </div>
   );
 };
+
+function isLastMessageReasoning(messages: Message[]): boolean {
+  if (messages.length === 0) {
+    return false;
+  }
+
+  const lastMessage = messages.at(-1);
+  if (!lastMessage) {
+    return false;
+  }
+
+  if (lastMessage.role !== "assistant" || !lastMessage.parts) {
+    return false;
+  }
+
+  const parts = lastMessage.parts;
+  if (parts.length === 0) {
+    return false;
+  }
+
+  // Check if the last part is reasoning
+  const lastPart = parts[parts.length - 1];
+  return lastPart.type === "reasoning";
+}

package/src/components/chat/markdown-renderer.tsx

@@ -81,19 +81,13 @@ function maybeTransform(
   };
 }
 
-const CodeBlock = ({ code, language }: CodeBlockProps) => {
-  const { theme } = useTheme();
+const InsertCodeBlockButton = ({ code, language }: CodeBlockProps) => {
   const { createNewCell } = useCellActions();
   const lastFocusedCellId = useLastFocusedCellId();
   const autoInstantiate = useAtomValue(autoInstantiateAtom);
-  const [value, setValue] = useState(code);
-
-  useEffect(() => {
-    setValue(code);
-  }, [code]);
 
   const handleInsertCode = () => {
-    const result = maybeTransform(language, value);
+    const result = maybeTransform(language, code);
 
     if (language === "sql") {
       maybeAddMarimoImport({
@@ -109,6 +103,22 @@ const CodeBlock = ({ code, language }: CodeBlockProps) => {
     });
   };
 
+  return (
+    <Button size="xs" variant="outline" onClick={handleInsertCode}>
+      Add to Notebook
+      <BetweenHorizontalStartIcon className="ml-2 h-4 w-4" />
+    </Button>
+  );
+};
+
+const CodeBlock = ({ code, language }: CodeBlockProps) => {
+  const { theme } = useTheme();
+  const [value, setValue] = useState(code);
+
+  useEffect(() => {
+    setValue(code);
+  }, [code]);
+
   const handleCopyCode = async () => {
     await copyToClipboard(value);
   };
@@ -132,10 +142,7 @@ const CodeBlock = ({ code, language }: CodeBlockProps) => {
         <CopyButton size="xs" variant="outline" onClick={handleCopyCode}>
           Copy
         </CopyButton>
-        <Button size="xs" variant="outline" onClick={handleInsertCode}>
-          Add to Notebook
-          <BetweenHorizontalStartIcon className="ml-2 h-4 w-4" />
-        </Button>
+        <InsertCodeBlockButton code={value} language={language} />
       </div>
     </div>
   );
@@ -164,17 +171,21 @@ const CopyButton: React.FC<ButtonProps> = ({ onClick, ...props }) => {
 };
 
 const COMPONENTS: Components = {
-  code: ({ children, className }) => {
+  code: ({ children, className, key }) => {
     const language = className?.replace("language-", "");
     if (language && typeof children === "string") {
       return (
-        <div>
+        <div key={key}>
           <div className="text-xs text-muted-foreground pl-1">{language}</div>
           <CodeBlock code={children.trim()} language={language} />
         </div>
       );
     }
-    return <code className={className}>{children}</code>;
+    return (
+      <code key={key} className={className}>
+        {children}
+      </code>
+    );
   },
 };
 
@@ -183,12 +194,14 @@ function parseMarkdownIntoBlocks(markdown: string): string[] {
   return tokens.map((token) => token.raw);
 }
 
+const PLUGINS = [remarkGfm];
+
 const MemoizedMarkdownBlock = memo(
   ({ content }: { content: string }) => {
     return (
       <Markdown
         components={COMPONENTS}
-        remarkPlugins={[remarkGfm]}
+        remarkPlugins={PLUGINS}
         className="prose dark:prose-invert max-w-none prose-pre:pl-0"
       >
         {content}