@marimo-team/frontend 0.23.7-dev14 → 0.23.7-dev16

package/dist/assets/ws-BV7dcs53.js

@@ -0,0 +1,22 @@
+var p=Object.defineProperty;var y=(s,i,e)=>i in s?p(s,i,{enumerable:!0,configurable:!0,writable:!0,value:e}):s[i]=e;var o=(s,i,e)=>y(s,typeof i!="symbol"?i+"":i,e);var m;(!globalThis.EventTarget||!globalThis.Event)&&console.error(`
+  PartySocket requires a global 'EventTarget' class to be available!
+  You can polyfill this global by adding this to your code before any partysocket imports: 
+  
+  \`\`\`
+  import 'partysocket/event-target-polyfill';
+  \`\`\`
+  Please file an issue at https://github.com/partykit/partykit if you're still having trouble.
+`);var _=class extends Event{constructor(i,e){super("error",e);o(this,"message");o(this,"error");this.message=i.message,this.error=i}},d=class extends Event{constructor(i=1e3,e="",t){super("close",t);o(this,"code");o(this,"reason");o(this,"wasClean",!0);this.code=i,this.reason=e}},u={Event,ErrorEvent:_,CloseEvent:d};function w(s,i){if(!s)throw Error(i)}function b(s){return new s.constructor(s.type,s)}function E(s){return"data"in s?new MessageEvent(s.type,s):"code"in s||"reason"in s?new d(s.code||1999,s.reason||"unknown reason",s):"error"in s?new _(s.error,s):new Event(s.type,s)}var v=typeof process<"u"&&((m=process.versions)==null?void 0:m.node)!==void 0&&typeof document>"u",f=typeof navigator<"u"&&navigator.product==="ReactNative",c=v||f?E:b,h={maxReconnectionDelay:1e4,minReconnectionDelay:1e3+Math.random()*4e3,minUptime:5e3,reconnectionDelayGrowFactor:1.3,connectionTimeout:4e3,maxRetries:1/0,maxEnqueuedMessages:1/0,startClosed:!1,debug:!1},g=!1,T=class a extends EventTarget{constructor(e,t,n={}){super();o(this,"_ws");o(this,"_retryCount",-1);o(this,"_uptimeTimeout");o(this,"_connectTimeout");o(this,"_shouldReconnect",!0);o(this,"_connectLock",!1);o(this,"_binaryType","blob");o(this,"_closeCalled",!1);o(this,"_messageQueue",[]);o(this,"_debugLogger",console.log.bind(console));o(this,"_url");o(this,"_protocols");o(this,"_options");o(this,"onclose",null);o(this,"onerror",null);o(this,"onmessage",null);o(this,"onopen",null);o(this,"_handleOpen",e=>{this._debug("open event");let{minUptime:t=h.minUptime}=this._options;clearTimeout(this._connectTimeout),this._uptimeTimeout=setTimeout(()=>this._acceptOpen(),t),w(this._ws,"WebSocket is not defined"),this._ws.binaryType=this._binaryType,this._messageQueue.forEach(n=>{var r;(r=this._ws)==null||r.send(n)}),this._messageQueue=[],this.onopen&&this.onopen(e),this.dispatchEvent(c(e))});o(this,"_handleMessage",e=>{this._debug("message event"),this.onmessage&&this.onmessage(e),this.dispatchEvent(c(e))});o(this,"_handleError",e=>{this._debug("error event",e.message),this._disconnect(void 0,e.message==="TIMEOUT"?"timeout":void 0),this.onerror&&this.onerror(e),this._debug("exec error listeners"),this.dispatchEvent(c(e)),this._connect()});o(this,"_handleClose",e=>{this._debug("close event"),this._clearTimeouts(),this._shouldReconnect&&this._connect(),this.onclose&&this.onclose(e),this.dispatchEvent(c(e))});this._url=e,this._protocols=t,this._options=n,this._options.startClosed&&(this._shouldReconnect=!1),this._options.debugLogger&&(this._debugLogger=this._options.debugLogger),this._connect()}static get CONNECTING(){return 0}static get OPEN(){return 1}static get CLOSING(){return 2}static get CLOSED(){return 3}get CONNECTING(){return a.CONNECTING}get OPEN(){return a.OPEN}get CLOSING(){return a.CLOSING}get CLOSED(){return a.CLOSED}get binaryType(){return this._ws?this._ws.binaryType:this._binaryType}set binaryType(e){this._binaryType=e,this._ws&&(this._ws.binaryType=e)}get retryCount(){return Math.max(this._retryCount,0)}get bufferedAmount(){return this._messageQueue.reduce((e,t)=>(typeof t=="string"?e+=t.length:t instanceof Blob?e+=t.size:e+=t.byteLength,e),0)+(this._ws?this._ws.bufferedAmount:0)}get extensions(){return this._ws?this._ws.extensions:""}get protocol(){return this._ws?this._ws.protocol:""}get readyState(){return this._ws?this._ws.readyState:this._options.startClosed?a.CLOSED:a.CONNECTING}get url(){return this._ws?this._ws.url:""}get shouldReconnect(){return this._shouldReconnect}close(e=1e3,t){if(this._closeCalled=!0,this._shouldReconnect=!1,this._clearTimeouts(),!this._ws){this._debug("close enqueued: no ws instance");return}if(this._ws.readyState===this.CLOSED){this._debug("close: already closed");return}this._ws.close(e,t)}reconnect(e,t){this._shouldReconnect=!0,this._closeCalled=!1,this._retryCount=-1,!this._ws||this._ws.readyState===this.CLOSED||this._disconnect(e,t),this._connect()}send(e){if(this._ws&&this._ws.readyState===this.OPEN)this._debug("send",e),this._ws.send(e);else{let{maxEnqueuedMessages:t=h.maxEnqueuedMessages}=this._options;this._messageQueue.length<t&&(this._debug("enqueue",e),this._messageQueue.push(e))}}_debug(...e){this._options.debug&&this._debugLogger("RWS>",...e)}_getNextDelay(){let{reconnectionDelayGrowFactor:e=h.reconnectionDelayGrowFactor,minReconnectionDelay:t=h.minReconnectionDelay,maxReconnectionDelay:n=h.maxReconnectionDelay}=this._options,r=0;return this._retryCount>0&&(r=t*e**(this._retryCount-1),r>n&&(r=n)),this._debug("next delay",r),r}_wait(){return new Promise(e=>{setTimeout(e,this._getNextDelay())})}_getNextProtocols(e){if(!e)return Promise.resolve(null);if(typeof e=="string"||Array.isArray(e))return Promise.resolve(e);if(typeof e=="function"){let t=e();if(!t)return Promise.resolve(null);if(typeof t=="string"||Array.isArray(t))return Promise.resolve(t);if(t.then)return t}throw Error("Invalid protocols")}_getNextUrl(e){if(typeof e=="string")return Promise.resolve(e);if(typeof e=="function"){let t=e();if(typeof t=="string")return Promise.resolve(t);if(t.then)return t}throw Error("Invalid URL")}_connect(){if(this._connectLock||!this._shouldReconnect)return;this._connectLock=!0;let{maxRetries:e=h.maxRetries,connectionTimeout:t=h.connectionTimeout}=this._options;if(this._retryCount>=e){this._debug("max retries reached",this._retryCount,">=",e),this._connectLock=!1;return}this._retryCount++,this._debug("connect",this._retryCount),this._removeListeners(),this._wait().then(()=>Promise.all([this._getNextUrl(this._url),this._getNextProtocols(this._protocols||null)])).then(([n,r])=>{if(this._closeCalled){this._connectLock=!1;return}!this._options.WebSocket&&typeof WebSocket>"u"&&!g&&(console.error(`\u203C\uFE0F No WebSocket implementation available. You should define options.WebSocket. 
+
+For example, if you're using node.js, run \`npm install ws\`, and then in your code:
+
+import PartySocket from 'partysocket';
+import WS from 'ws';
+
+const partysocket = new PartySocket({
+  host: "127.0.0.1:1999",
+  room: "test-room",
+  WebSocket: WS
+});
+
+`),g=!0);let l=this._options.WebSocket||WebSocket;this._debug("connect",{url:n,protocols:r}),this._ws=r?new l(n,r):new l(n),this._ws.binaryType=this._binaryType,this._connectLock=!1,this._addListeners(),this._connectTimeout=setTimeout(()=>this._handleTimeout(),t)}).catch(n=>{this._connectLock=!1,this._handleError(new u.ErrorEvent(Error(n.message),this))})}_handleTimeout(){this._debug("timeout event"),this._handleError(new u.ErrorEvent(Error("TIMEOUT"),this))}_disconnect(e=1e3,t){if(this._clearTimeouts(),this._ws){this._removeListeners();try{(this._ws.readyState===this.OPEN||this._ws.readyState===this.CONNECTING)&&this._ws.close(e,t),this._handleClose(new u.CloseEvent(e,t,this))}catch{}}}_acceptOpen(){this._debug("accept open"),this._retryCount=0}_removeListeners(){this._ws&&(this._debug("removeListeners"),this._ws.removeEventListener("open",this._handleOpen),this._ws.removeEventListener("close",this._handleClose),this._ws.removeEventListener("message",this._handleMessage),this._ws.removeEventListener("error",this._handleError))}_addListeners(){this._ws&&(this._debug("addListeners"),this._ws.addEventListener("open",this._handleOpen),this._ws.addEventListener("close",this._handleClose),this._ws.addEventListener("message",this._handleMessage),this._ws.addEventListener("error",this._handleError))}_clearTimeouts(){clearTimeout(this._connectTimeout),clearTimeout(this._uptimeTimeout)}};export{T as t};

