@marimo-team/islands 0.15.2 → 0.15.3

package/src/__mocks__/notebook.ts

@@ -104,12 +104,14 @@ export const MockNotebook = {
       type: "exception",
       msg,
       exception_type,
+      raising_cell: null,
     }),
 
     strictException: (msg: string, ref: string): MarimoError => ({
       type: "strict-exception",
       msg,
       ref,
+      blamed_cell: null,
     }),
 
     interruption: (): MarimoError => ({
@@ -124,6 +126,7 @@ export const MockNotebook = {
     unknown: (msg: string): MarimoError => ({
       type: "unknown",
       msg,
+      error_type: null,
     }),
   },
 

package/src/__mocks__/requests.ts

@@ -44,6 +44,9 @@ export const MockRequestClient = {
       getUsageStats: vi.fn().mockResolvedValue({}),
       sendPdb: vi.fn().mockResolvedValue({}),
       sendListFiles: vi.fn().mockResolvedValue({ files: [] }),
+      sendSearchFiles: vi
+        .fn()
+        .mockResolvedValue({ files: [], query: "", total_found: 0 }),
       sendCreateFileOrFolder: vi.fn().mockResolvedValue({}),
       sendDeleteFileOrFolder: vi.fn().mockResolvedValue({}),
       sendRenameFileOrFolder: vi.fn().mockResolvedValue({}),

package/src/__tests__/__snapshots__/CellStatus.test.tsx.snap

@@ -39,14 +39,14 @@ exports[`CellStatusComponent > renders disabled and stale state 1`] = `
         width="24"
         xmlns="http://www.w3.org/2000/svg"
       >
+        <path
+          d="M4.929 4.929 19.07 19.071"
+        />
         <circle
           cx="12"
           cy="12"
           r="10"
         />
-        <path
-          d="m4.9 4.9 14.2 14.2"
-        />
       </svg>
       <div
         class="second-icon absolute bottom-[-2px] right-[-2px] rounded-full"
@@ -104,14 +104,14 @@ exports[`CellStatusComponent > renders disabled state 1`] = `
       width="24"
       xmlns="http://www.w3.org/2000/svg"
     >
+      <path
+        d="M4.929 4.929 19.07 19.071"
+      />
       <circle
         cx="12"
         cy="12"
         r="10"
       />
-      <path
-        d="m4.9 4.9 14.2 14.2"
-      />
     </svg>
   </div>
 </div>
@@ -175,14 +175,14 @@ exports[`CellStatusComponent > renders disabled transitively state 1`] = `
           width="24"
           xmlns="http://www.w3.org/2000/svg"
         >
+          <path
+            d="M4.929 4.929 19.07 19.071"
+          />
           <circle
             cx="12"
             cy="12"
             r="10"
           />
-          <path
-            d="m4.9 4.9 14.2 14.2"
-          />
         </svg>
       </div>
     </div>
@@ -403,14 +403,14 @@ exports[`CellStatusComponent > renders stale and disabled transitively state 1`]
           width="24"
           xmlns="http://www.w3.org/2000/svg"
         >
+          <path
+            d="M4.929 4.929 19.07 19.071"
+          />
           <circle
             cx="12"
             cy="12"
             r="10"
           />
-          <path
-            d="m4.9 4.9 14.2 14.2"
-          />
         </svg>
       </div>
     </div>

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

@@ -1,22 +1,17 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
+import type { UIMessage } from "ai";
 import { describe, expect, it } from "vitest";
 import { Maps } from "@/utils/maps";
-import { addMessageToChat } from "../core/ai/chat-utils";
+import { replaceMessagesInChat } from "../core/ai/chat-utils";
 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", () => {
+describe("replaceMessagesInChat", () => {
   const mockChatState: ChatState = {
     chats: asMap([
       {
@@ -26,223 +21,43 @@ describe("addMessageToChat", () => {
           {
             id: "msg-1",
             role: "user",
-            content: "Hello",
-            timestamp: 1000,
-          },
-          {
-            id: "msg-2",
-            role: "assistant",
-            content: "Hi there!",
-            timestamp: 2000,
+            parts: [{ type: "text", text: "Hello" }],
+            metadata: { timestamp: 1000 },
           },
         ],
         createdAt: 1000,
         updatedAt: 2000,
       },
-      {
-        id: CHAT_2,
-        title: "Test Chat 2",
-        messages: [
-          {
-            id: "msg-3",
-            role: "user",
-            content: "How are you?",
-            timestamp: 3000,
-          },
-        ],
-        createdAt: 3000,
-        updatedAt: 3000,
-      },
     ]),
     activeChatId: CHAT_1,
   };
 
-  it("should add a new message to an existing chat", () => {
-    const result = addMessageToChat(
-      mockChatState,
-      CHAT_1,
-      "msg-4",
-      "user",
-      "New message",
-    );
-
-    expect(result.chats).toHaveLength(2);
-    const updatedChat = result.chats.get(CHAT_1);
-    expect(updatedChat?.messages).toHaveLength(3);
-    expect(updatedChat?.messages[2]).toEqual({
-      id: "msg-4",
-      role: "user",
-      content: "New message",
-      timestamp: expect.any(Number),
+  it("replaces messages in a chat", () => {
+    const newMessages: UIMessage[] = [
+      {
+        id: "msg-2",
+        role: "assistant",
+        parts: [{ type: "text", text: "Hi there!" }],
+        metadata: { timestamp: 2000 },
+      },
+    ];
+    const result = replaceMessagesInChat({
+      chatState: mockChatState,
+      chatId: CHAT_1,
+      messages: newMessages,
     });
-    expect(updatedChat?.updatedAt).toBeGreaterThan(
-      first(mockChatState.chats).updatedAt,
+    expect(result.chats.get(CHAT_1)?.messages).toEqual(newMessages);
+    expect(result.chats.get(CHAT_1)?.updatedAt).toBeGreaterThan(
+      mockChatState.chats.get(CHAT_1)?.updatedAt ?? 0,
     );
   });
 
-  it("should update an existing message", () => {
-    const result = addMessageToChat(
-      mockChatState,
-      CHAT_1,
-      "msg-1",
-      "user",
-      "Updated content",
-    );
-
-    expect(result.chats).toHaveLength(2);
-    const updatedChat = result.chats.get(CHAT_1);
-    expect(updatedChat?.messages).toHaveLength(2);
-    expect(updatedChat?.messages[0]).toEqual({
-      id: "msg-1",
-      role: "user",
-      content: "Updated content",
-      timestamp: 1000,
+  it("returns unchanged state if chatId is null", () => {
+    const result = replaceMessagesInChat({
+      chatState: mockChatState,
+      chatId: null,
+      messages: [],
     });
-    expect(updatedChat?.updatedAt).toBeGreaterThan(
-      first(mockChatState.chats).updatedAt,
-    );
-  });
-
-  it("should handle message parts", () => {
-    const parts = [{ type: "text" as const, text: "Part content" }];
-    const result = addMessageToChat(
-      mockChatState,
-      CHAT_1,
-      "msg-5",
-      "assistant",
-      "Message with parts",
-      parts,
-    );
-
-    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: asMap([
-        {
-          ...chats[0],
-          messages: [
-            {
-              ...chats[0].messages[0],
-              parts: originalParts,
-            },
-            chats[0].messages[1],
-          ],
-        },
-        chats[1],
-      ]),
-    };
-
-    const result = addMessageToChat(
-      stateWithParts,
-      CHAT_1,
-      "msg-1",
-      "user",
-      "Updated content",
-      updatedParts,
-    );
-
-    const updatedChat = result.chats.get(CHAT_1);
-    expect(updatedChat?.messages[0].parts).toEqual(updatedParts);
-  });
-
-  it("should return unchanged state when chatId is null", () => {
-    const result = addMessageToChat(
-      mockChatState,
-      null,
-      "msg-4",
-      "user",
-      "New message",
-    );
-
-    expect(result).toEqual(mockChatState);
-  });
-
-  it("should return unchanged state when chatId does not exist", () => {
-    const result = addMessageToChat(
-      mockChatState,
-      "non-existent-chat" as ChatId,
-      "msg-4",
-      "user",
-      "New message",
-    );
-
     expect(result).toEqual(mockChatState);
   });
-
-  it("should not modify other chats when updating a specific chat", () => {
-    const result = addMessageToChat(
-      mockChatState,
-      CHAT_1,
-      "msg-4",
-      "user",
-      "New message",
-    );
-
-    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,
-      "msg-4",
-      "user",
-      "New message",
-    );
-
-    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: asMap([
-        {
-          id: chatId,
-          title: "Empty Chat",
-          messages: [],
-          createdAt: 1000,
-          updatedAt: 1000,
-        },
-      ]),
-      activeChatId: chatId,
-    };
-
-    const result = addMessageToChat(
-      emptyChatState,
-      chatId,
-      "msg-1",
-      "user",
-      "First message",
-    );
-
-    const updatedChat = result.chats.get(chatId);
-    expect(updatedChat?.messages).toHaveLength(1);
-    expect(updatedChat?.messages[0].content).toBe("First message");
-  });
-
-  it("should handle different message roles", () => {
-    const result = addMessageToChat(
-      mockChatState,
-      CHAT_1,
-      "msg-4",
-      "assistant",
-      "Assistant response",
-    );
-
-    const updatedChat = result.chats.get(CHAT_1);
-    expect(updatedChat?.messages[2].role).toBe("assistant");
-  });
 });

package/src/components/ai/ai-model-dropdown.tsx

@@ -10,6 +10,7 @@ import {
   CircleHelpIcon,
 } from "lucide-react";
 import React from "react";
+import { type SupportedRole, useModelChange } from "@/core/ai/config";
 import {
   AiModelId,
   isKnownAIProvider,
@@ -36,12 +37,13 @@ import { getCurrentRoleTooltip, getTagColour } from "./display-helpers";
 interface AIModelDropdownProps {
   value?: string;
   placeholder?: string;
-  onSelect: (modelId: QualifiedModelId) => void;
+  onSelect?: (modelId: QualifiedModelId) => void;
   triggerClassName?: string;
   customDropdownContent?: React.ReactNode;
   iconSize?: "medium" | "small";
   showAddCustomModelDocs?: boolean;
-  forRole?: Role;
+  displayIconOnly?: boolean;
+  forRole: SupportedRole;
 }
 
 export const AIModelDropdown = ({
@@ -53,12 +55,13 @@ export const AIModelDropdown = ({
   iconSize = "medium",
   showAddCustomModelDocs = false,
   forRole,
+  displayIconOnly = false,
 }: AIModelDropdownProps) => {
-  const currentValue = value ? AiModelId.parse(value) : undefined;
   const [isOpen, setIsOpen] = React.useState(false);
 
   const ai = useAtomValue(aiAtom);
   const completion = useAtomValue(completionAtom);
+  const { saveModelChange } = useModelChange();
 
   // Only include autocompleteModel if copilot is set to "custom"
   const autocompleteModel =
@@ -88,6 +91,13 @@ export const AIModelDropdown = ({
           ? ai?.models?.edit_model
           : undefined;
 
+  // If value is provided, use it, otherwise use the active model
+  const currentValue = value
+    ? AiModelId.parse(value)
+    : activeModel
+      ? AiModelId.parse(activeModel)
+      : undefined;
+
   const iconSizeClass = iconSize === "medium" ? "h-4 w-4" : "h-3 w-3";
 
   const renderModelWithRole = (modelId: AiModelId, role: Role) => {
@@ -119,7 +129,11 @@ export const AIModelDropdown = ({
   };
 
   const handleSelect = (modelId: QualifiedModelId) => {
-    onSelect(modelId);
+    if (onSelect) {
+      onSelect(modelId);
+    } else {
+      saveModelChange(modelId, forRole);
+    }
     setIsOpen(false);
   };
 
@@ -136,11 +150,13 @@ export const AIModelDropdown = ({
                 provider={currentValue.providerId}
                 className={iconSizeClass}
               />
-              <span className="truncate">
-                {isKnownAIProvider(currentValue.providerId)
-                  ? currentValue.shortModelId
-                  : currentValue.id}
-              </span>
+              {displayIconOnly ? null : (
+                <span className="truncate">
+                  {isKnownAIProvider(currentValue.providerId)
+                    ? currentValue.shortModelId
+                    : currentValue.id}
+                </span>
+              )}
             </>
           ) : (
             <span className="text-muted-foreground truncate">

package/src/components/app-config/ai-config.tsx

@@ -31,6 +31,7 @@ import { Input } from "@/components/ui/input";
 import { Kbd } from "@/components/ui/kbd";
 import { NativeSelect } from "@/components/ui/native-select";
 import { Textarea } from "@/components/ui/textarea";
+import type { SupportedRole } from "@/core/ai/config";
 import {
   AiModelId,
   PROVIDERS,
@@ -225,6 +226,7 @@ interface ModelSelectorProps {
   description?: React.ReactNode;
   disabled?: boolean;
   label: string;
+  forRole: SupportedRole;
   onSubmit: (values: UserConfig) => void;
 }
 
@@ -237,6 +239,7 @@ export const ModelSelector: React.FC<ModelSelectorProps> = ({
   description,
   disabled = false,
   label,
+  forRole,
   onSubmit,
 }) => {
   return (
@@ -289,6 +292,7 @@ export const ModelSelector: React.FC<ModelSelectorProps> = ({
                     </div>
                   </>
                 }
+                forRole={forRole}
               />
             </FormControl>
             <FormMessage />
@@ -420,6 +424,7 @@ const renderCopilotProvider = ({
         testId="custom-model-input"
         description="Model to use for code completion when using a custom provider."
         onSubmit={onSubmit}
+        forRole="autocomplete"
       />
     );
   }
@@ -905,6 +910,7 @@ export const AiAssistConfig: React.FC<AiConfigProps> = ({
         description={
           <span>Model to use for chat conversations in the Chat panel.</span>
         }
+        forRole="chat"
         onSubmit={onSubmit}
       />
       <ModelSelector
@@ -921,6 +927,7 @@ export const AiAssistConfig: React.FC<AiConfigProps> = ({
             <Kbd className="inline">Generate with AI</Kbd> button.
           </span>
         }
+        forRole="edit"
         onSubmit={onSubmit}
       />
 

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

@@ -0,0 +1,71 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import type { FileUIPart } from "ai";
+import { FileIcon, FileTextIcon, ImageIcon, XIcon } from "lucide-react";
+import { useState } from "react";
+import { cn } from "@/utils/cn";
+
+export const AttachmentRenderer = ({
+  attachment,
+}: {
+  attachment: FileUIPart;
+}) => {
+  if (attachment.mediaType?.startsWith("image/")) {
+    return (
+      <img
+        src={attachment.url}
+        alt={attachment.filename}
+        className="max-h-[100px] max-w-[100px] object-contain mb-1.5"
+      />
+    );
+  }
+
+  return (
+    <div className="flex flex-row gap-1 items-center text-xs">
+      <FileIcon className="h-3 w-3 mt-0.5" />
+      {attachment.filename}
+    </div>
+  );
+};
+
+export const FileAttachmentPill = ({
+  file,
+  className,
+  onRemove,
+}: {
+  file: File;
+  className?: string;
+  onRemove: () => void;
+}) => {
+  const [isHovered, setIsHovered] = useState(false);
+
+  return (
+    <div
+      className={cn(
+        "py-1 px-1.5 bg-muted rounded-md cursor-pointer flex flex-row gap-1 items-center text-xs",
+        className,
+      )}
+      onMouseEnter={() => setIsHovered(true)}
+      onMouseLeave={() => setIsHovered(false)}
+    >
+      {isHovered ? (
+        <XIcon className="h-3 w-3 mt-0.5" onClick={onRemove} />
+      ) : (
+        renderFileIcon(file)
+      )}
+      {file.name}
+    </div>
+  );
+};
+
+function renderFileIcon(file: File): React.ReactNode {
+  const classNames = "h-3 w-3 mt-0.5";
+
+  if (file.type.startsWith("image/")) {
+    return <ImageIcon className={classNames} />;
+  } else if (file.type.startsWith("text/")) {
+    return <FileTextIcon className={classNames} />;
+  }
+
+  return <FileIcon className={classNames} />;
+}