@marimo-team/islands 0.25.2-dev9 → 0.25.2-dev91

package/src/core/websocket/connection-utils.ts

@@ -1,66 +1,66 @@
 /* Copyright 2026 Marimo. All rights reserved. */
-import { WebSocketState } from "./types";
+import { ConnectionState } from "./types";
 
 /**
  * Check if the app is in a closed/disconnected state
  */
-export function isAppClosed(state: WebSocketState): boolean {
-  return state === WebSocketState.CLOSED;
+export function isAppClosed(state: ConnectionState): boolean {
+  return state === ConnectionState.CLOSED;
 }
 
 /**
  * Check if the app is in a connecting state
  */
-export function isAppConnecting(state: WebSocketState): boolean {
-  return state === WebSocketState.CONNECTING;
+export function isAppConnecting(state: ConnectionState): boolean {
+  return state === ConnectionState.CONNECTING;
 }
 
 /**
  * Check if the app is in an open/connected state
  */
-export function isAppConnected(state: WebSocketState): boolean {
-  return state === WebSocketState.OPEN;
+export function isAppConnected(state: ConnectionState): boolean {
+  return state === ConnectionState.OPEN;
 }
 
 /**
  * Check if the app is in a closing state
  */
-export function isAppClosing(state: WebSocketState): boolean {
-  return state === WebSocketState.CLOSING;
+export function isAppClosing(state: ConnectionState): boolean {
+  return state === ConnectionState.CLOSING;
 }
 
 /**
  * Check if the app is in a not started state
  */
-export function isAppNotStarted(state: WebSocketState): boolean {
-  return state === WebSocketState.NOT_STARTED;
+export function isAppNotStarted(state: ConnectionState): boolean {
+  return state === ConnectionState.NOT_STARTED;
 }
 
 /**
  * Check if the app is in a state where user interactions should be disabled
  */
-export function isAppInteractionDisabled(state: WebSocketState): boolean {
+export function isAppInteractionDisabled(state: ConnectionState): boolean {
   return (
-    state === WebSocketState.CLOSED ||
-    state === WebSocketState.CLOSING ||
-    state === WebSocketState.CONNECTING
+    state === ConnectionState.CLOSED ||
+    state === ConnectionState.CLOSING ||
+    state === ConnectionState.CONNECTING
   );
 }
 
 /**
  * Get a human-readable tooltip message for the connection state
  */
-export function getConnectionTooltip(state: WebSocketState): string {
+export function getConnectionTooltip(state: ConnectionState): string {
   switch (state) {
-    case WebSocketState.CLOSED:
+    case ConnectionState.CLOSED:
       return "App disconnected";
-    case WebSocketState.CONNECTING:
+    case ConnectionState.CONNECTING:
       return "Connecting to a runtime ...";
-    case WebSocketState.CLOSING:
+    case ConnectionState.CLOSING:
       return "App disconnecting...";
-    case WebSocketState.OPEN:
+    case ConnectionState.OPEN:
       return "";
-    case WebSocketState.NOT_STARTED:
+    case ConnectionState.NOT_STARTED:
       return "Click to connect to a runtime";
     default:
       return "";

package/src/core/websocket/types.ts

@@ -6,8 +6,7 @@ export type ConnectionPhase =
   | NotificationMessageData<"startup-progress">["phase"]
   | "reconnecting";
 
-// TODO: rename to ConnectionState
-export const WebSocketState = {
+export const ConnectionState = {
   NOT_STARTED: "NOT_STARTED",
   CONNECTING: "CONNECTING",
   OPEN: "OPEN",
@@ -15,8 +14,8 @@ export const WebSocketState = {
   CLOSED: "CLOSED",
 } as const;
 
-export type WebSocketState =
-  (typeof WebSocketState)[keyof typeof WebSocketState];
+export type ConnectionState =
+  (typeof ConnectionState)[keyof typeof ConnectionState];
 
 export const WebSocketClosedReason = {
   KERNEL_DISCONNECTED: "KERNEL_DISCONNECTED",
@@ -28,7 +27,7 @@ export type WebSocketClosedReason =
 
 export type ConnectionStatus =
   | {
-      state: typeof WebSocketState.CLOSED;
+      state: typeof ConnectionState.CLOSED;
       code: WebSocketClosedReason;
       /**
        * Human-readable reason for closing the connection.
@@ -37,12 +36,12 @@ export type ConnectionStatus =
       phase?: ConnectionPhase;
     }
   | {
-      state: typeof WebSocketState.CONNECTING;
+      state: typeof ConnectionState.CONNECTING;
       phase?: ConnectionPhase;
     }
   | {
       state:
-        | typeof WebSocketState.OPEN
-        | typeof WebSocketState.CLOSING
-        | typeof WebSocketState.NOT_STARTED;
+        | typeof ConnectionState.OPEN
+        | typeof ConnectionState.CLOSING
+        | typeof ConnectionState.NOT_STARTED;
     };

package/src/core/websocket/useMarimoKernelConnection.tsx

@@ -60,7 +60,11 @@ import {
   handleRemoveUIElements,
 } from "../kernel/handlers";
 import { queryParamHandlers } from "../kernel/queryParamHandlers";
-import type { SessionId } from "../kernel/session";
+import {
+  type SessionId,
+  type StableSessionId,
+  stableSessionIdAtom,
+} from "../kernel/session";
 import { initialRunCompletedAtom, kernelStateAtom } from "../kernel/state";
 import { type LayoutState, useLayoutActions } from "../layout/state";
 import { kioskModeAtom } from "../mode";
@@ -80,7 +84,7 @@ import { isWasm } from "../wasm/utils";
 import {
   type ConnectionStatus,
   WebSocketClosedReason,
-  WebSocketState,
+  ConnectionState,
 } from "./types";
 
 const SUPPORTS_LAZY_KERNELS = true;
@@ -112,7 +116,7 @@ export function classifyCloseEvent(event: { reason?: string }): CloseDecision {
       return {
         kind: "gave-up",
         status: {
-          state: WebSocketState.CLOSED,
+          state: ConnectionState.CLOSED,
           code: WebSocketClosedReason.KERNEL_DISCONNECTED,
           reason: "kernel not found",
         },
@@ -124,7 +128,7 @@ export function classifyCloseEvent(event: { reason?: string }): CloseDecision {
       return {
         kind: "terminal",
         status: {
-          state: WebSocketState.CLOSED,
+          state: ConnectionState.CLOSED,
           code: WebSocketClosedReason.KERNEL_DISCONNECTED,
           reason: "kernel not found",
         },
@@ -134,7 +138,7 @@ export function classifyCloseEvent(event: { reason?: string }): CloseDecision {
       return {
         kind: "terminal",
         status: {
-          state: WebSocketState.CLOSED,
+          state: ConnectionState.CLOSED,
           code: WebSocketClosedReason.KERNEL_DISCONNECTED,
           reason: "not authorized",
         },
@@ -144,7 +148,7 @@ export function classifyCloseEvent(event: { reason?: string }): CloseDecision {
       return {
         kind: "terminal",
         status: {
-          state: WebSocketState.CLOSED,
+          state: ConnectionState.CLOSED,
           code: WebSocketClosedReason.KERNEL_DISCONNECTED,
           reason: "kiosk mode is not available for this session",
         },
@@ -154,7 +158,7 @@ export function classifyCloseEvent(event: { reason?: string }): CloseDecision {
       return {
         kind: "terminal",
         status: {
-          state: WebSocketState.CLOSED,
+          state: ConnectionState.CLOSED,
           code: WebSocketClosedReason.KERNEL_STARTUP_ERROR,
           reason: "Failed to start kernel sandbox",
         },
@@ -171,7 +175,7 @@ export function classifyCloseEvent(event: { reason?: string }): CloseDecision {
 
   return {
     kind: "retry",
-    status: { state: WebSocketState.CONNECTING },
+    status: { state: ConnectionState.CONNECTING },
   };
 }
 
@@ -231,6 +235,7 @@ export function useMarimoKernelConnection(opts: {
     addStartupLog,
   } = useAlertActions();
   const setKioskMode = useSetAtom(kioskModeAtom);
+  const setStableSessionId = useSetAtom(stableSessionIdAtom);
   const setCapabilities = useSetAtom(capabilitiesAtom);
   const runtimeManager = useRuntimeManager();
   const transportType = useAtomValue(connectionTransportTypeAtom);
@@ -251,17 +256,17 @@ export function useMarimoKernelConnection(opts: {
         const { phase } = msg.data;
         updateStartupProgress(msg.data);
         setConnection((previous) =>
-          previous.state === WebSocketState.OPEN ||
-          (previous.state === WebSocketState.CONNECTING &&
+          previous.state === ConnectionState.OPEN ||
+          (previous.state === ConnectionState.CONNECTING &&
             previous.phase === phase)
             ? previous
-            : { state: WebSocketState.CONNECTING, phase },
+            : { state: ConnectionState.CONNECTING, phase },
         );
         return;
       }
       case "kernel-ready": {
         setKernelStartupError(null);
-        setConnection({ state: WebSocketState.OPEN });
+        setConnection({ state: ConnectionState.OPEN });
         setInitialRunCompleted(
           Boolean(msg.data.resumed || msg.data.auto_instantiated),
         );
@@ -284,6 +289,10 @@ export function useMarimoKernelConnection(opts: {
           invalidateDataSourceDiscovery();
         }
         setKioskMode(msg.data.kiosk);
+        setStableSessionId(
+          (msg.data.stable_session_id as StableSessionId | null | undefined) ??
+            null,
+        );
         // A freshly started kernel has no breakpoints of its own; re-push
         // the client's set so they still apply, and clear the stale
         // highlighted line from the previous kernel instance.
@@ -295,6 +304,8 @@ export function useMarimoKernelConnection(opts: {
       case "completed-run":
         setInitialRunCompleted(true);
         return;
+      case "execution":
+        return;
       case "interrupted":
         return;
 
@@ -470,7 +481,7 @@ export function useMarimoKernelConnection(opts: {
 
       case "reconnected":
         setKernelStartupError(null);
-        setConnection({ state: WebSocketState.OPEN });
+        setConnection({ state: ConnectionState.OPEN });
         return;
 
       case "focus-cell":
@@ -515,12 +526,12 @@ export function useMarimoKernelConnection(opts: {
     }
     shouldTryReconnecting.current = true;
     setKernelStartupError(null);
-    setConnection({ state: WebSocketState.CONNECTING });
+    setConnection({ state: ConnectionState.CONNECTING });
     const healthy = await runtimeManager.reconcileFromHealth();
     if (!healthy) {
       shouldTryReconnecting.current = false;
       setConnection({
-        state: WebSocketState.CLOSED,
+        state: ConnectionState.CLOSED,
         code: WebSocketClosedReason.KERNEL_DISCONNECTED,
         reason: "kernel not found",
       });
@@ -596,17 +607,17 @@ export function useMarimoKernelConnection(opts: {
       setConnection((previous) => {
         const status = decision.status;
         if (
-          status.state === WebSocketState.CLOSED &&
+          status.state === ConnectionState.CLOSED &&
           status.code === WebSocketClosedReason.KERNEL_STARTUP_ERROR &&
-          (previous.state === WebSocketState.CONNECTING ||
-            previous.state === WebSocketState.CLOSED)
+          (previous.state === ConnectionState.CONNECTING ||
+            previous.state === ConnectionState.CLOSED)
         ) {
           return { ...status, phase: previous.phase };
         }
         if (
-          status.state === WebSocketState.CONNECTING &&
-          (previous.state === WebSocketState.OPEN ||
-            (previous.state === WebSocketState.CONNECTING &&
+          status.state === ConnectionState.CONNECTING &&
+          (previous.state === ConnectionState.OPEN ||
+            (previous.state === ConnectionState.CONNECTING &&
               previous.phase === "reconnecting"))
         ) {
           return { ...status, phase: "reconnecting" };
@@ -631,7 +642,7 @@ export function useMarimoKernelConnection(opts: {
     onError: (e) => {
       Logger.warn("WebSocket error", e);
       setConnection({
-        state: WebSocketState.CLOSED,
+        state: ConnectionState.CLOSED,
         code: WebSocketClosedReason.KERNEL_DISCONNECTED,
         reason: "kernel not found",
       });