package/dist/index.html

@@ -66,7 +66,7 @@
     <marimo-server-token data-token="{{ server_token }}" hidden></marimo-server-token>
     <!-- /TODO -->
     <title>{{ title }}</title>
-    <script type="module" crossorigin src="./assets/index-5VlGukaI.js"></script>
+    <script type="module" crossorigin src="./assets/index-DLZItUgg.js"></script>
     <link rel="modulepreload" crossorigin href="./assets/preload-helper-DItdS47A.js">
     <link rel="modulepreload" crossorigin href="./assets/chunk-LvLJmgfZ.js">
     <link rel="modulepreload" crossorigin href="./assets/react-Bj1aDYRI.js">
@@ -220,7 +220,7 @@
     <link rel="modulepreload" crossorigin href="./assets/memoize-Tp7rARFe.js">
     <link rel="modulepreload" crossorigin href="./assets/get-C-qh_et5.js">
     <link rel="modulepreload" crossorigin href="./assets/_baseSet-CxV9N1bc.js">
-    <link rel="modulepreload" crossorigin href="./assets/state-BaHWS0qi.js">
+    <link rel="modulepreload" crossorigin href="./assets/state-xEqF8Q3P.js">
     <link rel="modulepreload" crossorigin href="./assets/label-DTR8T0AE.js">
     <link rel="modulepreload" crossorigin href="./assets/textarea-bAp21zYj.js">
     <link rel="modulepreload" crossorigin href="./assets/refresh-ccw-C-n2VFP5.js">

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@marimo-team/frontend",
-  "version": "0.23.7-dev14",
+  "version": "0.23.7-dev16",
   "main": "dist/main.js",
   "types": "dist/index.d.ts",
   "type": "module",
