improvement: fallback to available storage when localStorage is not available (#6879)

When iframing marimo in a sandboxed iframe, depending on configuration
`localStorage` may not be available so this adds some appropriate
fallbacks

this also adds a lint rule (grit) to make sure we pass the optional
storage into `atomWithStorage`

Author: mscolnick

biome.jsonc

@@ -3,7 +3,8 @@
   "plugins": [
     "frontend/lint/addEventListenerObject.grit",
     "frontend/lint/removeEventListenerObject.grit",
-    "frontend/lint/preferObjectParams.grit"
+    "frontend/lint/preferObjectParams.grit",
+    "frontend/lint/atomWithStorageArgs.grit"
   ],
   "files": {
     "includes": ["frontend/**/*", "packages/**/*", "!*.json", "!*.jsonc"]

frontend/lint/atomWithStorageArgs.grit

@@ -0,0 +1,12 @@
+// Enforce atomWithStorage to use at least 3 arguments
+// This ensures the storage parameter is always explicitly provided
+or {
+  `atomWithStorage($arg1, $arg2)`,
+  `atomWithStorage<$_>($arg1, $arg2)`
+} where {
+  register_diagnostic(
+    span=$arg1,
+    message="atomWithStorage requires at least 3 arguments (key, defaultValue, storage). Provide the storage parameter explicitly.",
+    severity="error"
+  )
+}

frontend/src/components/chat/acp/state.ts

@@ -4,6 +4,7 @@ import { atom } from "jotai";
 import { atomWithStorage } from "jotai/utils";
 import { capitalize } from "lodash-es";
 import { isPlatformWindows } from "@/core/hotkeys/shortcuts";
+import { jotaiJsonStorage } from "@/utils/storage/jotai";
 import type { TypedString } from "@/utils/typed";
 import { generateUUID } from "@/utils/uuid";
 import type { ExternalAgentSessionId, SessionSupportType } from "./types";
@@ -41,6 +42,7 @@ export const agentSessionStateAtom = atomWithStorage<AgentSessionState>(
     sessions: [],
     activeTabId: null,
   },
+  jotaiJsonStorage,
 );
 
 export const selectedTabAtom = atom(

frontend/src/components/data-table/charts/__tests__/storage.test.ts

@@ -2,8 +2,18 @@
 
 import { getDefaultStore } from "jotai";
 import { beforeEach, describe, expect, it, vi } from "vitest";
-import { asMock, SetupMocks } from "@/__mocks__/common";
+
+// Mock the storage module to use our mock storage
+vi.mock("@/utils/storage/storage", () => ({
+  availableStorage: {
+    getItem: vi.fn(),
+    setItem: vi.fn(),
+    removeItem: vi.fn(),
+  },
+}));
+
 import type { CellId } from "@/core/cells/ids";
+import { availableStorage } from "@/utils/storage/storage";
 import { ChartSchema } from "../schemas";
 import type { TabName } from "../storage";
 import { KEY, tabsStorageAtom } from "../storage";
@@ -12,7 +22,6 @@ import { ChartType } from "../types";
 describe("Chart Transforms Storage", () => {
   beforeEach(() => {
     vi.clearAllMocks();
-    SetupMocks.localStorage();
     // Reset the store before each test
     const store = getDefaultStore();
     store.set(tabsStorageAtom, new Map());
@@ -87,14 +96,11 @@ describe("Chart Transforms Storage", () => {
 
   describe("LocalStorage integration", () => {
     it("should call localStorage.setItem when setting data", () => {
-      // Reset the mock to ensure clean state
-      asMock(window.localStorage.setItem).mockReset();
-
       const store = getDefaultStore();
       const newMap = new Map();
       store.set(tabsStorageAtom, newMap);
 
-      expect(window.localStorage.setItem).toHaveBeenCalledWith(
+      expect(availableStorage.setItem).toHaveBeenCalledWith(
         KEY,
         JSON.stringify([]),
       );

frontend/src/components/data-table/charts/storage.ts

@@ -5,7 +5,7 @@ import { capitalize } from "lodash-es";
 import { z } from "zod";
 import type { CellId } from "@/core/cells/ids";
 import { Logger } from "@/utils/Logger";
-import { NotebookScopedLocalStorage } from "@/utils/localStorage";
+import { NotebookScopedLocalStorage } from "@/utils/storage/typed";
 import type { TypedString } from "@/utils/typed";
 import { ChartSchema, type ChartSchemaType } from "./schemas";
 import type { ChartType } from "./types";

frontend/src/components/editor/ai/add-cell-with-ai.tsx

@@ -53,7 +53,8 @@ import { useRuntimeManager } from "@/core/runtime/config";
 import { useTheme } from "@/theme/useTheme";
 import { cn } from "@/utils/cn";
 import { prettyError } from "@/utils/errors";
-import { ZodLocalStorage } from "@/utils/localStorage";
+import { jotaiJsonStorage } from "@/utils/storage/jotai";
+import { ZodLocalStorage } from "@/utils/storage/typed";
 import { PythonIcon } from "../cell/code/icons";
 import {
   CompletionActions,
@@ -66,6 +67,7 @@ import { StreamingChunkTransport } from "./transport/chat-transport";
 const languageAtom = atomWithStorage<"python" | "sql">(
   "marimo:ai-language",
   "python",
+  jotaiJsonStorage,
 );
 
 const KEY = "marimo:ai-prompt-history";

frontend/src/components/editor/chrome/state.ts

@@ -3,7 +3,7 @@
 import { useAtomValue } from "jotai";
 import { z } from "zod";
 import { createReducerAndAtoms } from "@/utils/createReducer";
-import { ZodLocalStorage } from "@/utils/localStorage";
+import { ZodLocalStorage } from "@/utils/storage/typed";
 import type { PanelType } from "./types";
 
 export interface ChromeState {

frontend/src/components/editor/columns/storage.ts

@@ -2,7 +2,7 @@
 
 import { z } from "zod";
 import { arrayMove } from "@/utils/arrays";
-import { NotebookScopedLocalStorage } from "@/utils/localStorage";
+import { NotebookScopedLocalStorage } from "@/utils/storage/typed";
 
 const BASE_KEY = "marimo:notebook-col-sizes";
 

frontend/src/components/editor/errors/fix-mode.ts

@@ -2,12 +2,17 @@
 
 import { useAtom } from "jotai";
 import { atomWithStorage } from "jotai/utils";
+import { jotaiJsonStorage } from "@/utils/storage/jotai";
 
 export type FixMode = "prompt" | "autofix";
 
 const BASE_KEY = "marimo:ai-autofix-mode";
 
-const fixModeAtom = atomWithStorage<FixMode>(BASE_KEY, "autofix");
+const fixModeAtom = atomWithStorage<FixMode>(
+  BASE_KEY,
+  "autofix",
+  jotaiJsonStorage,
+);
 
 export function useFixMode() {
   const [fixMode, setFixMode] = useAtom(fixModeAtom);

frontend/src/components/editor/output/useWrapText.ts

@@ -2,11 +2,16 @@
 
 import { useAtom } from "jotai";
 import { atomWithStorage } from "jotai/utils";
+import { jotaiJsonStorage } from "@/utils/storage/jotai";
 
 const WRAP_TEXT_KEY = "marimo:console:wrapText";
 
 // Atom for storing wrap text preference - shared across all console outputs
-const wrapTextAtom = atomWithStorage<boolean>(WRAP_TEXT_KEY, false);
+const wrapTextAtom = atomWithStorage<boolean>(
+  WRAP_TEXT_KEY,
+  false,
+  jotaiJsonStorage,
+);
 
 export function useWrapText() {
   const [wrapText, setWrapText] = useAtom(wrapTextAtom);