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

package/src/core/ai/chat-utils.ts

@@ -2,11 +2,11 @@
 
 import type { Message as AIMessage } from "@ai-sdk/react";
 import { Logger } from "@/utils/Logger";
-import type { ChatState } from "./state";
+import type { ChatId, ChatState } from "./state";
 
 export const addMessageToChat = (
   chatState: ChatState,
-  chatId: string | null,
+  chatId: ChatId | null,
   messageId: string,
   role: "user" | "assistant",
   content: string,
@@ -16,26 +16,23 @@ export const addMessageToChat = (
     Logger.warn("No active chat");
     return chatState;
   }
-  const activeChatIndex = chatState.chats.findIndex(
-    (chat) => chat.id === chatId,
-  );
-  if (activeChatIndex === -1) {
+  const chat = chatState.chats.get(chatId);
+  if (!chat) {
     Logger.warn("No active chat");
     return chatState;
   }
 
-  const chat = chatState.chats[activeChatIndex];
   const messageIndex = chat.messages.findIndex(
     (message) => message.id === messageId,
   );
 
   // Create copy of chats to modify
-  const newChats = [...chatState.chats];
-  const CURRENT_TIMESTAMP = Date.now();
+  const newChats = new Map(chatState.chats);
+  const timestamp = Date.now();
 
   if (messageIndex === -1) {
     // Handle new message
-    newChats[activeChatIndex] = {
+    newChats.set(chatId, {
       ...chat,
       messages: [
         ...chat.messages,
@@ -43,12 +40,12 @@ export const addMessageToChat = (
           id: messageId,
           role,
           content,
-          timestamp: CURRENT_TIMESTAMP,
+          timestamp: timestamp,
           parts,
         },
       ],
-      updatedAt: CURRENT_TIMESTAMP,
-    };
+      updatedAt: timestamp,
+    });
   } else {
     // Handle update message
     const newMessages = [...chat.messages];
@@ -57,11 +54,11 @@ export const addMessageToChat = (
       content,
       parts,
     };
-    newChats[activeChatIndex] = {
+    newChats.set(chat.id, {
       ...chat,
       messages: newMessages,
-      updatedAt: CURRENT_TIMESTAMP,
-    };
+      updatedAt: timestamp,
+    });
   }
 
   return {

package/src/core/ai/state.ts

@@ -3,9 +3,13 @@
 import type { Message as AIMessage } from "@ai-sdk/react";
 import { atom } from "jotai";
 import { atomWithStorage } from "jotai/utils";
+import { adaptForLocalStorage } from "@/utils/storage";
+import type { TypedString } from "@/utils/typed";
 import type { CellId } from "../cells/ids";
 
-const KEY = "marimo:ai:chatState:v3";
+const KEY = "marimo:ai:chatState:v4";
+
+export type ChatId = TypedString<"ChatId">;
 
 export const aiCompletionCellAtom = atom<{
   cellId: CellId;
@@ -27,7 +31,7 @@ export interface Message {
 }
 
 export interface Chat {
-  id: string;
+  id: ChatId;
   title: string;
   messages: Message[];
   createdAt: number;
@@ -35,21 +39,37 @@ export interface Chat {
 }
 
 export interface ChatState {
-  chats: Chat[];
-  activeChatId: string | null;
+  chats: Map<ChatId, Chat>;
+  activeChatId: ChatId | null;
 }
 
-export const chatStateAtom = atomWithStorage<ChatState>(KEY, {
-  chats: [],
-  activeChatId: null,
-});
+export const chatStateAtom = atomWithStorage<ChatState>(
+  KEY,
+  {
+    chats: new Map(),
+    activeChatId: null,
+  },
+  adaptForLocalStorage({
+    toSerializable: (value: ChatState) => ({
+      chats: [...value.chats.entries()],
+      activeChatId: value.activeChatId,
+    }),
+    fromSerializable: (value) => ({
+      chats: new Map(value.chats),
+      activeChatId: value.activeChatId,
+    }),
+  }),
+);
 
 export const activeChatAtom = atom(
   (get) => {
     const state = get(chatStateAtom);
-    return state.chats.find((chat) => chat.id === state.activeChatId);
+    if (!state.activeChatId) {
+      return null;
+    }
+    return state.chats.get(state.activeChatId);
   },
-  (get, set, chatId: string | null) => {
+  (get, set, chatId: ChatId | null) => {
     set(chatStateAtom, (prev) => ({
       ...prev,
       activeChatId: chatId,

package/src/utils/__tests__/storage.test.ts

@@ -0,0 +1,144 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import { adaptForLocalStorage } from "../storage";
+
+// Mock localStorage
+const localStorageMock = {
+  getItem: vi.fn(),
+  setItem: vi.fn(),
+  removeItem: vi.fn(),
+};
+
+Object.defineProperty(global, "localStorage", {
+  value: localStorageMock,
+  writable: true,
+});
+
+// Mock Logger
+vi.mock("../Logger", () => ({
+  Logger: {
+    warn: vi.fn(),
+  },
+}));
+
+interface TestValue {
+  id: string;
+  data: Map<string, number>;
+}
+
+interface SerializableValue {
+  id: string;
+  data: Array<[string, number]>;
+}
+
+describe("adaptForLocalStorage", () => {
+  const storage = adaptForLocalStorage<TestValue, SerializableValue>({
+    toSerializable: (v) => ({
+      id: v.id,
+      data: [...v.data.entries()],
+    }),
+    fromSerializable: (s) => ({
+      id: s.id,
+      data: new Map(s.data),
+    }),
+  });
+
+  const testValue: TestValue = {
+    id: "test",
+    data: new Map([
+      ["key1", 1],
+      ["key2", 2],
+    ]),
+  };
+
+  const initialValue: TestValue = {
+    id: "initial",
+    data: new Map(),
+  };
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+  });
+
+  describe("getItem", () => {
+    it("should return parsed value when localStorage contains valid data", () => {
+      const serialized = JSON.stringify({
+        id: "test",
+        data: [
+          ["key1", 1],
+          ["key2", 2],
+        ],
+      });
+      localStorageMock.getItem.mockReturnValue(serialized);
+
+      const result = storage.getItem("test-key", initialValue);
+
+      expect(localStorageMock.getItem).toHaveBeenCalledWith("test-key");
+      expect(result.id).toBe("test");
+      expect(result.data.get("key1")).toBe(1);
+      expect(result.data.get("key2")).toBe(2);
+    });
+
+    it("should return initial value when localStorage returns null", () => {
+      localStorageMock.getItem.mockReturnValue(null);
+
+      const result = storage.getItem("test-key", initialValue);
+
+      expect(result).toBe(initialValue);
+    });
+
+    it("should return initial value when localStorage returns empty string", () => {
+      localStorageMock.getItem.mockReturnValue("");
+
+      const result = storage.getItem("test-key", initialValue);
+
+      expect(result).toBe(initialValue);
+    });
+
+    it("should return initial value and log warning when JSON parsing fails", () => {
+      localStorageMock.getItem.mockReturnValue("invalid-json");
+
+      const result = storage.getItem("test-key", initialValue);
+
+      expect(result).toBe(initialValue);
+      // Logger.warn should have been called but we're not testing the exact call
+    });
+
+    it("should return initial value and log warning when deserialization fails", () => {
+      // This JSON is parseable but will cause the Map constructor to throw
+      localStorageMock.getItem.mockReturnValue(
+        '{"id": "test", "data": "not-an-array"}',
+      );
+
+      const result = storage.getItem("test-key", initialValue);
+
+      expect(result).toEqual(initialValue);
+    });
+  });
+
+  describe("setItem", () => {
+    it("should serialize and store value in localStorage", () => {
+      storage.setItem("test-key", testValue);
+
+      expect(localStorageMock.setItem).toHaveBeenCalledWith(
+        "test-key",
+        JSON.stringify({
+          id: "test",
+          data: [
+            ["key1", 1],
+            ["key2", 2],
+          ],
+        }),
+      );
+    });
+  });
+
+  describe("removeItem", () => {
+    it("should call localStorage.removeItem with the correct key", () => {
+      storage.removeItem("test-key");
+
+      expect(localStorageMock.removeItem).toHaveBeenCalledWith("test-key");
+    });
+  });
+});

package/src/utils/storage.ts

@@ -0,0 +1,38 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import type { SyncStorage } from "jotai/vanilla/utils/atomWithStorage";
+import { Logger } from "./Logger";
+
+/**
+ * Converts a value to and from a serializable format for storage.
+ * Useful for data structures that are not natively supported by localStorage (e.g., Maps)
+ *
+ * @param opts - The options for the storage adapter.
+ * @returns - The storage adapter.
+ */
+export function adaptForLocalStorage<Value, Serializable>(opts: {
+  toSerializable: (v: Value) => Serializable;
+  fromSerializable: (s: Serializable) => Value;
+}): SyncStorage<Value> {
+  return {
+    getItem(key, initialValue) {
+      try {
+        const value = localStorage.getItem(key);
+        if (!value) {
+          return initialValue;
+        }
+        const parsed = JSON.parse(value) as Serializable;
+        return opts.fromSerializable(parsed);
+      } catch (error) {
+        Logger.warn(`Error getting ${key} from storage`, error);
+        return initialValue;
+      }
+    },
+    setItem: (key, value): void => {
+      localStorage.setItem(key, JSON.stringify(opts.toSerializable(value)));
+    },
+    removeItem: (key: string): void => {
+      localStorage.removeItem(key);
+    },
+  };
+}

package/dist/assets/_baseMap-CeNM_WhA.js

@@ -1 +0,0 @@
-import{b as m}from"./_baseEach-Ca2O9Wne.js";import{x as s}from"./index-BhSNFJlD.js";function e(r,o){var a=-1,t=s(r)?Array(r.length):[];return m(r,function(n,f,i){t[++a]=o(n,f,i)}),t}export{e as b};

package/dist/assets/channel-fQDJ54fB.js

@@ -1 +0,0 @@
-import{U as s,C as o}from"./mermaid-DsKZ0k6o.js";const n=(a,r)=>s.lang.round(o.parse(a)[r]);export{n as c};

package/dist/assets/classDiagram-3BZAVTQC-CFX5hd8N.js

@@ -1 +0,0 @@
-import{s as a,c as s,a as t,C as o}from"./chunk-JBRWN2VN-DMbINMTT.js";import{_ as e}from"./mermaid-DsKZ0k6o.js";import"./transform-B8bpuzxV.js";import"./chunk-GLLZNHP4-BqFTg_51.js";import"./chunk-WVR4S24B-B9Jid-bk.js";import"./chunk-NRVI72HA-CQ-LV0Dy.js";import"./index-BhSNFJlD.js";import"./step-BwsUM5iJ.js";import"./timer-BwIYMJWC.js";var i={parser:t,get db(){return new o},renderer:s,styles:a,init:e(r=>{r.class||(r.class={}),r.class.arrowMarkerAbsolute=r.arrowMarkerAbsolute},"init")};export{i as diagram};

package/dist/assets/classDiagram-v2-QTMF73CY-CFX5hd8N.js

@@ -1 +0,0 @@
-import{s as a,c as s,a as t,C as o}from"./chunk-JBRWN2VN-DMbINMTT.js";import{_ as e}from"./mermaid-DsKZ0k6o.js";import"./transform-B8bpuzxV.js";import"./chunk-GLLZNHP4-BqFTg_51.js";import"./chunk-WVR4S24B-B9Jid-bk.js";import"./chunk-NRVI72HA-CQ-LV0Dy.js";import"./index-BhSNFJlD.js";import"./step-BwsUM5iJ.js";import"./timer-BwIYMJWC.js";var i={parser:t,get db(){return new o},renderer:s,styles:a,init:e(r=>{r.class||(r.class={}),r.class.arrowMarkerAbsolute=r.arrowMarkerAbsolute},"init")};export{i as diagram};

package/dist/assets/clone-sKX7VVYQ.js

@@ -1 +0,0 @@
-import{b as n}from"./_baseUniq-D4iU71YJ.js";function o(r){return n(r,4)}export{o as c};

package/dist/assets/infoDiagram-6WOFNB3A-EUOwUSFP.js

@@ -1,2 +0,0 @@
-import{_ as a,l as i,I as s,d as p,K as n}from"./mermaid-DsKZ0k6o.js";import{p as m}from"./treemap-75Q7IDZK-BrZyfDQ2.js";import"./index-BhSNFJlD.js";import"./transform-B8bpuzxV.js";import"./timer-BwIYMJWC.js";import"./step-BwsUM5iJ.js";import"./_baseEach-Ca2O9Wne.js";import"./_baseUniq-D4iU71YJ.js";import"./min-CnCkyEM4.js";import"./_baseMap-CeNM_WhA.js";import"./clone-sKX7VVYQ.js";import"./_createAggregator-QKQ5gj2H.js";var d={parse:a(async r=>{const t=await m("info",r);i.debug(t)},"parse")},g={version:n.version+""},v={parser:d,db:{getVersion:a(()=>g.version,"getVersion")},renderer:{draw:a((r,t,o)=>{i.debug(`rendering info diagram
-`+r);const e=s(t);p(e,100,400,!0),e.append("g").append("text").attr("x",100).attr("y",40).attr("class","version").attr("font-size",32).style("text-anchor","middle").text(`v${o}`)},"draw")}};export{v as diagram};

package/dist/assets/stateDiagram-v2-EYPG3UTE-CQ0TxHAR.js

@@ -1 +0,0 @@
-import{s as t,b as a,a as e,S as o}from"./chunk-LXBSTHXV-C6VjboaW.js";import{_ as s}from"./mermaid-DsKZ0k6o.js";import"./transform-B8bpuzxV.js";import"./chunk-WVR4S24B-B9Jid-bk.js";import"./chunk-NRVI72HA-CQ-LV0Dy.js";import"./index-BhSNFJlD.js";import"./step-BwsUM5iJ.js";import"./timer-BwIYMJWC.js";var i={parser:e,get db(){return new o(2)},renderer:a,styles:t,init:s(r=>{r.state||(r.state={}),r.state.arrowMarkerAbsolute=r.arrowMarkerAbsolute},"init")};export{i as diagram};