@marimo-team/islands 0.25.3-dev2 → 0.25.3-dev200

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

@@ -0,0 +1,257 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+import {
+  act,
+  fireEvent,
+  render,
+  renderHook,
+  screen,
+} from "@testing-library/react";
+import {
+  Profiler,
+  startTransition,
+  StrictMode,
+  Suspense,
+  useLayoutEffect,
+  useState,
+} from "react";
+import { describe, expect, it, vi } from "vitest";
+import { Deferred } from "@/utils/Deferred";
+import { useStateForKey } from "../useStateForKey";
+
+describe("useStateForKey", () => {
+  it("preserves edits for the same key and creates fresh state on key changes", () => {
+    const initialize = vi.fn(() => new Set<string>());
+    const { result, rerender } = renderHook(
+      ({ key }) => useStateForKey(key, initialize),
+      { initialProps: { key: "first" } },
+    );
+    const initial = result.current[0];
+    act(() => result.current[1](new Set(["expanded"])));
+    rerender({ key: "first" });
+    expect(result.current[0]).toEqual(new Set(["expanded"]));
+    expect(initialize).toHaveBeenCalledOnce();
+
+    rerender({ key: "second" });
+    expect(result.current[0]).toEqual(new Set());
+    expect(result.current[0]).not.toBe(initial);
+    expect(initialize).toHaveBeenCalledTimes(2);
+
+    rerender({ key: "first" });
+    expect(result.current[0]).toEqual(new Set());
+    expect(initialize).toHaveBeenCalledTimes(3);
+  });
+
+  it("preserves queued functional updates and ignores setters from another key", () => {
+    const { result, rerender } = renderHook(
+      ({ key }) => useStateForKey(key, () => 0),
+      { initialProps: { key: "first" } },
+    );
+    const oldSetter = result.current[1];
+    act(() => {
+      result.current[1]((value) => value + 1);
+      result.current[1]((value) => value + 1);
+    });
+    expect(result.current[0]).toBe(2);
+
+    rerender({ key: "second" });
+    const staleUpdate = vi.fn(() => 99);
+    act(() => oldSetter(staleUpdate));
+    act(() => oldSetter(99));
+    expect(result.current[0]).toBe(0);
+    expect(staleUpdate).not.toHaveBeenCalled();
+  });
+
+  it("uses Object.is for NaN and signed zero keys", () => {
+    const { result, rerender } = renderHook(
+      ({ key }) => useStateForKey(key, () => "initial"),
+      { initialProps: { key: Number.NaN } },
+    );
+    act(() => result.current[1]("draft"));
+    rerender({ key: Number.NaN });
+    expect(result.current[0]).toBe("draft");
+
+    rerender({ key: 0 });
+    expect(result.current[0]).toBe("initial");
+    act(() => result.current[1]("draft"));
+    rerender({ key: -0 });
+    expect(result.current[0]).toBe("initial");
+  });
+
+  it("stores function values without invoking them", () => {
+    const initial = vi.fn();
+    const changed = vi.fn();
+    const { result, rerender } = renderHook(
+      ({ key, value }) => useStateForKey(key, () => value),
+      { initialProps: { key: "first", value: initial } },
+    );
+    expect(result.current[0]).toBe(initial);
+    act(() => result.current[1](() => changed));
+    expect(result.current[0]).toBe(changed);
+    rerender({ key: "second", value: initial });
+    expect(result.current[0]).toBe(initial);
+    expect(initial).not.toHaveBeenCalled();
+    expect(changed).not.toHaveBeenCalled();
+  });
+
+  it("uses a changed initializer only after the key changes", () => {
+    const first = vi.fn(() => "first default");
+    const second = vi.fn(() => "second default");
+    const { result, rerender } = renderHook(
+      ({ key, initialize }) => useStateForKey(key, initialize),
+      { initialProps: { key: "A", initialize: first } },
+    );
+    act(() => result.current[1]("draft"));
+    rerender({ key: "A", initialize: second });
+    expect(result.current[0]).toBe("draft");
+    expect(second).not.toHaveBeenCalled();
+    rerender({ key: "B", initialize: second });
+    expect(result.current[0]).toBe("second default");
+    expect(second).toHaveBeenCalledOnce();
+  });
+
+  it("compares object keys by reference rather than their contents", () => {
+    const key = { id: "A" };
+    const { result, rerender } = renderHook(
+      ({ key }) => useStateForKey(key, () => 0),
+      { initialProps: { key } },
+    );
+    act(() => result.current[1](1));
+    rerender({ key });
+    expect(result.current[0]).toBe(1);
+    rerender({ key: { id: "A" } });
+    expect(result.current[0]).toBe(0);
+  });
+
+  it("keeps setters stable for the same key and accepts old setters when that key returns", () => {
+    const { result, rerender } = renderHook(
+      ({ key }) => useStateForKey(key, () => 0),
+      { initialProps: { key: "A" } },
+    );
+    const setA = result.current[1];
+    act(() => setA(1));
+    rerender({ key: "A" });
+    expect(result.current[1]).toBe(setA);
+    rerender({ key: "B" });
+    expect(result.current[1]).not.toBe(setA);
+    act(() => setA(2));
+    expect(result.current[0]).toBe(0);
+    rerender({ key: "A" });
+    act(() => setA(3));
+    expect(result.current[0]).toBe(3);
+  });
+
+  it("retries the owner without rendering or committing stale state to children", () => {
+    const onOwnerRender = vi.fn();
+    const onChildRender = vi.fn();
+    const onCommit = vi.fn();
+    const onLayout = vi.fn();
+    function Child({ owner, value }: { owner: string; value: string }) {
+      onChildRender(owner, value);
+      useLayoutEffect(() => onLayout(owner, value), [owner, value]);
+      return <span>{value}</span>;
+    }
+    function Owner({ owner }: { owner: string }) {
+      onOwnerRender();
+      const [value, setValue] = useStateForKey(owner, () => `${owner} default`);
+      return (
+        <>
+          <button onClick={() => setValue("draft")}>Edit</button>
+          <Child owner={owner} value={value} />
+        </>
+      );
+    }
+    const content = (owner: string) => (
+      <Profiler id="owner" onRender={onCommit}>
+        <Owner owner={owner} />
+      </Profiler>
+    );
+    const { rerender } = render(content("A"));
+    const button = screen.getByRole("button", { name: "Edit" });
+    fireEvent.click(button);
+    expect(screen.getByText("draft")).toBeInTheDocument();
+    onOwnerRender.mockClear();
+    onChildRender.mockClear();
+    onCommit.mockClear();
+    onLayout.mockClear();
+
+    rerender(content("B"));
+    expect(onOwnerRender.mock.calls.length).toBeGreaterThan(1);
+    expect(onChildRender.mock.calls).toEqual([["B", "B default"]]);
+    expect(onLayout.mock.calls).toEqual([["B", "B default"]]);
+    expect(onCommit).toHaveBeenCalledOnce();
+    expect(screen.getByRole("button", { name: "Edit" })).toBe(button);
+  });
+
+  it("keeps state correct when Strict Mode repeats initialization and updates", () => {
+    const initialize = vi.fn(() => new Set<string>());
+    const { result, rerender } = renderHook(
+      ({ key }) => useStateForKey(key, initialize),
+      { initialProps: { key: "A" }, wrapper: StrictMode },
+    );
+    expect(initialize.mock.calls.length).toBeGreaterThan(1);
+    act(() => result.current[1]((previous) => new Set([...previous, "draft"])));
+    expect(result.current[0]).toEqual(new Set(["draft"]));
+    rerender({ key: "B" });
+    expect(result.current[0]).toEqual(new Set());
+  });
+
+  it.each(["completes", "is canceled"])(
+    "preserves visible edits when a suspended transition %s",
+    async (outcome) => {
+      const pending = new Deferred<void>();
+      let blocked = true;
+      const suspended = vi.fn();
+      function Pending({ owner }: { owner: string }) {
+        if (owner === "B" && blocked) {
+          suspended();
+          throw pending.promise;
+        }
+        return null;
+      }
+      function Owner({ owner }: { owner: string }) {
+        const [value, setValue] = useStateForKey(owner, () => 0);
+        return (
+          <>
+            <button onClick={() => setValue((previous) => previous + 1)}>
+              {owner}:{value}
+            </button>
+            <Pending owner={owner} />
+          </>
+        );
+      }
+      function Harness() {
+        const [owner, setOwner] = useState("A");
+        return (
+          <>
+            <button onClick={() => startTransition(() => setOwner("B"))}>
+              Switch owner
+            </button>
+            <button onClick={() => setOwner("A")}>Keep owner</button>
+            <Suspense fallback="Loading">
+              <Owner owner={owner} />
+            </Suspense>
+          </>
+        );
+      }
+      render(<Harness />);
+      fireEvent.click(screen.getByRole("button", { name: "A:0" }));
+      fireEvent.click(screen.getByRole("button", { name: "Switch owner" }));
+      expect(suspended).toHaveBeenCalled();
+      fireEvent.click(screen.getByRole("button", { name: "A:1" }));
+      expect(screen.getByRole("button", { name: "A:2" })).toBeInTheDocument();
+      if (outcome === "is canceled") {
+        fireEvent.click(screen.getByRole("button", { name: "Keep owner" }));
+      }
+      await act(async () => {
+        blocked = false;
+        pending.resolve();
+        await pending.promise;
+      });
+      expect(
+        screen.getByRole("button", {
+          name: outcome === "completes" ? "B:0" : "A:2",
+        }),
+      ).toBeInTheDocument();
+    },
+  );
+});

