@marimo-team/frontend 0.25.3-dev2 → 0.25.3-dev21

package/src/core/websocket/{useMarimoKernelConnection.tsx → useNotebookSession.tsx}

@@ -1,8 +1,8 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import { invalidatePackageData } from "@/core/packages/package-data";
-import { useAtom, useAtomValue, useSetAtom } from "jotai";
-import { useRef } from "react";
+import { useAtom, useAtomValue, useSetAtom, useStore } from "jotai";
+import { useLayoutEffect } from "react";
 import { useErrorBoundary } from "react-error-boundary";
 import { toast } from "@/components/ui/use-toast";
 import { getNotebook, useCellActions } from "@/core/cells/cells";
@@ -12,8 +12,7 @@ import type {
   NotificationMessageData,
   NotificationPayload,
 } from "@/core/kernel/messages";
-import { TRANSPORT_EXHAUSTED_REASON } from "@/core/websocket/transports/ws";
-import { useConnectionTransport } from "@/core/websocket/useWebSocket";
+import { useNotebookClient } from "@/core/websocket/useNotebookClient";
 import { renderHTML } from "@/plugins/core/RenderHTML";
 import {
   handleWidgetMessage,
@@ -21,11 +20,7 @@ import {
 } from "@/plugins/impl/anywidget/registry";
 import { logNever } from "@/utils/assertNever";
 import { prettyError } from "@/utils/errors";
-import {
-  type JsonString,
-  safeExtractSetUIElementMessageBuffers,
-} from "@/utils/json/base64";
-import { jsonParseWithSpecialChar } from "@/utils/json/json-parser";
+import { safeExtractSetUIElementMessageBuffers } from "@/utils/json/base64";
 import { Logger } from "@/utils/Logger";
 import { reloadSafe } from "@/utils/reload-safe";
 import { useAlertActions } from "../alerts/state";
@@ -70,7 +65,7 @@ import { type LayoutState, useLayoutActions } from "../layout/state";
 import { kioskModeAtom } from "../mode";
 import { connectionAtom, startupProgressAtom } from "../network/connection";
 import type { RequestId } from "../network/DeferredRequestRegistry";
-import { useRuntimeManager } from "../runtime/config";
+import { runtimeConfigAtom, useRuntimeManager } from "../runtime/config";
 import { SECRETS_REGISTRY } from "../secrets/request-registry";
 import { isStaticNotebook } from "../static/static-state";
 import {
@@ -81,104 +76,10 @@ import { useStorageActions } from "../storage/state";
 import { useVariablesActions } from "../variables/state";
 import type { VariableName } from "../variables/types";
 import { isWasm } from "../wasm/utils";
-import {
-  type ConnectionStatus,
-  WebSocketClosedReason,
-  ConnectionState,
-} from "./types";
+import { WebSocketClosedReason, ConnectionState } from "./types";
 
 const SUPPORTS_LAZY_KERNELS = true;
 
-// All MARIMO_* reasons except TRANSPORT_EXHAUSTED are emitted by the backend
-// (marimo/_server/api/endpoints/ws_endpoint.py and ws/*.py). Keep in sync with
-// the backend literals. Over the SSE transport these arrive as in-band
-// `close` events rather than WebSocket close frames, so every terminal
-// reason must be classified here — an unclassified reason falls into the
-// retry path.
-export type CloseReason =
-  | "MARIMO_NO_FILE_KEY"
-  | "MARIMO_NO_SESSION_ID"
-  | "MARIMO_NO_SESSION"
-  | "MARIMO_SHUTDOWN"
-  | "MARIMO_KERNEL_STARTUP_ERROR"
-  | "MARIMO_UNAUTHORIZED"
-  | "MARIMO_KIOSK_NOT_ALLOWED"
-  | typeof TRANSPORT_EXHAUSTED_REASON;
-
-export type CloseDecision =
-  | { kind: "terminal"; status: ConnectionStatus; closeTransport: boolean }
-  | { kind: "gave-up"; status: ConnectionStatus }
-  | { kind: "retry"; status: ConnectionStatus };
-
-export function classifyCloseEvent(event: { reason?: string }): CloseDecision {
-  switch (event.reason as CloseReason | undefined) {
-    case TRANSPORT_EXHAUSTED_REASON:
-      return {
-        kind: "gave-up",
-        status: {
-          state: ConnectionState.CLOSED,
-          code: WebSocketClosedReason.KERNEL_DISCONNECTED,
-          reason: "kernel not found",
-        },
-      };
-    case "MARIMO_NO_FILE_KEY":
-    case "MARIMO_NO_SESSION_ID":
-    case "MARIMO_NO_SESSION":
-    case "MARIMO_SHUTDOWN":
-      return {
-        kind: "terminal",
-        status: {
-          state: ConnectionState.CLOSED,
-          code: WebSocketClosedReason.KERNEL_DISCONNECTED,
-          reason: "kernel not found",
-        },
-        closeTransport: true,
-      };
-    case "MARIMO_UNAUTHORIZED":
-      return {
-        kind: "terminal",
-        status: {
-          state: ConnectionState.CLOSED,
-          code: WebSocketClosedReason.KERNEL_DISCONNECTED,
-          reason: "not authorized",
-        },
-        closeTransport: true,
-      };
-    case "MARIMO_KIOSK_NOT_ALLOWED":
-      return {
-        kind: "terminal",
-        status: {
-          state: ConnectionState.CLOSED,
-          code: WebSocketClosedReason.KERNEL_DISCONNECTED,
-          reason: "kiosk mode is not available for this session",
-        },
-        closeTransport: true,
-      };
-    case "MARIMO_KERNEL_STARTUP_ERROR":
-      return {
-        kind: "terminal",
-        status: {
-          state: ConnectionState.CLOSED,
-          code: WebSocketClosedReason.KERNEL_STARTUP_ERROR,
-          reason: "Failed to start kernel sandbox",
-        },
-        closeTransport: true,
-      };
-    default:
-      // Empty/undefined reasons are normal transient closes. Anything else is
-      // an unknown server reason; warn so a new MARIMO_* reason doesn't fall
-      // silently into the retry path.
-      if (event.reason) {
-        logNever(event.reason as never);
-      }
-  }
-
-  return {
-    kind: "retry",
-    status: { state: ConnectionState.CONNECTING },
-  };
-}
-
 function getExistingCells(): CellData[] | undefined {
   if (!SUPPORTS_LAZY_KERNELS) {
     return undefined;
@@ -190,16 +91,12 @@ function getExistingCells(): CellData[] | undefined {
   );
 }
 
-/**
- * Creates a connection to the Marimo kernel and handles incoming messages.
- */
-export function useMarimoKernelConnection(opts: {
+/** Coordinate session startup and apply notifications to the editor. */
+export function useNotebookSession(opts: {
   sessionId: SessionId;
   autoInstantiate: boolean;
   setCells: (cells: CellData[], layout: LayoutState) => void;
 }) {
-  // Track whether we want to try reconnecting.
-  const shouldTryReconnecting = useRef<boolean>(true);
   const { autoInstantiate, sessionId, setCells } = opts;
   const { showBoundary } = useErrorBoundary();
 
@@ -238,6 +135,7 @@ export function useMarimoKernelConnection(opts: {
   const setStableSessionId = useSetAtom(stableSessionIdAtom);
   const setCapabilities = useSetAtom(capabilitiesAtom);
   const runtimeManager = useRuntimeManager();
+  const { notebookId } = useAtomValue(runtimeConfigAtom);
   const transportType = useAtomValue(connectionTransportTypeAtom);
   const setCacheInfo = useSetAtom(cacheInfoAtom);
   const setKernelStartupError = useSetAtom(kernelStartupErrorAtom);
@@ -246,8 +144,7 @@ export function useMarimoKernelConnection(opts: {
     filterFromVariables: filterStorageFromVariables,
   } = useStorageActions();
 
-  const handleMessage = (e: MessageEvent<JsonString<NotificationPayload>>) => {
-    const msg = jsonParseWithSpecialChar(e.data);
+  const handleNotification = (msg: NotificationPayload) => {
     switch (msg.data.op) {
       case "reload":
         reloadSafe();
@@ -505,107 +402,58 @@ export function useMarimoKernelConnection(opts: {
     }
   };
 
-  const tryReconnecting = (code?: number, reason?: string) => {
-    // If not properly gated, we could try reconnecting forever if the
-    // issue is not transient. So we want to try reconnecting only once after an
-    // open connection is closed.
-    if (shouldTryReconnecting.current) {
-      shouldTryReconnecting.current = false;
-      ws.reconnect(code, reason);
-    }
-  };
-
-  // Manual reconnect. Probes /health first to fail fast when the runtime
-  // is unreachable, instead of waiting on partysocket's retry budget.
-  const reconnect = async () => {
-    if (
-      ws.readyState === WebSocket.OPEN ||
-      ws.readyState === WebSocket.CONNECTING
-    ) {
-      return;
-    }
-    shouldTryReconnecting.current = true;
-    setKernelStartupError(null);
-    setConnection({ state: ConnectionState.CONNECTING });
-    const healthy = await runtimeManager.reconcileFromHealth();
-    if (!healthy) {
-      shouldTryReconnecting.current = false;
-      setConnection({
-        state: ConnectionState.CLOSED,
-        code: WebSocketClosedReason.KERNEL_DISCONNECTED,
-        reason: "kernel not found",
-      });
-      return;
-    }
-    ws.reconnect();
+  const store = useStore();
+  const onMessageError = (error: unknown) => {
+    toast({
+      title: "Failed to handle message",
+      description: prettyError(error),
+      variant: "danger",
+    });
   };
-
-  const ws = useConnectionTransport({
+  const notebook = useNotebookClient({
+    sessionId,
     static: isStaticNotebook(),
     transportType,
-    /**
-     * Unique URL for this session.
-     */
+    messageFormat:
+      notebookId && transportType === "websocket" ? "host" : "session",
     url: () =>
       transportType === "sse"
         ? runtimeManager.getSseURL(sessionId).toString()
         : runtimeManager.getWsURL(sessionId).toString(),
-    /**
-     * Auth headers for the SSE transport. WebSockets cannot send headers
-     * and instead put the access token in the URL.
-     */
+    // SSE uses headers for auth. WebSockets carry the token in the URL.
     headers: () => runtimeManager.headers(),
-
-    /**
-     * The transport is open; kernel-ready establishes session readiness.
-     */
-    onOpen: async () => {
-      updateStartupProgress(null);
-      // If we are open, we can reset our reconnecting flag.
-      shouldTryReconnecting.current = true;
-    },
-
-    /**
-     * Wait to connect, in case the remote kernel still starting up.
-     */
-    waitToConnect: async () => {
-      if (isStaticNotebook()) {
-        return;
-      }
-      if (isWasm()) {
-        return;
-      }
-
-      if (runtimeManager.isSameOrigin) {
-        return;
+    waitUntilReady: async () => {
+      if (!isStaticNotebook() && !isWasm() && !runtimeManager.isSameOrigin) {
+        await runtimeManager.waitForHealthy();
       }
-      await runtimeManager.waitForHealthy();
     },
-
-    /**
-     * Handle messages sent by the kernel.
-     */
-    onMessage: (e) => {
+    checkHealth: () => runtimeManager.reconcileFromHealth(),
+    onMessageError,
+    onNotification: (notification) => {
       try {
-        handleMessage(e);
+        handleNotification(notification);
       } catch (error) {
-        Logger.error("Failed to handle message", e.data, error);
-        toast({
-          title: "Failed to handle message",
-          description: prettyError(error),
-          variant: "danger",
-        });
+        Logger.error("Failed to handle notification", notification, error);
+        onMessageError(error);
       }
     },
+  });
 
-    /**
-     * Handle a close event. We may want to reconnect.
-     */
-    onClose: (e) => {
-      Logger.warn("WebSocket closed", e.code, e.reason);
-      const decision = classifyCloseEvent(e);
+  // Subscribe before passive connection setup: static/WASM transports can emit synchronously.
+  useLayoutEffect(() => {
+    const updateConnection = () => {
+      const status = store.get(notebook.connectionAtom);
+      if (status.state === ConnectionState.NOT_STARTED) {
+        return;
+      }
+      if (status.state === ConnectionState.OPEN) {
+        updateStartupProgress(null);
+        return; // kernel-ready, not the socket, establishes session readiness.
+      }
+      if (status.state === ConnectionState.CONNECTING) {
+        setKernelStartupError(null);
+      }
       setConnection((previous) => {
-        const status = decision.status;
         if (
           status.state === ConnectionState.CLOSED &&
           status.code === WebSocketClosedReason.KERNEL_STARTUP_ERROR &&
@@ -624,31 +472,17 @@ export function useMarimoKernelConnection(opts: {
         }
         return status;
       });
-      if (decision.kind === "terminal" && decision.closeTransport) {
-        ws.close(); // close to prevent reconnecting
-        return;
-      }
-      if (decision.kind === "retry") {
-        // Session should be valid
-        // - browser tab might have been closed or re-opened
-        // - computer might have just woken from sleep
-        tryReconnecting(e.code, e.reason);
-      }
-    },
-
-    /**
-     * When we encounter an error, we should close the connection.
-     */
-    onError: (e) => {
-      Logger.warn("WebSocket error", e);
-      setConnection({
-        state: ConnectionState.CLOSED,
-        code: WebSocketClosedReason.KERNEL_DISCONNECTED,
-        reason: "kernel not found",
-      });
-      tryReconnecting();
-    },
-  });
-
-  return { connection, reconnect };
+    };
+    const unsubscribe = store.sub(notebook.connectionAtom, updateConnection);
+    updateConnection();
+    return unsubscribe;
+  }, [
+    notebook,
+    store,
+    setConnection,
+    setKernelStartupError,
+    updateStartupProgress,
+  ]);
+
+  return { notebook, connection, reconnect: notebook.reconnect };
 }

package/src/mount.tsx

@@ -10,7 +10,7 @@ import {
   userConfigAtom,
 } from "@/core/config/config";
 import { KnownQueryParams } from "@/core/constants";
-import { pairPreviewAtom, pairPreviewSchema } from "@/core/config/pair";
+import { pairConfigAtom, pairConfigSchema } from "@/core/config/pair";
 import { getFilenameFromDOM } from "@/core/dom/htmlUtils";
 import { getMarimoCode } from "@/core/meta/globals";
 import {
@@ -209,7 +209,7 @@ const mountOptionsSchema = z.object({
    * marimo config overrides
    */
   configOverrides: passthroughObject,
-  pairPreview: pairPreviewSchema.optional(),
+  pairConfig: pairConfigSchema.optional(),
   /**
    * marimo app config
    */
@@ -274,6 +274,8 @@ const mountOptionsSchema = z.object({
       .transform((val) => val as api.Notebook["NotebookV1"]),
   ]),
 
+  notebookId: z.string().optional(),
+
   /**
    * Runtime configs
    */
@@ -284,6 +286,7 @@ const mountOptionsSchema = z.object({
         // Lazy by default, but can be overridden by the runtime config
         lazy: z.boolean().default(true),
         authToken: z.string().nullish(),
+        notebookId: z.string().optional(),
       }),
     )
     .nullish()
@@ -347,7 +350,7 @@ function initStore(options: unknown) {
   );
   store.set(userConfigAtom, parseUserConfig(parsedOptions.data.config));
   store.set(appConfigAtom, parseAppConfig(parsedOptions.data.appConfig));
-  store.set(pairPreviewAtom, parsedOptions.data.pairPreview);
+  store.set(pairConfigAtom, parsedOptions.data.pairConfig);
 
   // Runtime config
   if (parsedOptions.data.runtimeConfig.length > 0) {
@@ -364,6 +367,7 @@ function initStore(options: unknown) {
   } else {
     store.set(runtimeConfigAtom, {
       ...DEFAULT_RUNTIME_CONFIG,
+      notebookId: parsedOptions.data.notebookId,
       lazy: false,
       serverToken: parsedOptions.data.serverToken,
     });

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

@@ -136,6 +136,90 @@ describe("LoadingDataTableComponent", () => {
     expect(downloadGeoJSON).toHaveBeenCalledTimes(1);
   });
 
+  it("refreshes export metadata when a dataframe transform changes the data", async () => {
+    const host = document.createElement("div");
+    const metadataFor = (name: string) => ({
+      geometry_columns: [{ name, encoding: "objects", crs: "EPSG:4326" }],
+      primary_geometry_column: name,
+      default_geometry_column: name,
+      formats: {
+        parquet: { available: true, reason: null, missing_packages: [] },
+      },
+    });
+    const getExportMetadata = vi
+      .fn()
+      .mockResolvedValueOnce(metadataFor("before"))
+      .mockResolvedValueOnce(metadataFor("after"));
+    const view = (data: string) => (
+      <Provider store={store}>
+        <TooltipProvider>
+          <LoadingDataTableComponent
+            label={null}
+            totalRows={1}
+            pagination={false}
+            pageSize={10}
+            selection={null}
+            showDownload={true}
+            showFilters={false}
+            showColumnSummaries={false}
+            showDataTypes={false}
+            showPageSizeSelector={false}
+            showColumnExplorer={false}
+            showRowExplorer={false}
+            showChartBuilder={false}
+            rowHeaders={[]}
+            fieldTypes={undefined}
+            totalColumns={1}
+            maxColumns="all"
+            hasStableRowId={false}
+            lazy={false}
+            host={host}
+            showSearch={false}
+            value={[]}
+            setValue={vi.fn()}
+            data={data}
+            search={vi.fn().mockResolvedValue({
+              data,
+              total_rows: 1,
+              cell_styles: null,
+              cell_hover_texts: null,
+            })}
+            download_as={vi.fn()}
+            download_geojson={vi.fn()}
+            get_export_metadata={getExportMetadata}
+            get_column_summaries={vi.fn()}
+            get_data_url={vi.fn() as GetDataUrl}
+            get_row_ids={vi.fn() as GetRowIds}
+            get_size_bytes={vi.fn().mockResolvedValue({ size_bytes: null })}
+          />
+        </TooltipProvider>
+      </Provider>
+    );
+    const { rerender } = render(
+      view(JSON.stringify([{ before: "POINT (1 2)" }])),
+    );
+    await waitFor(() =>
+      expect(
+        screen.getByRole("columnheader", { name: "before" }),
+      ).toBeInTheDocument(),
+    );
+    fireEvent.click(screen.getByTestId("export-button"));
+    await waitFor(() =>
+      expect(screen.getByTestId("export-summary-parquet")).toHaveTextContent(
+        "before",
+      ),
+    );
+
+    rerender(view(JSON.stringify([{ after: "POINT (3 4)" }])));
+
+    await waitFor(() =>
+      expect(screen.getByTestId("export-summary-parquet")).toHaveTextContent(
+        "after",
+      ),
+    );
+    expect(getExportMetadata).toHaveBeenCalledTimes(2);
+  });
+
   it("does not log duplicate keys when a pivot becomes empty", async () => {
     const host = document.createElement("div");
     const consoleError = vi

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

@@ -4,6 +4,9 @@ import {
   exactDateTime,
   getDateFormat,
   getShortTimeZone,
+  parseLocalISODate,
+  parseLocalISODateTime,
+  parseTimeToDate,
   prettyDate,
   timeAgo,
 } from "../dates";
@@ -191,6 +194,119 @@ describe("dates", () => {
     });
   });
 
+  describe("parseTimeToDate", () => {
+    it("parses 24-hour times", () => {
+      expect(parseTimeToDate("14:30")).toEqual(new Date(1970, 0, 1, 14, 30));
+      expect(parseTimeToDate("14:30:05")).toEqual(
+        new Date(1970, 0, 1, 14, 30, 5),
+      );
+    });
+
+    it("parses times with an am/pm suffix", () => {
+      expect(parseTimeToDate("2:30 PM")).toEqual(new Date(1970, 0, 1, 14, 30));
+      expect(parseTimeToDate("12:00 am")).toEqual(new Date(1970, 0, 1, 0, 0));
+      expect(parseTimeToDate("12:00 pm")).toEqual(new Date(1970, 0, 1, 12, 0));
+    });
+
+    it("rejects an out-of-range hour instead of letting Date normalize it", () => {
+      // new Date(1970, 0, 1, 25, 0) would silently become 01:00 the next day.
+      expect(parseTimeToDate("25:00")).toBeUndefined();
+    });
+
+    it("rejects an out-of-range minute instead of letting Date normalize it", () => {
+      // new Date(1970, 0, 1, 12, 99) would silently become 13:39.
+      expect(parseTimeToDate("12:99")).toBeUndefined();
+    });
+
+    it("rejects an out-of-range second", () => {
+      expect(parseTimeToDate("12:30:99")).toBeUndefined();
+    });
+
+    it("requires an hour from 1 to 12 with an am/pm suffix", () => {
+      expect(parseTimeToDate("0:30 am")).toBeUndefined();
+      expect(parseTimeToDate("13:30 pm")).toBeUndefined();
+    });
+
+    it("parses milliseconds", () => {
+      expect(parseTimeToDate("14:30:05.123")).toEqual(
+        new Date(1970, 0, 1, 14, 30, 5, 123),
+      );
+    });
+
+    it("pads a short millisecond fraction and truncates a long one", () => {
+      expect(parseTimeToDate("14:30:05.5")).toEqual(
+        new Date(1970, 0, 1, 14, 30, 5, 500),
+      );
+      expect(parseTimeToDate("14:30:05.123456")).toEqual(
+        new Date(1970, 0, 1, 14, 30, 5, 123),
+      );
+    });
+  });
+
+  describe("parseLocalISODate", () => {
+    it("parses a date-only string into local midnight", () => {
+      expect(parseLocalISODate("2024-01-15")).toEqual(new Date(2024, 0, 15));
+    });
+
+    it("returns undefined for a non-date-only string", () => {
+      expect(parseLocalISODate("2024-01-15T00:00:00")).toBeUndefined();
+      expect(parseLocalISODate("")).toBeUndefined();
+      expect(parseLocalISODate("not-a-date")).toBeUndefined();
+    });
+
+    it("returns undefined for a calendar date that doesn't exist", () => {
+      expect(parseLocalISODate("2024-02-30")).toBeUndefined();
+      expect(parseLocalISODate("2024-13-01")).toBeUndefined();
+    });
+  });
+
+  describe("parseLocalISODateTime", () => {
+    it("parses a 'T'-separated datetime into local wall-clock time", () => {
+      expect(parseLocalISODateTime("2024-01-15T08:30:00")).toEqual(
+        new Date(2024, 0, 15, 8, 30, 0),
+      );
+    });
+
+    it("parses a space-separated datetime", () => {
+      expect(parseLocalISODateTime("2024-01-15 08:30:00")).toEqual(
+        new Date(2024, 0, 15, 8, 30, 0),
+      );
+    });
+
+    it("defaults seconds to 0 when omitted", () => {
+      expect(parseLocalISODateTime("2024-01-15T08:30")).toEqual(
+        new Date(2024, 0, 15, 8, 30, 0),
+      );
+    });
+
+    it("ignores a trailing Z instead of shifting to UTC", () => {
+      expect(parseLocalISODateTime("2024-01-15T08:30:00Z")).toEqual(
+        new Date(2024, 0, 15, 8, 30, 0),
+      );
+    });
+
+    it("ignores a trailing UTC offset instead of shifting the time", () => {
+      expect(parseLocalISODateTime("2024-01-15T08:30:00+05:00")).toEqual(
+        new Date(2024, 0, 15, 8, 30, 0),
+      );
+      expect(parseLocalISODateTime("2024-01-15T08:30:00-0500")).toEqual(
+        new Date(2024, 0, 15, 8, 30, 0),
+      );
+    });
+
+    it("ignores sub-second precision", () => {
+      expect(parseLocalISODateTime("2024-01-15T08:30:00.123")).toEqual(
+        new Date(2024, 0, 15, 8, 30, 0),
+      );
+    });
+
+    it("returns undefined for a date-only or unparsable string", () => {
+      expect(parseLocalISODateTime("2024-01-15")).toBeUndefined();
+      expect(parseLocalISODateTime("not-a-datetime")).toBeUndefined();
+      expect(parseLocalISODateTime("")).toBeUndefined();
+    });
+  });
+
   describe("getDateFormat", () => {
     it("returns the correct format for date", () => {
       expect(getDateFormat("2023-05-15")).toBe("yyyy-MM-dd");