@marimo-team/frontend 0.20.2 → 0.20.3-dev77

package/src/core/storage/state.ts

@@ -0,0 +1,132 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { atom, useAtomValue } from "jotai";
+import { useAsyncData } from "@/hooks/useAsyncData";
+import { createReducerAndAtoms } from "@/utils/createReducer";
+import type { NotificationMessageData } from "../kernel/messages";
+import type { VariableName } from "../variables/types";
+import { ListStorageEntries } from "./request-registry";
+import type { StorageEntry, StorageState } from "./types";
+import {
+  DEFAULT_FETCH_LIMIT,
+  ROOT_PATH,
+  storageNamespacePrefix,
+  storagePathKey,
+} from "./types";
+
+function initialState(): StorageState {
+  return {
+    namespaces: [],
+    entriesByPath: new Map(),
+  };
+}
+
+const {
+  reducer,
+  createActions,
+  valueAtom: storageAtom,
+  useActions,
+} = createReducerAndAtoms(initialState, {
+  setNamespaces: (
+    state,
+    data: NotificationMessageData<"storage-namespaces">,
+  ) => {
+    // Merge/replace namespaces by name
+    const existingMap = new Map(state.namespaces.map((ns) => [ns.name, ns]));
+    for (const ns of data.namespaces) {
+      existingMap.set(ns.name, ns);
+    }
+    return {
+      ...state,
+      namespaces: [...existingMap.values()],
+    };
+  },
+
+  setEntries: (
+    state,
+    opts: {
+      namespace: string;
+      prefix: string | null | undefined;
+      entries: StorageEntry[];
+    },
+  ) => {
+    const key = storagePathKey(opts.namespace, opts.prefix);
+    const entriesByPath = new Map(state.entriesByPath);
+    entriesByPath.set(key, opts.entries);
+    return { ...state, entriesByPath };
+  },
+
+  clearNamespaceCache: (state, namespace: string) => {
+    const entriesByPath = new Map(state.entriesByPath);
+    const prefix = storageNamespacePrefix(namespace);
+    for (const key of entriesByPath.keys()) {
+      if (key.startsWith(prefix)) {
+        entriesByPath.delete(key);
+      }
+    }
+    return { ...state, entriesByPath };
+  },
+
+  filterFromVariables: (state, variableNames: VariableName[]) => {
+    const names = new Set(variableNames);
+    // Filter out namespaces whose backing variable is no longer in scope
+    const namespaces = state.namespaces.filter((ns) => {
+      return names.has(ns.name as VariableName);
+    });
+    return { ...state, namespaces };
+  },
+});
+
+/**
+ * React hook to get the storage state.
+ */
+export const useStorage = () => useAtomValue(storageAtom);
+
+export const storageNamespacesAtom = atom((get) => get(storageAtom).namespaces);
+
+/**
+ * React hook to get the storage actions.
+ */
+export function useStorageActions() {
+  return useActions();
+}
+
+export { storageAtom };
+
+/**
+ * Hook that fetches and caches storage entries for a given namespace/prefix.
+ * Entries are fetched on first access and cached in the store for subsequent renders.
+ */
+export function useStorageEntries(namespace: string, prefix?: string) {
+  const { entriesByPath } = useStorage();
+  const { setEntries } = useStorageActions();
+  const cached = entriesByPath.get(storagePathKey(namespace, prefix));
+
+  const { isPending, error, refetch } = useAsyncData(async () => {
+    if (cached) {
+      return;
+    }
+    const result = await ListStorageEntries.request({
+      namespace,
+      prefix: prefix ?? ROOT_PATH,
+      limit: DEFAULT_FETCH_LIMIT,
+    });
+    if (result.error) {
+      throw new Error(result.error);
+    }
+    setEntries({ namespace, prefix, entries: result.entries });
+  }, [namespace, prefix, cached === undefined]);
+
+  return {
+    entries: cached ?? [],
+    isPending: isPending && !cached,
+    error: cached ? undefined : error,
+    refetch,
+  };
+}
+
+export const exportedForTesting = {
+  reducer,
+  createActions,
+  initialState,
+};