package/src/hooks/useAsyncData.ts

@@ -322,7 +322,7 @@ export function useAsyncData<T>(
   const [nonce, setNonce] = useState(0);
   const [result, setResult] = useState<
     PendingResult<T> | LoadingResult<T> | ErrorResult<T> | SuccessResult<T>
-  >(Result.pending());
+  >(() => Result.pending());
 
   const asProps =
     typeof loaderOrProps === "function"
@@ -339,12 +339,12 @@ export function useAsyncData<T>(
         keepPrevious = true;
       },
     };
+    // Activity can restart this effect without changing dependencies, so
+    // loading state must follow the request lifecycle.
     setResult((prevResult) => {
-      // If we have previous data, show reloading state
       if (prevResult.status === "success" || prevResult.status === "loading") {
         return Result.loading(prevResult.data);
       }
-      // Otherwise, show initial loading state
       return Result.pending();
     });
     fetchStable(context)

package/src/hooks/useDataSourceDiscovery.ts

@@ -42,11 +42,13 @@ export function useDetectedDataSources(
       dialects: [...connections.values()].map((connection) =>
         connection.dialect.toLowerCase(),
       ),
-      storageProtocols: storageNamespaces.map((namespace) =>
-        namespace.protocol.toLowerCase(),
+      storageProtocols: new Set(
+        storageNamespaces.map((namespace) => namespace.protocol.toLowerCase()),
       ),
-      storageBackendTypes: storageNamespaces.map((namespace) =>
-        namespace.backendType.toLowerCase(),
+      storageBackendTypes: new Set(
+        storageNamespaces.map((namespace) =>
+          namespace.backendType.toLowerCase(),
+        ),
       ),
     };
 

package/src/hooks/useDebounce.ts

@@ -2,6 +2,7 @@
 import { debounce } from "lodash-es";
 import { useEffect, useMemo, useState } from "react";
 import useEvent from "react-use-event-hook";
+import { useInternalStateWithSync } from "./useInternalStateWithSync";
 
 export function useDebounce<T>(value: T, delay: number): T {
   const [debouncedValue, setDebouncedValue] = useState<T>(value);
@@ -41,18 +42,14 @@ export function useDebounceControlledState<T>(opts: {
   disabled?: boolean;
 }) {
   const { initialValue, onChange, delay, disabled } = opts;
-  const [internalValue, setInternalValue] = useState<T>(initialValue);
+  const [internalValue, setInternalValue] = useInternalStateWithSync(
+    initialValue,
+    Object.is,
+  );
   const debouncedValue = useDebounce(internalValue, delay || 200);
 
   const onUpdate = useEvent(onChange);
 
-  // Handle external changes:
-  useEffect(() => {
-    // TODO: Replace this synchronization with a controlled-state primitive.
-    // oxlint-disable-next-line react/no-deriving-state-in-effects -- Synchronize the debounced draft when the controlled value changes.
-    setInternalValue(initialValue);
-  }, [initialValue]);
-
   // Handle debounced updates for user input
   useEffect(() => {
     if (disabled) {
@@ -80,8 +77,7 @@ export function useDebounceControlledState<T>(opts: {
   };
 }
 
-// oxlint-disable-next-line typescript/no-explicit-any
-export function useDebouncedCallback<T extends (...args: any[]) => unknown>(
+export function useDebouncedCallback<T extends (...args: never[]) => unknown>(
   callback: T,
   delay: number,
 ) {

package/src/hooks/useDeepCompareMemoize.ts

@@ -12,7 +12,7 @@ import React from "react";
  * Useful when the re-render if more expensive than the deep compare.
  */
 export function useDeepCompareMemoize<T>(value: T) {
-  const ref = React.useRef<T>(value);
+  const ref = React.useRef(value);
 
   if (!dequal(value, ref.current)) {
     ref.current = value;

package/src/hooks/useDelayElapsed.ts

@@ -1,6 +1,6 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import { useLayoutEffect, useState } from "react";
+import { useEffect, useState } from "react";
 
 /**
  * Returns `false` until `delayMs` has elapsed since mount (or since `delayMs`
@@ -8,20 +8,24 @@ import { useLayoutEffect, useState } from "react";
  * `delayMs <= 0`. Useful for suppressing spinner/fallback flashes.
  */
 export function useDelayElapsed(delayMs: number): boolean {
-  const [elapsed, setElapsed] = useState(delayMs <= 0);
+  const [timer, setTimer] = useState({ delayMs, elapsed: delayMs <= 0 });
 
-  // Layout effect (not a plain effect) so that when `delayMs` flips from `<= 0`
-  // to `> 0` — e.g. a gate re-closing — we reset `elapsed` to `false` before
-  // the browser paints, rather than briefly flashing the fallback for a frame.
-  useLayoutEffect(() => {
+  // Reset during render so a changed delay never commits a stale `elapsed`
+  // (e.g. a gate re-closing shows no fallback flash before paint)
+  if (!Object.is(timer.delayMs, delayMs)) {
+    setTimer({ delayMs, elapsed: delayMs <= 0 });
+  }
+
+  useEffect(() => {
     if (delayMs <= 0) {
-      setElapsed(true);
       return;
     }
-    setElapsed(false);
-    const timeout = setTimeout(() => setElapsed(true), delayMs);
+    const timeout = setTimeout(
+      () => setTimer({ delayMs, elapsed: true }),
+      delayMs,
+    );
     return () => clearTimeout(timeout);
   }, [delayMs]);
 
-  return elapsed;
+  return timer.elapsed;
 }

package/src/hooks/useElapsedTime.ts

@@ -1,13 +1,13 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import { useEffect, useRef, useState } from "react";
+import { useEffect, useState } from "react";
 import { type Milliseconds, Time } from "@/utils/time";
 
 /**
  * Returns the elapsed time since mount, in milliseconds.
  */
 export function useElapsedTime(initialStartTimeMs: Milliseconds) {
-  const startTime = useRef(initialStartTimeMs);
+  const [startTime] = useState(initialStartTimeMs);
   const [endTime, setEndTime] = useState(initialStartTimeMs);
 
   useEffect(() => {
@@ -21,5 +21,5 @@ export function useElapsedTime(initialStartTimeMs: Milliseconds) {
     return () => clearInterval(interval);
   }, []);
 
-  return endTime - startTime.current;
+  return endTime - startTime;
 }

package/src/hooks/useEventListener.ts

@@ -62,8 +62,7 @@ export function useEventListener(
       return;
     }
 
-    // oxlint-disable-next-line typescript/no-explicit-any
-    const eventListener = (event: any) => savedListener.current(event);
+    const eventListener = (event: Event) => savedListener.current(event);
     target.addEventListener(type, eventListener, options);
 
     return () => {

package/src/hooks/useInputHistory.ts

@@ -37,12 +37,9 @@ export function useInputHistory({
   const historyIndexRef = React.useRef<number>(-1);
   // Store the current input when navigating history
   const pendingInputRef = React.useRef<string>("");
-  // Keep a ref to history for use in callbacks without causing re-renders
-  const historyRef = React.useRef(history);
-  historyRef.current = history;
 
   const navigateUp = React.useCallback(() => {
-    const currentHistory = historyRef.current;
+    const currentHistory = history;
     if (currentHistory.length === 0) {
       return;
     }
@@ -59,10 +56,10 @@ export function useInputHistory({
       historyIndexRef.current = newIndex;
       setValue(currentHistory[currentHistory.length - 1 - newIndex]);
     }
-  }, [value, setValue]);
+  }, [history, value, setValue]);
 
   const navigateDown = React.useCallback(() => {
-    const currentHistory = historyRef.current;
+    const currentHistory = history;
     if (currentHistory.length === 0) {
       return;
     }
@@ -79,7 +76,7 @@ export function useInputHistory({
       historyIndexRef.current = -1;
       setValue(pendingInputRef.current);
     }
-  }, [setValue]);
+  }, [history, setValue]);
 
   const addToHistory = React.useCallback((command: string) => {
     setHistory((prev) => {

package/src/hooks/useInternalStateWithSync.ts

@@ -1,5 +1,5 @@
 /* Copyright 2026 Marimo. All rights reserved. */
-import { useRef, useState } from "react";
+import { useState } from "react";
 
 type EqualityFn<T> = (a: T, b: T) => boolean;
 
@@ -13,13 +13,13 @@ export function useInternalStateWithSync<T>(
   initial: T,
   equalityFn: EqualityFn<T> = defaultEqualityFn,
 ) {
-  const [state, setState] = useState(initial);
-  const ref = useRef(initial);
+  const [state, setState] = useState(() => initial);
+  const [previousInitial, setPreviousInitial] = useState(() => initial);
 
   // Whenever the outside state changes, update the internal state
-  if (!equalityFn(ref.current, initial)) {
-    setState(initial); // react is ok to setState during render
-    ref.current = initial;
+  if (!equalityFn(previousInitial, initial)) {
+    setState(() => initial);
+    setPreviousInitial(() => initial);
   }
 
   return [state, setState] as const;

package/src/hooks/useResizeObserver.ts

@@ -6,6 +6,7 @@
 
 import type { RefObject } from "react";
 import { useEffect, useRef } from "react";
+import useEvent from "react-use-event-hook";
 
 interface Size {
   width: number | undefined;
@@ -30,9 +31,7 @@ export function useResizeObserver<T extends HTMLElement = HTMLElement>(
   const { ref, box = "content-box", skip } = options;
   // oxlint-disable-next-line react/hook-use-state
   const previousSize = useRef<Size>({ ...initialSize });
-  const onResize = useRef<((size: Size) => void) | undefined>(undefined);
-  // oxlint-disable-next-line react-hooks/react-compiler
-  onResize.current = options.onResize;
+  const onResize = useEvent((size: Size) => options.onResize?.(size));
 
   useEffect(() => {
     if (!ref.current || skip) {
@@ -63,9 +62,7 @@ export function useResizeObserver<T extends HTMLElement = HTMLElement>(
         previousSize.current.width = newWidth;
         previousSize.current.height = newHeight;
 
-        if (onResize.current) {
-          onResize.current(newSize);
-        }
+        onResize(newSize);
       }
     });
 
@@ -74,7 +71,7 @@ export function useResizeObserver<T extends HTMLElement = HTMLElement>(
     return () => {
       observer.disconnect();
     };
-  }, [box, ref, skip]);
+  }, [box, ref, skip, onResize]);
 }
 
 type BoxSizesKey = keyof Pick<