improvement: fallback to available storage when localStorage is not available

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);

frontend/src/components/home/state.ts

@@ -4,6 +4,7 @@ import { atomWithStorage } from "jotai/utils";
 import React from "react";
 import type { MarimoFile } from "@/core/network/types";
 import { Functions } from "@/utils/functions";
+import { jotaiJsonStorage } from "@/utils/storage/jotai";
 
 export type RunningNotebooksMap = Map<string, MarimoFile>;
 
@@ -19,12 +20,12 @@ export const WorkspaceRootContext = React.createContext<string>("");
 export const includeMarkdownAtom = atomWithStorage<boolean>(
   "marimo:home:include-markdown",
   false,
-  undefined,
+  jotaiJsonStorage,
   { getOnInit: true },
 );
 export const expandedFoldersAtom = atomWithStorage<Record<string, boolean>>(
   "marimo:home:expanded-folders",
   {},
-  undefined,
+  jotaiJsonStorage,
   { getOnInit: true },
 );

frontend/src/components/scratchpad/scratchpad-history.ts

@@ -1,12 +1,17 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 import { atom } from "jotai";
 import { atomWithStorage } from "jotai/utils";
+import { jotaiJsonStorage } from "@/utils/storage/jotai";
 
 const MAX_HISTORY_ITEMS = 15;
 const KEY = "marimo:scratchpadHistory:v1";
 
 // Atom for storing the history
-export const scratchpadHistoryAtom = atomWithStorage<string[]>(KEY, []);
+export const scratchpadHistoryAtom = atomWithStorage<string[]>(
+  KEY,
+  [],
+  jotaiJsonStorage,
+);
 
 // Atom for controlling the visibility of the history list
 export const historyVisibleAtom = atom(false);

frontend/src/core/ai/state.ts

@@ -5,7 +5,7 @@ import type { FileUIPart } from "ai";
 import { atom } from "jotai";
 import { atomWithStorage } from "jotai/utils";
 import { uniqueBy } from "@/utils/arrays";
-import { adaptForLocalStorage } from "@/utils/storage";
+import { adaptForLocalStorage, jotaiJsonStorage } from "@/utils/storage/jotai";
 import type { TypedString } from "@/utils/typed";
 import type { CellId } from "../cells/ids";
 
@@ -25,6 +25,7 @@ const INCLUDE_OTHER_CELLS_KEY = "marimo:ai:includeOtherCells";
 export const includeOtherCellsAtom = atomWithStorage<boolean>(
   INCLUDE_OTHER_CELLS_KEY,
   true,
+  jotaiJsonStorage,
 );
 
 export interface Message {

frontend/src/core/codemirror/copilot/state.ts

@@ -7,13 +7,15 @@ import {
   resolvedMarimoConfigAtom,
 } from "@/core/config/config";
 import { store, waitFor } from "@/core/state/jotai";
+import { jotaiJsonStorage } from "@/utils/storage/jotai";
+import { availableStorage } from "@/utils/storage/storage";
 
 const KEY = "marimo:copilot:signedIn";
 
 export const isGitHubCopilotSignedInState = atomWithStorage<boolean | null>(
   KEY,
   null,
-  undefined,
+  jotaiJsonStorage,
   {
     getOnInit: true,
   },
@@ -49,7 +51,7 @@ export function clearGitHubCopilotLoadingVersion(expectedVersion: number) {
 }
 
 function getIsLastSignedIn() {
-  const lastSignedIn = localStorage.getItem(KEY);
+  const lastSignedIn = availableStorage.getItem(KEY);
   return lastSignedIn === "true";
 }
 

frontend/src/core/codemirror/language/languages/sql/sql-mode.ts

@@ -3,12 +3,17 @@
 import { useAtom } from "jotai";
 import { atomWithStorage } from "jotai/utils";
 import { store } from "@/core/state/jotai";
+import { jotaiJsonStorage } from "@/utils/storage/jotai";
 
 const BASE_KEY = "marimo:notebook-sql-mode";
 
 export type SQLMode = "validate" | "default";
 
-const sqlModeAtom = atomWithStorage<SQLMode>(BASE_KEY, "default");
+const sqlModeAtom = atomWithStorage<SQLMode>(
+  BASE_KEY,
+  "default",
+  jotaiJsonStorage,
+);
 
 export function useSQLMode() {
   const [sqlMode, setSQLMode] = useAtom(sqlModeAtom);

frontend/src/core/rtc/state.ts

@@ -1,13 +1,18 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 import { atomWithStorage } from "jotai/utils";
 import { once } from "lodash-es";
+import { jotaiJsonStorage } from "@/utils/storage/jotai";
 import { getFeatureFlag } from "../config/feature-flag";
 
 /**
  * The username for the current user when using real-time collaboration.
  * This is stored in localStorage.
  */
-export const usernameAtom = atomWithStorage<string>("marimo:rtc:username", "");
+export const usernameAtom = atomWithStorage<string>(
+  "marimo:rtc:username",
+  "",
+  jotaiJsonStorage,
+);
 
 /**
  * Whether RTC is enabled.

frontend/src/core/wasm/store.ts

@@ -4,7 +4,7 @@ import {
   compressToEncodedURIComponent,
   decompressFromEncodedURIComponent,
 } from "lz-string";
-import { TypedLocalStorage } from "@/utils/localStorage";
+import { TypedLocalStorage } from "@/utils/storage/typed";
 import { PyodideRouter } from "./router";
 
 export interface FileStore {

frontend/src/hooks/useLocalStorage.ts

@@ -1,7 +1,7 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
 import { useState } from "react";
-import { TypedLocalStorage } from "@/utils/localStorage";
+import { TypedLocalStorage } from "@/utils/storage/typed";
 
 /**
  * React hook to use localStorage

frontend/src/utils/__tests__/local-storage.test.ts

@@ -5,7 +5,7 @@ import { z } from "zod";
 import { TestUtils } from "@/__tests__/test-helpers";
 import { filenameAtom } from "@/core/saving/file-state";
 import { store } from "@/core/state/jotai";
-import { NotebookScopedLocalStorage } from "../localStorage";
+import { NotebookScopedLocalStorage } from "../storage/typed";
 
 describe("NotebookScopedLocalStorage", () => {
   const schema = z.object({