@marimo-team/islands 0.25.3-dev26 → 0.25.3-dev261

package/src/stories/editor.stories.tsx

@@ -5,7 +5,7 @@ import { EditorState, type Extension } from "@codemirror/state";
 import { EditorView } from "@codemirror/view";
 import type { Meta, StoryObj } from "@storybook/react-vite";
 import CodeMirror from "@uiw/react-codemirror";
-import { useEffect, useRef } from "react";
+import { useEffect, useMemo, useRef } from "react";
 import { CopilotConfig } from "@/core/codemirror/copilot/copilot-config";
 import { copilotBundle } from "@/core/codemirror/copilot/extension";
 import { OverridingHotkeyProvider } from "@/core/hotkeys/hotkeys";
@@ -52,28 +52,35 @@ const Editor = (opts: { extensions?: Extension[] }): React.ReactNode => {
     });
 
     return () => view.destroy();
-    // oxlint-disable-next-line react-hooks/exhaustive-deps
-  }, [ref.current]);
+  }, [opts.extensions]);
 
   return <div className="cm" ref={ref} />;
 };
 
-export const Primary: Story = {
-  render: (_args, ctx) => (
+const PrimaryEditor = ({ theme }: Pick<CodeMirrorSetupOpts, "theme">) => {
+  const extensions = useMemo(
+    () =>
+      basicBundle({
+        completionConfig: {
+          activate_on_typing: false,
+          copilot: false,
+          signature_hint_on_typing: false,
+        },
+        theme,
+        hotkeys: new OverridingHotkeyProvider({}),
+      } as unknown as CodeMirrorSetupOpts),
+    [theme],
+  );
+
+  return (
     <div className="Cell m-20 w-[60%] overflow-hidden">
-      <Editor
-        extensions={basicBundle({
-          completionConfig: {
-            activate_on_typing: false,
-            copilot: false,
-            signature_hint_on_typing: false,
-          },
-          theme: ctx.globals.theme,
-          hotkeys: new OverridingHotkeyProvider({}),
-        } as unknown as CodeMirrorSetupOpts)}
-      />
+      <Editor extensions={extensions} />
     </div>
-  ),
+  );
+};
+
+export const Primary: Story = {
+  render: (_args, ctx) => <PrimaryEditor theme={ctx.globals.theme} />,
 };
 
 export const DefaultPython: Story = {

package/src/stories/form-wrapper.stories.tsx

@@ -20,7 +20,7 @@ type Story = StoryObj<typeof FormWrapper>;
 const props: FormWrapperProps<string> = {
   currentValue: "currentValue",
   newValue: "currentValue",
-  setValue: (v) => {
+  setValue: () => {
     toast({
       title: "Form submitted",
     });
@@ -133,7 +133,7 @@ export const Validate: Story = {
         <FormWrapper
           {...props}
           shouldValidate={true}
-          validate={async ({ value }) => {
+          validate={async () => {
             const random = Math.random();
             if (random < 0.8) {
               return `Failed to validate. Random number was ${random}`;

package/src/stories/layout/vertical/one-column.stories.tsx

@@ -3,16 +3,17 @@
 import type { Meta } from "@storybook/react-vite";
 import { createStore, Provider } from "jotai";
 import { createRef } from "react";
+import { mockDocumentModel } from "@/__mocks__/document";
 import { cellId } from "@/__tests__/branded";
 import { type NotebookState, notebookAtom } from "@/core/cells/cells";
 import { createCellRuntimeState } from "@/core/cells/types";
 import { defaultUserConfig, parseAppConfig } from "@/core/config/config-schema";
 import { showCodeInRunModeAtom } from "@/core/meta/state";
+import { DocumentProvider } from "@/core/notebook/document";
 import { connectionAtom } from "@/core/network/connection";
 import { requestClientAtom } from "@/core/network/requests";
 import { resolveRequestClient } from "@/core/network/resolve";
-import { ConnectionState } from "@/core/websocket/types";
-import { MultiColumn } from "@/utils/id-tree";
+import { ConnectionState } from "@/core/session/types";
 import type { Milliseconds, Seconds } from "@/utils/time";
 import { CellArray } from "../../../components/editor/renderers/cell-array";
 import { CellsRenderer } from "../../../components/editor/renderers/cells-renderer";
@@ -36,7 +37,6 @@ const createLongReprNotebook = (
 print('hello')
 LongReprA()`,
       edited: false,
-      serializedEditorState: null,
       config: {
         hide_code: hideCode,
         disabled: false,
@@ -53,7 +53,7 @@ LongReprA()`,
       lastExecutionTime: null,
     },
   },
-  cellIds: MultiColumn.from([[cellId]]),
+  cellIds: [cellId],
   cellRuntime: {
     [cellId]: createCellRuntimeState({
       output: {
@@ -88,7 +88,6 @@ LongReprA()`,
   cellLogs: [],
   history: [],
   scrollKey: null,
-  untouchedNewCells: new Set(),
 });
 
 export default {
@@ -108,18 +107,21 @@ const EditModeCodeShown = () => {
   store.set(connectionAtom, { state: ConnectionState.OPEN });
   store.set(requestClientAtom, resolveRequestClient());
   store.set(showCodeInRunModeAtom, true);
+  const { model } = mockDocumentModel(store, null);
 
   return (
     <Provider store={store}>
-      <TooltipProvider>
-        <CellsRenderer appConfig={parseAppConfig({})} mode="edit">
-          <CellArray
-            mode="edit"
-            userConfig={defaultUserConfig()}
-            appConfig={parseAppConfig({})}
-          />
-        </CellsRenderer>
-      </TooltipProvider>
+      <DocumentProvider model={model}>
+        <TooltipProvider>
+          <CellsRenderer appConfig={parseAppConfig({})} mode="edit">
+            <CellArray
+              mode="edit"
+              userConfig={defaultUserConfig()}
+              appConfig={parseAppConfig({})}
+            />
+          </CellsRenderer>
+        </TooltipProvider>
+      </DocumentProvider>
     </Provider>
   );
 };
@@ -133,18 +135,21 @@ const EditModeCodeHidden = () => {
   store.set(connectionAtom, { state: ConnectionState.OPEN });
   store.set(requestClientAtom, resolveRequestClient());
   store.set(showCodeInRunModeAtom, true);
+  const { model } = mockDocumentModel(store, null);
 
   return (
     <Provider store={store}>
-      <TooltipProvider>
-        <CellsRenderer appConfig={parseAppConfig({})} mode="edit">
-          <CellArray
-            mode="edit"
-            userConfig={defaultUserConfig()}
-            appConfig={parseAppConfig({})}
-          />
-        </CellsRenderer>
-      </TooltipProvider>
+      <DocumentProvider model={model}>
+        <TooltipProvider>
+          <CellsRenderer appConfig={parseAppConfig({})} mode="edit">
+            <CellArray
+              mode="edit"
+              userConfig={defaultUserConfig()}
+              appConfig={parseAppConfig({})}
+            />
+          </CellsRenderer>
+        </TooltipProvider>
+      </DocumentProvider>
     </Provider>
   );
 };

package/src/stories/runtime-story.tsx

@@ -1,7 +1,7 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 import type { Meta } from "@storybook/react-vite";
 import { createStore, Provider, useAtomValue, useStore } from "jotai";
-import { useEffect, useRef, useState } from "react";
+import { createRef, useEffect, useRef, useState } from "react";
 import { Button } from "@/components/ui/button";
 import { ConnectionNotice } from "@/components/editor/alerts/connection-notice";
 import PackagesPanel from "@/components/editor/chrome/panels/packages-panel";
@@ -13,11 +13,9 @@ import { Cell } from "@/components/editor/notebook-cell";
 import { VerticalLayoutWrapper } from "@/components/editor/renderers/vertical-layout/vertical-layout-wrapper";
 import { alertAtom } from "@/core/alerts/state";
 import type { EnvironmentOperation } from "@/core/alerts/environment";
-import {
-  createNotebookActions,
-  notebookAtom,
-  notebookReducer,
-} from "@/core/cells/cells";
+import { notebookAtom } from "@/core/cells/cells";
+import { CellId } from "@/core/cells/ids";
+import { createCell, createCellRuntimeState } from "@/core/cells/types";
 import { userConfigAtom } from "@/core/config/config";
 import {
   type AppConfig,
@@ -30,7 +28,7 @@ import { requestClientAtom } from "@/core/network/requests";
 import { createStaticRequests } from "@/core/network/requests-static";
 import { sandboxActionsAtom, sandboxAtom } from "@/core/packages/sandbox-state";
 import { filenameAtom } from "@/core/saving/file-state";
-import { WebSocketClosedReason, ConnectionState } from "@/core/websocket/types";
+import { WebSocketClosedReason, ConnectionState } from "@/core/session/types";
 import { HTTPError } from "@/utils/errors";
 
 const manifest = `requires-python = ">=3.13"
@@ -112,7 +110,7 @@ function RuntimePreview({
       {props.surface === "notebook" && (
         <div className="flex-1 min-w-0 pt-12">
           <ConnectionNotice appConfig={appConfig} onRetry={reconnect} />
-          {notebook.cellIds.inOrderIds.map((cellId) => (
+          {notebook.cellIds.map((cellId) => (
             <VerticalLayoutWrapper key={cellId} appConfig={appConfig}>
               <Cell
                 cellId={cellId}
@@ -349,15 +347,23 @@ export function RuntimeStory(props: Props) {
       log_mode: "replace",
     });
     if (props.existingCells) {
-      const actions = createNotebookActions((action) =>
-        state.set(notebookAtom, (value) => notebookReducer(value, action)),
-      );
-      actions.createNewCell({
-        cellId: "__end__",
-        before: false,
-        code: 'import polars as pl\n\ntrips = pl.read_csv("bike_trips.csv")\ntrips.head()',
-        autoFocus: false,
-      });
+      const cellId = CellId.create();
+      state.set(notebookAtom, (value) => ({
+        ...value,
+        cellIds: [cellId],
+        cellData: {
+          ...value.cellData,
+          [cellId]: createCell({
+            id: cellId,
+            code: 'import polars as pl\n\ntrips = pl.read_csv("bike_trips.csv")\ntrips.head()',
+          }),
+        },
+        cellRuntime: {
+          ...value.cellRuntime,
+          [cellId]: createCellRuntimeState(),
+        },
+        cellHandles: { ...value.cellHandles, [cellId]: createRef() },
+      }));
     }
     state.set(requestClientAtom, {
       ...createStaticRequests(),

package/src/stories/select.stories.tsx

@@ -24,7 +24,7 @@ export default meta;
 type Story = StoryObj;
 
 export const Primary: Story = {
-  render: (args, ctx) => (
+  render: () => (
     <Select>
       <SelectTrigger className="w-64">
         <SelectValue placeholder="Select a fruit" />
@@ -63,7 +63,7 @@ export const Primary: Story = {
 };
 
 export const Clearable: Story = {
-  render: (args, ctx) => {
+  render: () => {
     const [value, setValue] = useState<string>("");
     return (
       <>

package/src/stories/switchable-multi-select.stories.tsx

@@ -15,7 +15,7 @@ export default meta;
 type Story = StoryObj;
 
 export const Primary: Story = {
-  render: (args, ctx) => {
+  render: () => {
     const [value, setValue] = useState<string[]>([]);
     return (
       <TooltipProvider>

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

@@ -16,6 +16,41 @@ import {
 } from "../arrays";
 
 describe("arrays", () => {
+  describe.each([
+    ["intersection", Arrays.intersection, [3, 3, 2], [Number.NaN]],
+    ["difference", Arrays.difference, [1], [1]],
+  ] as const)("%s", (_name, operation, expected, expectedWithNaN) => {
+    it("preserves order and duplicates without mutating readonly inputs", () => {
+      const values = Object.freeze([3, 1, 3, 2]);
+      const other = Object.freeze([2, 3, 3]);
+      expect(operation(values, other)).toEqual(expected);
+      expect(values).toEqual([3, 1, 3, 2]);
+      expect(other).toEqual([2, 3, 3]);
+    });
+
+    it("uses the same equality as array membership, including NaN", () => {
+      expect(operation([Number.NaN, 1], [Number.NaN])).toEqual(expectedWithNaN);
+    });
+
+    it("handles an empty source", () => {
+      expect(operation([], [1])).toEqual([]);
+    });
+  });
+
+  it("handles empty membership lists", () => {
+    expect(Arrays.intersection([1, 1, 2], [])).toEqual([]);
+    expect(Arrays.difference([1, 1, 2], [])).toEqual([1, 1, 2]);
+  });
+
+  it("compares objects by identity", () => {
+    const shared = { id: 1 };
+    const distinct = { id: 1 };
+    expect(Arrays.intersection([shared, distinct], [shared])).toEqual([shared]);
+    expect(Arrays.difference([shared, distinct], [shared])).toEqual([distinct]);
+    expect(Arrays.intersection([shared, distinct], [shared])[0]).toBe(shared);
+    expect(Arrays.difference([shared, distinct], [shared])[0]).toBe(distinct);
+  });
+
   describe("arrayDelete", () => {
     it("should delete an element at the specified index", () => {
       expect(arrayDelete([1, 2, 3], 1)).toEqual([1, 3]);

package/src/utils/__tests__/batch-requests.test.ts

@@ -8,7 +8,7 @@ async function sleep(ms: number) {
 
 describe("batch", () => {
   it("should return the same promise for calls with the same key", async () => {
-    const loader = vi.fn().mockImplementation(async (arg) => {
+    const loader = vi.fn().mockImplementation(async () => {
       await sleep(10);
       return "response";
     });
@@ -25,7 +25,7 @@ describe("batch", () => {
   });
 
   it("should call loader again after the first promise resolves", async () => {
-    const loader = vi.fn().mockImplementation(async (arg) => {
+    const loader = vi.fn().mockImplementation(async () => {
       await sleep(10);
       return "response";
     });
@@ -40,7 +40,7 @@ describe("batch", () => {
   });
 
   it("should handle different keys separately", async () => {
-    const loader = vi.fn().mockImplementation(async (arg) => {
+    const loader = vi.fn().mockImplementation(async () => {
       await sleep(10);
       return "response";
     });

package/src/utils/__tests__/{createReducer.test.ts → createReducer.test.tsx}

@@ -1,5 +1,8 @@
 /* Copyright 2026 Marimo. All rights reserved. */
-import { describe, expect, it, vi } from "vitest";
+import { act, renderHook } from "@testing-library/react";
+import { createStore, getDefaultStore, Provider } from "jotai";
+import type { PropsWithChildren } from "react";
+import { describe, expect, expectTypeOf, it, vi } from "vitest";
 import { Logger } from "@/utils/Logger";
 import { createReducer, createReducerAndAtoms } from "../createReducer";
 
@@ -35,6 +38,68 @@ describe("createReducer", () => {
     expect(dispatch).toHaveBeenCalledWith({ type: "decrement", payload: 3 });
   });
 
+  it("preserves action names and payload types through reducers and atoms", () => {
+    const initialState = () => ({ count: 0 });
+    const { reducer, createActions, useActions } = createReducerAndAtoms(
+      initialState,
+      {
+        increment: (state, payload: number) => ({
+          count: state.count + payload,
+        }),
+        setFromText: (_state, payload: string) => ({ count: payload.length }),
+        reset: () => initialState(),
+      },
+      [
+        (_prevState, _newState, action) => {
+          if (action.type === "increment") {
+            expectTypeOf(action.payload).toEqualTypeOf<number>();
+          }
+          if (action.type === "setFromText") {
+            expectTypeOf(action.payload).toEqualTypeOf<string>();
+          }
+        },
+      ],
+    );
+    const dispatch = vi.fn();
+    const actions = createActions(dispatch);
+
+    expectTypeOf(actions).toEqualTypeOf<{
+      increment: (payload: number) => void;
+      setFromText: (payload: string) => void;
+      reset: () => void;
+    }>();
+    expectTypeOf<ReturnType<typeof useActions>>().toEqualTypeOf<
+      typeof actions
+    >();
+    expectTypeOf(reducer)
+      .parameter(1)
+      .toEqualTypeOf<
+        | { type: "increment"; payload: number }
+        | { type: "setFromText"; payload: string }
+        | { type: "reset"; payload: unknown }
+      >();
+
+    // @ts-expect-error - action names determine their payload types
+    expectTypeOf(reducer).toBeCallableWith(initialState(), {
+      type: "increment",
+      payload: "5",
+    });
+    // @ts-expect-error - action creators reject incorrect payloads
+    expectTypeOf(actions.setFromText).toBeCallableWith(5);
+
+    actions.increment(5);
+    actions.setFromText("hello");
+    actions.reset();
+    expect(dispatch.mock.calls).toEqual([
+      [{ type: "increment", payload: 5 }],
+      [{ type: "setFromText", payload: "hello" }],
+      [{ type: "reset", payload: undefined }],
+    ]);
+    expect(
+      reducer(initialState(), { type: "setFromText", payload: "hello" }),
+    ).toEqual({ count: 5 });
+  });
+
   it("actions are the same reference", () => {
     const initialState = () => ({ count: 0 });
     const reducers = {
@@ -93,6 +158,7 @@ describe("createReducer", () => {
     const loggerSpy = vi.spyOn(Logger, "error");
 
     const state = { count: 0 };
+    // @ts-expect-error - exercise runtime handling of an invalid action
     const newState = reducer(state, { type: "nonexistent", payload: null });
 
     expect(newState).toBe(state);
@@ -119,6 +185,41 @@ describe("createReducerAndAtoms", () => {
     expect(typeof result.useActions).toBe("function");
   });
 
+  it("caches hook actions and supports bypassing middleware", () => {
+    const middleware = vi.fn();
+    const { useActions, valueAtom } = createReducerAndAtoms(
+      () => ({ count: 0 }),
+      {
+        increment: (state, payload: number) => ({
+          count: state.count + payload,
+        }),
+      },
+      [middleware],
+    );
+    const { result, rerender } = renderHook((options) => useActions(options), {
+      initialProps: { skipMiddleware: false },
+    });
+    const actions = result.current;
+    act(() => actions.increment(2));
+    rerender({ skipMiddleware: false });
+    expect(result.current).toBe(actions);
+    expect(getDefaultStore().get(valueAtom)).toEqual({ count: 2 });
+    expect(middleware).toHaveBeenCalledExactlyOnceWith(
+      { count: 0 },
+      { count: 2 },
+      { type: "increment", payload: 2 },
+    );
+
+    middleware.mockClear();
+    rerender({ skipMiddleware: true });
+    const plainActions = result.current;
+    rerender({ skipMiddleware: true });
+    expect(result.current).toBe(plainActions);
+    act(() => result.current.increment(3));
+    expect(getDefaultStore().get(valueAtom)).toEqual({ count: 5 });
+    expect(middleware).not.toHaveBeenCalled();
+  });
+
   it("should apply middleware", () => {
     const initialState = () => ({ count: 0 });
     const reducers = {
@@ -139,6 +240,31 @@ describe("createReducerAndAtoms", () => {
     expect(middleware).toHaveBeenCalledWith(state, { count: 5 }, action);
   });
 
+  it("runs middleware after the write commits, so it sees the new state", () => {
+    const initialState = () => ({ count: 0 });
+    const reducers = {
+      increment: (state: State, payload: number) => ({
+        count: state.count + payload,
+      }),
+    };
+    const seen: number[] = [];
+    const { valueAtom, useActions } = createReducerAndAtoms(
+      initialState,
+      reducers,
+      [(_prev, next) => seen.push(store.get(valueAtom).count, next.count)],
+    );
+    const store = createStore();
+    const wrapper = ({ children }: PropsWithChildren) => (
+      <Provider store={store}>{children}</Provider>
+    );
+    const { result } = renderHook(() => useActions(), { wrapper });
+
+    result.current.increment(5);
+
+    expect(seen).toEqual([5, 5]);
+    expect(store.get(valueAtom)).toEqual({ count: 5 });
+  });
+
   it("should log an error for non-existent action types", () => {
     const initialState = () => ({ count: 0 });
     const reducers = {