@@ -120,7 +120,7 @@
     "lz-string": "^1.5.0",
     "marked": "^15.0.12",
     "mermaid": "^11.12.3",
-    "partysocket": "1.1.10",
+    "partysocket": "1.1.13",
     "path-to-regexp": "^8.4.0",
     "plotly.js": "^3.3.1",
     "pyodide": "0.27.7",

package/src/components/editor/app-container.tsx

@@ -15,6 +15,7 @@ interface Props {
   connection: ConnectionStatus;
   isRunning: boolean;
   width: AppConfig["width"];
+  onReconnect?: () => void;
 }
 
 export const AppContainer: React.FC<PropsWithChildren<Props>> = ({
@@ -22,13 +23,18 @@ export const AppContainer: React.FC<PropsWithChildren<Props>> = ({
   connection,
   isRunning,
   children,
+  onReconnect,
 }) => {
   const connectionState = connection.state;
 
   return (
     <>
       <DynamicFavicon isRunning={isRunning} />
-      <StatusOverlay connection={connection} isRunning={isRunning} />
+      <StatusOverlay
+        connection={connection}
+        isRunning={isRunning}
+        onReconnect={onReconnect}
+      />
       <PyodideLoader>
         <WrappedWithSidebar>
           {/** oxlint-ignore-next-line -- ID is used by other components to grab the DOM element */}

package/src/components/editor/header/__tests__/status.test.tsx

@@ -0,0 +1,108 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+// @vitest-environment jsdom
+
+import { fireEvent, render } from "@testing-library/react";
+import { createStore, Provider as JotaiProvider } from "jotai";
+import type React from "react";
+import { describe, expect, it, vi } from "vitest";
+import { TooltipProvider } from "@/components/ui/tooltip";
+import { viewStateAtom } from "@/core/mode";
+import {
+  type ConnectionStatus,
+  WebSocketClosedReason,
+  WebSocketState,
+} from "@/core/websocket/types";
+import { StatusOverlay } from "../status";
+
+function renderOverlay(
+  connection: ConnectionStatus,
+  onReconnect?: () => void,
+): ReturnType<typeof render> {
+  const store = createStore();
+  store.set(viewStateAtom, { mode: "edit", cellAnchor: null });
+  const wrapper: React.FC<React.PropsWithChildren> = ({ children }) => (
+    <JotaiProvider store={store}>
+      <TooltipProvider>{children}</TooltipProvider>
+    </JotaiProvider>
+  );
+  return render(
+    <StatusOverlay
+      connection={connection}
+      isRunning={false}
+      onReconnect={onReconnect}
+    />,
+    { wrapper },
+  );
+}
+
+describe("StatusOverlay disconnect indicator", () => {
+  it("invokes onReconnect when the disconnect icon is clicked", () => {
+    const onReconnect = vi.fn();
+    const { getByTestId } = renderOverlay(
+      {
+        state: WebSocketState.CLOSED,
+        code: WebSocketClosedReason.KERNEL_DISCONNECTED,
+        reason: "kernel not found",
+      },
+      onReconnect,
+    );
+
+    const icon = getByTestId("disconnected-indicator") as HTMLButtonElement;
+    expect(icon.tagName).toBe("BUTTON");
+    expect(icon.disabled).toBe(false);
+    expect(icon.getAttribute("aria-label")).toBe("Reconnect to app");
+    fireEvent.click(icon);
+    expect(onReconnect).toHaveBeenCalledTimes(1);
+  });
+
+  it("renders a disabled button when no onReconnect is provided", () => {
+    const { getByTestId } = renderOverlay({
+      state: WebSocketState.CLOSED,
+      code: WebSocketClosedReason.KERNEL_DISCONNECTED,
+      reason: "kernel not found",
+    });
+
+    const button = getByTestId("disconnected-indicator");
+    expect((button as HTMLButtonElement).disabled).toBe(true);
+  });
+
+  it.each([
+    [
+      WebSocketClosedReason.MALFORMED_QUERY,
+      "the kernel did not recognize a request; please file a bug with marimo",
+    ],
+    [
+      WebSocketClosedReason.KERNEL_STARTUP_ERROR,
+      "Failed to start kernel sandbox",
+    ],
+  ])(
+    "renders a disabled button for non-recoverable close reason %s",
+    (code, reason) => {
+      const onReconnect = vi.fn();
+      const { getByTestId } = renderOverlay(
+        { state: WebSocketState.CLOSED, code, reason },
+        onReconnect,
+      );
+
+      const button = getByTestId("disconnected-indicator") as HTMLButtonElement;
+      expect(button.disabled).toBe(true);
+      fireEvent.click(button);
+      expect(onReconnect).not.toHaveBeenCalled();
+    },
+  );
+
+  it("does not render the disconnect icon when another tab has taken over", () => {
+    const onReconnect = vi.fn();
+    const { queryByTestId } = renderOverlay(
+      {
+        state: WebSocketState.CLOSED,
+        code: WebSocketClosedReason.ALREADY_RUNNING,
+        reason: "another browser tab is already connected to the kernel",
+        canTakeover: true,
+      },
+      onReconnect,
+    );
+
+    expect(queryByTestId("disconnected-indicator")).toBeNull();
+  });
+});

package/src/components/editor/header/status.tsx

@@ -7,16 +7,26 @@ import { Tooltip } from "@/components/ui/tooltip";
 import { notebookScrollToRunning } from "@/core/cells/actions";
 import { onlyScratchpadIsRunningAtom } from "@/core/cells/cells";
 import { viewStateAtom } from "@/core/mode";
-import { type ConnectionStatus, WebSocketState } from "@/core/websocket/types";
+import {
+  type ConnectionStatus,
+  WebSocketClosedReason,
+  WebSocketState,
+} from "@/core/websocket/types";
 import { cn } from "@/utils/cn";
 
 export const StatusOverlay: React.FC<{
   connection: ConnectionStatus;
   isRunning: boolean;
-}> = ({ connection, isRunning }) => {
+  onReconnect?: () => void;
+}> = ({ connection, isRunning, onReconnect }) => {
   const { mode } = useAtomValue(viewStateAtom);
   const isClosed = connection.state === WebSocketState.CLOSED;
   const isOpen = connection.state === WebSocketState.OPEN;
+  // Only KERNEL_DISCONNECTED is recoverable by a retry. Other terminal
+  // reasons (MALFORMED_QUERY, KERNEL_STARTUP_ERROR) would deterministically
+  // fail the same way; ALREADY_RUNNING is handled by `LockedIcon` below.
+  const canReconnect =
+    isClosed && connection.code === WebSocketClosedReason.KERNEL_DISCONNECTED;
 
   return (
     <>
@@ -28,7 +38,11 @@ export const StatusOverlay: React.FC<{
         )}
       >
         {isOpen && isRunning && <RunningIcon />}
-        {isClosed && !connection.canTakeover && <DisconnectedIcon />}
+        {isClosed && !connection.canTakeover && (
+          <DisconnectedIcon
+            onReconnect={canReconnect ? onReconnect : undefined}
+          />
+        )}
         {isClosed && connection.canTakeover && <LockedIcon />}
       </div>
     </>
@@ -37,13 +51,33 @@ export const StatusOverlay: React.FC<{
 
 const topLeftStatus = "print:hidden pointer-events-auto hover:cursor-pointer";
 
-const DisconnectedIcon = () => (
-  <Tooltip content="App disconnected">
-    <div className={topLeftStatus}>
-      <UnlinkIcon className="w-[25px] h-[25px] text-(--red-11)" />
-    </div>
-  </Tooltip>
-);
+const DisconnectedIcon: React.FC<{ onReconnect?: () => void }> = ({
+  onReconnect,
+}) => {
+  const disabled = !onReconnect;
+  return (
+    <Tooltip
+      content={
+        disabled ? "App disconnected" : "App disconnected — click to reconnect"
+      }
+    >
+      {/* Wrapper span keeps the tooltip reachable when the button is
+          disabled — a disabled <button> swallows pointer events. */}
+      <span tabIndex={disabled ? 0 : -1}>
+        <button
+          type="button"
+          className={cn(topLeftStatus, "bg-transparent border-0 p-0")}
+          aria-label={disabled ? "App disconnected" : "Reconnect to app"}
+          data-testid="disconnected-indicator"
+          onClick={onReconnect}
+          disabled={disabled}
+        >
+          <UnlinkIcon className="w-[25px] h-[25px] text-(--red-11)" />
+        </button>
+      </span>
+    </Tooltip>
+  );
+};
 
 const LockedIcon = () => (
   <Tooltip content="Notebook locked">

package/src/components/pages/run-page.tsx

@@ -34,7 +34,10 @@ const RunPage = (props: Props) => {
 
 const Watermark = () => {
   return (
-    <div className="fixed bottom-0 right-0 z-50" data-testid="watermark">
+    <div
+      className="fixed bottom-0 right-0 z-50 print:hidden"
+      data-testid="watermark"
+    >
       <a
         href={Constants.githubPage}
         target="_blank"

package/src/core/edit-app.tsx

@@ -79,7 +79,7 @@ export const EditApp: React.FC<AppProps> = ({
     };
   }, []);
 
-  const { connection } = useMarimoKernelConnection({
+  const { connection, reconnect } = useMarimoKernelConnection({
     autoInstantiate: userConfig.runtime.auto_instantiate,
     setCells: (cells, layout) => {
       setCells(cells);
@@ -147,6 +147,7 @@ export const EditApp: React.FC<AppProps> = ({
         connection={connection}
         isRunning={isRunning}
         width={appConfig.width}
+        onReconnect={reconnect}
       >
         <AppHeader
           connection={connection}

package/src/core/run-app.tsx

@@ -38,7 +38,7 @@ export const RunApp: React.FC<AppProps> = ({ appConfig }) => {
     };
   }, []);
 
-  const { connection } = useMarimoKernelConnection({
+  const { connection, reconnect } = useMarimoKernelConnection({
     autoInstantiate: true,
     setCells: setCells,
     sessionId: getSessionId(),
@@ -84,6 +84,7 @@ export const RunApp: React.FC<AppProps> = ({ appConfig }) => {
       connection={connection}
       isRunning={isRunning}
       width={appConfig.width}
+      onReconnect={reconnect}
     >
       <AppHeader connection={connection} className="sm:pt-8">
         {galleryHref && (

package/src/core/websocket/__tests__/useMarimoKernelConnection.hook.test.tsx

@@ -0,0 +1,155 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+// @vitest-environment jsdom
+
+import { act, renderHook } from "@testing-library/react";
+import { createStore, Provider as JotaiProvider } from "jotai";
+import type React from "react";
+import { ErrorBoundary } from "react-error-boundary";
+import { beforeEach, describe, expect, it, vi } from "vitest";
+
+vi.mock("@/core/websocket/useWebSocket", async () => {
+  const actual =
+    await vi.importActual<typeof import("../useWebSocket")>("../useWebSocket");
+  return {
+    ...actual,
+    useConnectionTransport: vi.fn(),
+  };
+});
+
+vi.mock("@/core/runtime/config", async () => {
+  const actual = await vi.importActual<typeof import("@/core/runtime/config")>(
+    "@/core/runtime/config",
+  );
+  return {
+    ...actual,
+    useRuntimeManager: vi.fn(),
+  };
+});
+
+import { useRuntimeManager } from "@/core/runtime/config";
+import { connectionAtom } from "../../network/connection";
+import type { SessionId } from "../../kernel/session";
+import { WebSocketClosedReason, WebSocketState } from "../types";
+import { useMarimoKernelConnection } from "../useMarimoKernelConnection";
+import { useConnectionTransport } from "../useWebSocket";
+
+interface MockTransport {
+  readyState: 0 | 1 | 2 | 3;
+  retryCount: number;
+  reconnect: ReturnType<typeof vi.fn>;
+  close: ReturnType<typeof vi.fn>;
+  send: ReturnType<typeof vi.fn>;
+  addEventListener: ReturnType<typeof vi.fn>;
+  removeEventListener: ReturnType<typeof vi.fn>;
+}
+
+function makeTransport(
+  readyState: 0 | 1 | 2 | 3 = WebSocket.CLOSED,
+): MockTransport {
+  return {
+    readyState,
+    retryCount: 0,
+    reconnect: vi.fn(),
+    close: vi.fn(),
+    send: vi.fn(),
+    addEventListener: vi.fn(),
+    removeEventListener: vi.fn(),
+  };
+}
+
+function makeRuntimeManager(isHealthy = vi.fn().mockResolvedValue(true)) {
+  return {
+    isHealthy,
+    getWsURL: () => new URL("ws://localhost/ws"),
+    waitForHealthy: vi.fn().mockResolvedValue(undefined),
+    isSameOrigin: true,
+  };
+}
+
+describe("useMarimoKernelConnection.reconnect()", () => {
+  let transport: MockTransport;
+  let isHealthy: ReturnType<typeof vi.fn>;
+  let store: ReturnType<typeof createStore>;
+
+  beforeEach(() => {
+    transport = makeTransport(WebSocket.CLOSED);
+    isHealthy = vi.fn().mockResolvedValue(true);
+    store = createStore();
+    store.set(connectionAtom, {
+      state: WebSocketState.CLOSED,
+      code: WebSocketClosedReason.KERNEL_DISCONNECTED,
+      reason: "kernel not found",
+    });
+    vi.mocked(useConnectionTransport).mockReturnValue(transport);
+    vi.mocked(useRuntimeManager).mockReturnValue(
+      makeRuntimeManager(isHealthy) as unknown as ReturnType<
+        typeof useRuntimeManager
+      >,
+    );
+  });
+
+  function renderUseHook() {
+    const wrapper: React.FC<React.PropsWithChildren> = ({ children }) => (
+      <JotaiProvider store={store}>
+        <ErrorBoundary fallback={null}>{children}</ErrorBoundary>
+      </JotaiProvider>
+    );
+    return renderHook(
+      () =>
+        useMarimoKernelConnection({
+          sessionId: "test-session" as SessionId,
+          autoInstantiate: false,
+          setCells: () => {},
+        }),
+      { wrapper },
+    );
+  }
+
+  it("is a no-op when the transport is already OPEN", async () => {
+    transport.readyState = WebSocket.OPEN;
+    const { result } = renderUseHook();
+    await act(async () => {
+      await result.current.reconnect();
+    });
+    expect(isHealthy).not.toHaveBeenCalled();
+    expect(transport.reconnect).not.toHaveBeenCalled();
+  });
+
+  it("is a no-op when the transport is already CONNECTING", async () => {
+    transport.readyState = WebSocket.CONNECTING;
+    const { result } = renderUseHook();
+    await act(async () => {
+      await result.current.reconnect();
+    });
+    expect(isHealthy).not.toHaveBeenCalled();
+    expect(transport.reconnect).not.toHaveBeenCalled();
+  });
+
+  it("probes /health and reconnects when the runtime is healthy", async () => {
+    isHealthy.mockResolvedValue(true);
+    const { result } = renderUseHook();
+    await act(async () => {
+      await result.current.reconnect();
+    });
+    expect(isHealthy).toHaveBeenCalledOnce();
+    expect(transport.reconnect).toHaveBeenCalledOnce();
+    expect(store.get(connectionAtom)).toEqual({
+      state: WebSocketState.CONNECTING,
+    });
+  });
+
+  it("transitions to CLOSED and does not call ws.reconnect when the probe fails", async () => {
+    isHealthy.mockResolvedValue(false);
+    const { result } = renderUseHook();
+    await act(async () => {
+      await result.current.reconnect();
+    });
+    expect(isHealthy).toHaveBeenCalledOnce();
+    expect(transport.reconnect).not.toHaveBeenCalled();
+    expect(store.get(connectionAtom)).toEqual({
+      state: WebSocketState.CLOSED,
+      code: WebSocketClosedReason.KERNEL_DISCONNECTED,
+      reason: "kernel not found",
+    });
+  });
+});

package/src/core/websocket/__tests__/useMarimoKernelConnection.test.ts

@@ -0,0 +1,137 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { afterEach, describe, expect, it, vi } from "vitest";
+import { Logger } from "@/utils/Logger";
+import { WebSocketClosedReason, WebSocketState } from "../types";
+import { classifyCloseEvent } from "../useMarimoKernelConnection";
+import { MAX_RETRIES } from "../useWebSocket";
+
+function classify(
+  reason: string | undefined,
+  retryCount = 0,
+  maxRetries = MAX_RETRIES,
+) {
+  return classifyCloseEvent({ reason }, { retryCount, maxRetries });
+}
+
+describe("classifyCloseEvent", () => {
+  describe("transient closes (default branch)", () => {
+    it("retries when retryCount < maxRetries", () => {
+      const decision = classify(undefined, 0);
+      expect(decision.kind).toBe("retry");
+      expect(decision.status).toEqual({ state: WebSocketState.CONNECTING });
+    });
+
+    it("retries on each intermediate close event during a retry storm", () => {
+      for (let n = 0; n < MAX_RETRIES; n++) {
+        const decision = classify(undefined, n);
+        expect(decision.kind).toBe("retry");
+        expect(decision.status).toEqual({ state: WebSocketState.CONNECTING });
+      }
+    });
+
+    it("transitions to CLOSED when retryCount reaches maxRetries", () => {
+      const decision = classify(undefined, MAX_RETRIES);
+      expect(decision.kind).toBe("gave-up");
+      expect(decision.status).toEqual({
+        state: WebSocketState.CLOSED,
+        code: WebSocketClosedReason.KERNEL_DISCONNECTED,
+        reason: "kernel not found",
+      });
+    });
+
+    it("transitions to CLOSED when retryCount exceeds maxRetries", () => {
+      const decision = classify(undefined, MAX_RETRIES + 5);
+      expect(decision.kind).toBe("gave-up");
+    });
+
+    it("treats unknown reason strings as transient and logs a warning", () => {
+      const logger = vi.spyOn(Logger, "warn").mockImplementation(() => {});
+      const decision = classify("something-else", 3);
+      expect(decision.kind).toBe("retry");
+      expect(logger).toHaveBeenCalled();
+      logger.mockRestore();
+    });
+  });
+
+  afterEach(() => {
+    vi.restoreAllMocks();
+  });
+
+  describe("terminal closes (server-initiated)", () => {
+    it("MARIMO_ALREADY_CONNECTED → terminal + closeTransport, with takeover", () => {
+      const decision = classify("MARIMO_ALREADY_CONNECTED", 0);
+      expect(decision.kind).toBe("terminal");
+      expect(decision.status).toMatchObject({
+        state: WebSocketState.CLOSED,
+        code: WebSocketClosedReason.ALREADY_RUNNING,
+        canTakeover: true,
+      });
+      if (decision.kind === "terminal") {
+        expect(decision.closeTransport).toBe(true);
+      }
+    });
+
+    it.each([
+      "MARIMO_WRONG_KERNEL_ID",
+      "MARIMO_NO_FILE_KEY",
+      "MARIMO_NO_SESSION_ID",
+      "MARIMO_NO_SESSION",
+      "MARIMO_SHUTDOWN",
+    ])("%s → terminal with KERNEL_DISCONNECTED, closes transport", (reason) => {
+      const decision = classify(reason, 0);
+      expect(decision.kind).toBe("terminal");
+      expect(decision.status).toMatchObject({
+        state: WebSocketState.CLOSED,
+        code: WebSocketClosedReason.KERNEL_DISCONNECTED,
+      });
+      if (decision.kind === "terminal") {
+        expect(decision.closeTransport).toBe(true);
+      }
+    });
+
+    it("MARIMO_MALFORMED_QUERY → terminal but does NOT close transport", () => {
+      const decision = classify("MARIMO_MALFORMED_QUERY", 0);
+      expect(decision.kind).toBe("terminal");
+      expect(decision.status).toMatchObject({
+        state: WebSocketState.CLOSED,
+        code: WebSocketClosedReason.MALFORMED_QUERY,
+      });
+      if (decision.kind === "terminal") {
+        expect(decision.closeTransport).toBe(false);
+      }
+    });
+
+    it("MARIMO_KERNEL_STARTUP_ERROR → terminal + closeTransport", () => {
+      const decision = classify("MARIMO_KERNEL_STARTUP_ERROR", 0);
+      expect(decision.kind).toBe("terminal");
+      expect(decision.status).toMatchObject({
+        state: WebSocketState.CLOSED,
+        code: WebSocketClosedReason.KERNEL_STARTUP_ERROR,
+      });
+      if (decision.kind === "terminal") {
+        expect(decision.closeTransport).toBe(true);
+      }
+    });
+
+    it("terminal closes ignore retryCount entirely", () => {
+      const decision = classify("MARIMO_SHUTDOWN", 99);
+      expect(decision.kind).toBe("terminal");
+    });
+  });
+
+  describe("retry budget exhaustion", () => {
+    it("yields retry on attempts 1..maxRetries-1 and gave-up on the final close", () => {
+      const states: string[] = [];
+      for (let attempt = 1; attempt <= MAX_RETRIES; attempt++) {
+        states.push(classify(undefined, attempt - 1).kind);
+      }
+      states.push(classify(undefined, MAX_RETRIES).kind);
+
+      expect(states).toEqual([
+        ...Array.from({ length: MAX_RETRIES }, () => "retry"),
+        "gave-up",
+      ]);
+    });
+  });
+});

package/src/core/websocket/transports/basic.ts

@@ -43,6 +43,8 @@ export class BasicTransport implements IConnectionTransport {
     return WebSocket.OPEN;
   }
 
+  readonly retryCount = 0;
+
   reconnect(code?: number | undefined, reason?: string | undefined): void {
     this.close();
     this.connect();

package/src/core/websocket/transports/transport.ts

@@ -25,6 +25,7 @@ export interface IConnectionTransport {
     callback: ConnectionTransportCallback<T>,
   ): void;
   readyState: WebSocket["readyState"];
+  readonly retryCount: number;
 }
 
 export class ConnectionSubscriptions {