package/src/core/storage/types.ts

@@ -0,0 +1,57 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+import type { components } from "@marimo-team/marimo-api";
+
+export type StorageEntry = components["schemas"]["StorageEntry"];
+export type StorageNamespace = components["schemas"]["StorageNamespace"];
+
+/**
+ * Key for looking up lazy-loaded entries: "namespace::prefix"
+ */
+export type StoragePathKey = `${string}::${string}`;
+
+export const STORAGE_PATH_SEPARATOR = "::";
+
+export function storagePathKey(
+  namespace: string,
+  prefix: string | null | undefined,
+): StoragePathKey {
+  return `${namespace}${STORAGE_PATH_SEPARATOR}${prefix ?? ""}`;
+}
+
+export function storageNamespacePrefix(namespace: string): string {
+  return `${namespace}${STORAGE_PATH_SEPARATOR}`;
+}
+
+/**
+ * The storage state.
+ */
+export interface StorageState {
+  /** Namespaces from the broadcast notification */
+  namespaces: StorageNamespace[];
+  /** Lazy-loaded entries keyed by "namespace::prefix" */
+  entriesByPath: ReadonlyMap<StoragePathKey, StorageEntry[]>;
+}
+
+export function storageUrl(
+  protocol: string,
+  rootPath: string,
+  entryPath: string,
+): URL {
+  const parts = [rootPath, entryPath].filter(Boolean);
+  const path = parts.join("/").replaceAll(/\/+/g, "/");
+  return new URL(`${protocol}://${path}`);
+}
+
+export const ROOT_PATH = "";
+export const DEFAULT_FETCH_LIMIT = 150;
+
+/** Non-exhaustive list of known storage protocols */
+export type KnownStorageProtocol =
+  | "s3"
+  | "gcs"
+  | "azure"
+  | "coreweave"
+  | "cloudflare"
+  | "http"
+  | "file"
+  | "in-memory";

package/src/core/wasm/bridge.ts

@@ -596,6 +596,8 @@ export class PyodideBridge implements RunRequests, EditRequests {
   invokeAiTool = throwNotImplemented;
   clearCache = throwNotImplemented;
   getCacheInfo = throwNotImplemented;
+  listStorageEntries = throwNotImplemented;
+  downloadStorage = throwNotImplemented;
 
   private async putControlRequest(operation: CommandMessage) {
     await this.rpc.proxy.request.bridge({

package/src/core/websocket/useMarimoKernelConnection.tsx

@@ -56,6 +56,11 @@ import type { RequestId } from "../network/DeferredRequestRegistry";
 import { useRuntimeManager } from "../runtime/config";
 import { SECRETS_REGISTRY } from "../secrets/request-registry";
 import { isStaticNotebook } from "../static/static-state";
+import {
+  DownloadStorage,
+  ListStorageEntries,
+} from "../storage/request-registry";
+import { useStorageActions } from "../storage/state";
 import { useVariablesActions } from "../variables/state";
 import type { VariableName } from "../variables/types";
 import { isWasm } from "../wasm/utils";
@@ -105,6 +110,10 @@ export function useMarimoKernelConnection(opts: {
   const runtimeManager = useRuntimeManager();
   const setCacheInfo = useSetAtom(cacheInfoAtom);
   const setKernelStartupError = useSetAtom(kernelStartupErrorAtom);
+  const {
+    setNamespaces: setStorageNamespaces,
+    filterFromVariables: filterStorageFromVariables,
+  } = useStorageActions();
 
   const handleMessage = (e: MessageEvent<JsonString<NotificationPayload>>) => {
     const msg = jsonParseWithSpecialChar(e.data);
@@ -196,6 +205,9 @@ export function useMarimoKernelConnection(opts: {
         filterDataSourcesFromVariables(
           msg.data.variables.map((v) => v.name as VariableName),
         );
+        filterStorageFromVariables(
+          msg.data.variables.map((v) => v.name as VariableName),
+        );
         return;
       case "variable-values":
         setMetadata(
@@ -285,7 +297,13 @@ export function useMarimoKernelConnection(opts: {
         });
         return;
       case "storage-namespaces":
-        // TODO: Handle storage namespaces
+        setStorageNamespaces(msg.data);
+        return;
+      case "storage-entries":
+        ListStorageEntries.resolve(msg.data.request_id as RequestId, msg.data);
+        return;
+      case "storage-download-ready":
+        DownloadStorage.resolve(msg.data.request_id as RequestId, msg.data);
         return;
 
       case "reconnected":

package/src/hooks/__tests__/useNonce.test.tsx

@@ -0,0 +1,127 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { act, renderHook } from "@testing-library/react";
+import { describe, expect, it } from "vitest";
+import { useNonce } from "../useNonce";
+
+describe("useNonce", () => {
+  it("should return a stable function reference", () => {
+    const { result, rerender } = renderHook(() => useNonce());
+    const firstRender = result.current;
+
+    rerender();
+
+    expect(result.current).toBe(firstRender);
+  });
+
+  it("should trigger re-render when called", () => {
+    let renderCount = 0;
+    const { result } = renderHook(() => {
+      renderCount++;
+      return useNonce();
+    });
+
+    const initialRenderCount = renderCount;
+
+    act(() => {
+      result.current();
+    });
+
+    expect(renderCount).toBe(initialRenderCount + 1);
+  });
+
+  it("should handle multiple rapid calls within the same act", () => {
+    // This is the core test for the stale closure fix.
+    // Before the fix, calling the function multiple times in the same
+    // synchronous block would always set nonce to the same value (e.g. 0+1=1),
+    // because the closure captured a stale `nonce`. With the functional update
+    // form `setNonce(n => n + 1)`, each call correctly increments from the
+    // latest pending state.
+    let renderCount = 0;
+    const { result } = renderHook(() => {
+      renderCount++;
+      return useNonce();
+    });
+
+    const initialRenderCount = renderCount;
+
+    // Call 3 times within the same act — only one batched re-render
+    act(() => {
+      result.current();
+      result.current();
+      result.current();
+    });
+
+    // React batches these into a single re-render, but the nonce should
+    // have incremented 3 times (0 → 1 → 2 → 3) thanks to functional updates.
+    // We verify a re-render happened (at least 1 extra render).
+    expect(renderCount).toBeGreaterThan(initialRenderCount);
+  });
+
+  it("should trigger separate re-renders for separate act calls", () => {
+    let renderCount = 0;
+    const { result } = renderHook(() => {
+      renderCount++;
+      return useNonce();
+    });
+
+    const initialRenderCount = renderCount;
+
+    act(() => {
+      result.current();
+    });
+
+    act(() => {
+      result.current();
+    });
+
+    act(() => {
+      result.current();
+    });
+
+    expect(renderCount).toBe(initialRenderCount + 3);
+  });
+
+  it("should maintain stable reference across multiple re-renders", () => {
+    const { result, rerender } = renderHook(() => useNonce());
+    const increment = result.current;
+
+    act(() => {
+      result.current();
+    });
+
+    rerender();
+
+    act(() => {
+      result.current();
+    });
+
+    rerender();
+
+    // Function reference should remain the same throughout
+    expect(result.current).toBe(increment);
+  });
+
+  it("should work correctly when called from event handlers", () => {
+    let renderCount = 0;
+    const { result } = renderHook(() => {
+      renderCount++;
+      return useNonce();
+    });
+
+    const initialRenderCount = renderCount;
+    const increment = result.current;
+
+    act(() => {
+      increment();
+    });
+
+    expect(renderCount).toBe(initialRenderCount + 1);
+
+    act(() => {
+      increment();
+    });
+
+    expect(renderCount).toBe(initialRenderCount + 2);
+  });
+});

package/src/hooks/__tests__/useOverflowDetection.test.tsx

@@ -0,0 +1,127 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { renderHook } from "@testing-library/react";
+import { useRef } from "react";
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import { useOverflowDetection } from "../useOverflowDetection";
+
+// Mock ResizeObserver
+const mockObserve = vi.fn();
+const mockDisconnect = vi.fn();
+
+global.ResizeObserver = class MockResizeObserver {
+  callback: ResizeObserverCallback;
+
+  constructor(callback: ResizeObserverCallback) {
+    this.callback = callback;
+  }
+
+  observe(target: Element) {
+    mockObserve(target);
+    // Simulate initial observation callback
+    this.callback([], this);
+  }
+
+  unobserve() {
+    // noop
+  }
+
+  disconnect() {
+    mockDisconnect();
+  }
+};
+
+describe("useOverflowDetection", () => {
+  beforeEach(() => {
+    mockObserve.mockClear();
+    mockDisconnect.mockClear();
+  });
+
+  it("should return false when ref is null", () => {
+    const { result } = renderHook(() => {
+      const ref = useRef<HTMLDivElement>(null);
+      return useOverflowDetection(ref);
+    });
+
+    expect(result.current).toBe(false);
+    expect(mockObserve).not.toHaveBeenCalled();
+  });
+
+  it("should observe element when ref is set", () => {
+    const element = document.createElement("div");
+
+    const { result } = renderHook(() => {
+      const ref = useRef<HTMLDivElement>(element);
+      return useOverflowDetection(ref);
+    });
+
+    expect(mockObserve).toHaveBeenCalledWith(element);
+    expect(result.current).toBe(false); // No overflow by default
+  });
+
+  it("should not observe when enabled is false", () => {
+    const element = document.createElement("div");
+
+    renderHook(() => {
+      const ref = useRef<HTMLDivElement>(element);
+      return useOverflowDetection(ref, false);
+    });
+
+    expect(mockObserve).not.toHaveBeenCalled();
+  });
+
+  it("should disconnect observer on unmount", () => {
+    const element = document.createElement("div");
+
+    const { unmount } = renderHook(() => {
+      const ref = useRef<HTMLDivElement>(element);
+      return useOverflowDetection(ref);
+    });
+
+    unmount();
+    expect(mockDisconnect).toHaveBeenCalled();
+  });
+
+  it("should re-observe when enabled changes from false to true", () => {
+    const element = document.createElement("div");
+
+    const { rerender } = renderHook(
+      ({ enabled }) => {
+        const ref = useRef<HTMLDivElement>(element);
+        return useOverflowDetection(ref, enabled);
+      },
+      { initialProps: { enabled: false } },
+    );
+
+    expect(mockObserve).not.toHaveBeenCalled();
+
+    rerender({ enabled: true });
+    expect(mockObserve).toHaveBeenCalledWith(element);
+  });
+
+  it("should detect overflow when scrollHeight > clientHeight", () => {
+    const element = document.createElement("div");
+    Object.defineProperty(element, "scrollHeight", { value: 500 });
+    Object.defineProperty(element, "clientHeight", { value: 200 });
+
+    const { result } = renderHook(() => {
+      const ref = useRef<HTMLDivElement>(element);
+      return useOverflowDetection(ref);
+    });
+
+    expect(result.current).toBe(true);
+  });
+
+  it("should not detect overflow when scrollHeight <= clientHeight", () => {
+    const element = document.createElement("div");
+    Object.defineProperty(element, "scrollHeight", { value: 100 });
+    Object.defineProperty(element, "clientHeight", { value: 200 });
+
+    const { result } = renderHook(() => {
+      const ref = useRef<HTMLDivElement>(element);
+      return useOverflowDetection(ref);
+    });
+
+    expect(result.current).toBe(false);
+  });
+});

package/src/hooks/useNonce.ts

@@ -1,9 +1,10 @@
 /* Copyright 2026 Marimo. All rights reserved. */
-import { useState } from "react";
+import { useCallback, useState } from "react";
 
 export function useNonce() {
-  const [nonce, setNonce] = useState(0);
-  return () => {
-    setNonce(nonce + 1);
-  };
+  // eslint-disable-next-line react/hook-use-state
+  const [, setNonce] = useState(0);
+  return useCallback(() => {
+    setNonce((n) => n + 1);
+  }, []);
 }

package/src/hooks/useOverflowDetection.ts

@@ -0,0 +1,36 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import type { RefObject } from "react";
+import { useEffect, useState } from "react";
+
+/**
+ * Hook to detect if an element's content overflows its visible area.
+ * Uses ResizeObserver to track size changes.
+ *
+ * @param ref - Ref to the element to observe
+ * @param enabled - When this changes, re-setup the observer (use to handle element recreation)
+ */
+export function useOverflowDetection(
+  ref: RefObject<HTMLElement | null>,
+  enabled = true,
+): boolean {
+  const [isOverflowing, setIsOverflowing] = useState(false);
+
+  useEffect(() => {
+    const el = ref.current;
+    if (!el || !enabled) {
+      return;
+    }
+
+    const resizeObserver = new ResizeObserver(() => {
+      setIsOverflowing(el.scrollHeight > el.clientHeight);
+    });
+    resizeObserver.observe(el);
+
+    return () => {
+      resizeObserver.disconnect();
+    };
+  }, [ref, enabled]);
+
+  return isOverflowing;
+}

package/src/plugins/impl/MatrixPlugin.tsx

@@ -1,5 +1,5 @@
 /* Copyright 2026 Marimo. All rights reserved. */
-import { type JSX, useCallback, useRef, useState } from "react";
+import { type JSX, useCallback, useEffect, useRef, useState } from "react";
 import { z } from "zod";
 import { cn } from "@/utils/cn";
 import type { IPlugin, IPluginProps, Setter } from "../types";
@@ -90,7 +90,10 @@ const MatrixComponent = ({
   // Outside of a drag we always read from the prop `value` directly,
   // which avoids stale-state bugs when the matrix shape changes.
   const [draft, setDraft] = useState(value);
-  const displayValue = activeCell != null ? draft : value;
+  useEffect(() => {
+    setDraft(value);
+  }, [value]);
+  const displayValue = activeCell == null ? value : draft;
 
   const formatValue = (val: number) =>
     scientific ? val.toExponential(precision) : val.toFixed(precision);

package/src/plugins/impl/__tests__/MatrixPlugin.test.tsx

@@ -206,7 +206,7 @@ describe("MatrixPlugin", () => {
     const plugin = new MatrixPlugin();
     const props = makeProps({
       value: [
-        [0.00153, 1234567],
+        [0.001_53, 1_234_567],
         [0, -0.042],
       ],
       data: {

package/src/plugins/impl/code/any-language-editor.tsx

@@ -17,10 +17,14 @@ import { ErrorBanner } from "../common/error-banner";
 
 export const LANGUAGE_MAP: Record<string, LanguageName | undefined> = {
   python: "py",
+  python3: "py",
   javascript: "js",
   typescript: "ts",
   shell: "sh",
   bash: "sh",
+  // Other fallbacks
+  unknown: "text",
+  undefined: "text",
 };
 
 function isSupportedLanguage(
@@ -41,10 +45,17 @@ function isSupportedLanguage(
 const AnyLanguageCodeMirror: React.FC<
   ReactCodeMirrorProps & {
     language: string | undefined;
+    hideUnsupportedLanguageErrors?: boolean;
     theme: ResolvedTheme;
     showCopyButton?: boolean;
   }
-> = ({ language, showCopyButton, extensions = [], ...props }) => {
+> = ({
+  language,
+  hideUnsupportedLanguageErrors,
+  showCopyButton,
+  extensions = [],
+  ...props
+}) => {
   // Maybe normalize the language to the extension
   language = LANGUAGE_MAP[language || ""] || language;
 
@@ -62,10 +73,10 @@ const AnyLanguageCodeMirror: React.FC<
 
   return (
     <div className="relative w-full group hover-actions-parent">
-      {!isSupported && (
+      {!isSupported && !hideUnsupportedLanguageErrors && (
         <ErrorBanner
           className="mb-1 rounded-sm"
-          error={`Language ${language} not supported. \n\nSupported languages are: ${Object.keys(
+          error={`Language ${language} not supported. Supported languages are: ${Object.keys(
             langs,
           ).join(", ")}`}
         />