@marimo-team/islands 0.23.7-dev55 → 0.23.7-dev57

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

@@ -4,50 +4,22 @@ 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 });
+function classify(reason: string | undefined) {
+  return classifyCloseEvent({ reason });
 }
 
 describe("classifyCloseEvent", () => {
   describe("transient closes (default branch)", () => {
-    it("retries when retryCount < maxRetries", () => {
-      const decision = classify(undefined, 0);
+    it("retries on empty/undefined reason", () => {
+      const decision = classify(undefined);
       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);
+      const decision = classify("something-else");
       expect(decision.kind).toBe("retry");
       expect(logger).toHaveBeenCalled();
       logger.mockRestore();
@@ -60,7 +32,7 @@ describe("classifyCloseEvent", () => {
 
   describe("terminal closes (server-initiated)", () => {
     it("MARIMO_ALREADY_CONNECTED → terminal + closeTransport, with takeover", () => {
-      const decision = classify("MARIMO_ALREADY_CONNECTED", 0);
+      const decision = classify("MARIMO_ALREADY_CONNECTED");
       expect(decision.kind).toBe("terminal");
       expect(decision.status).toMatchObject({
         state: WebSocketState.CLOSED,
@@ -79,7 +51,7 @@ describe("classifyCloseEvent", () => {
       "MARIMO_NO_SESSION",
       "MARIMO_SHUTDOWN",
     ])("%s → terminal with KERNEL_DISCONNECTED, closes transport", (reason) => {
-      const decision = classify(reason, 0);
+      const decision = classify(reason);
       expect(decision.kind).toBe("terminal");
       expect(decision.status).toMatchObject({
         state: WebSocketState.CLOSED,
@@ -91,7 +63,7 @@ describe("classifyCloseEvent", () => {
     });
 
     it("MARIMO_MALFORMED_QUERY → terminal but does NOT close transport", () => {
-      const decision = classify("MARIMO_MALFORMED_QUERY", 0);
+      const decision = classify("MARIMO_MALFORMED_QUERY");
       expect(decision.kind).toBe("terminal");
       expect(decision.status).toMatchObject({
         state: WebSocketState.CLOSED,
@@ -103,7 +75,7 @@ describe("classifyCloseEvent", () => {
     });
 
     it("MARIMO_KERNEL_STARTUP_ERROR → terminal + closeTransport", () => {
-      const decision = classify("MARIMO_KERNEL_STARTUP_ERROR", 0);
+      const decision = classify("MARIMO_KERNEL_STARTUP_ERROR");
       expect(decision.kind).toBe("terminal");
       expect(decision.status).toMatchObject({
         state: WebSocketState.CLOSED,
@@ -113,25 +85,17 @@ describe("classifyCloseEvent", () => {
         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",
-      ]);
+  describe("transport exhaustion", () => {
+    it("MARIMO_TRANSPORT_EXHAUSTED → gave-up with KERNEL_DISCONNECTED", () => {
+      const decision = classify("MARIMO_TRANSPORT_EXHAUSTED");
+      expect(decision.kind).toBe("gave-up");
+      expect(decision.status).toEqual({
+        state: WebSocketState.CLOSED,
+        code: WebSocketClosedReason.KERNEL_DISCONNECTED,
+        reason: "kernel not found",
+      });
     });
   });
 });

package/src/core/websocket/transports/__tests__/ws.test.ts

@@ -0,0 +1,125 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import { MAX_RETRIES, WsTransport, TRANSPORT_EXHAUSTED_REASON } from "../ws";
+
+let innerListeners: Record<string, ((e: unknown) => void)[]>;
+
+vi.mock("partysocket/ws", () => {
+  class FakeReconnectingWebSocket {
+    retryCount = 0;
+    readyState = WebSocket.CONNECTING;
+    constructor() {
+      innerListeners = { open: [], close: [], message: [], error: [] };
+    }
+    addEventListener(event: string, cb: (e: unknown) => void) {
+      innerListeners[event].push(cb);
+    }
+    removeEventListener(event: string, cb: (e: unknown) => void) {
+      innerListeners[event] = innerListeners[event].filter((c) => c !== cb);
+    }
+    reconnect() {}
+    close() {}
+    send() {}
+  }
+  return { default: FakeReconnectingWebSocket };
+});
+
+interface FakeReconnectingWebSocket {
+  retryCount: number;
+  readyState: number;
+}
+
+function dispatchClose(reason = "") {
+  const evt = new CloseEvent("close", { reason, code: 1006 });
+  for (const cb of innerListeners.close) {
+    cb(evt);
+  }
+}
+
+describe("WsTransport", () => {
+  let transport: WsTransport;
+  let inner: FakeReconnectingWebSocket;
+
+  beforeEach(() => {
+    transport = new WsTransport(() => "ws://example.invalid/ws");
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    inner = (transport as any).inner;
+  });
+
+  describe("close event reason", () => {
+    it("forwards the original reason before retry budget is exhausted", () => {
+      const seen: CloseEvent[] = [];
+      transport.addEventListener("close", (e) => seen.push(e));
+
+      inner.retryCount = MAX_RETRIES - 1;
+      dispatchClose("");
+
+      expect(seen).toHaveLength(1);
+      expect(seen[0].reason).toBe("");
+    });
+
+    it("rewrites reason to MARIMO_TRANSPORT_EXHAUSTED on exhaustion", () => {
+      const seen: CloseEvent[] = [];
+      transport.addEventListener("close", (e) => seen.push(e));
+
+      inner.retryCount = MAX_RETRIES;
+      dispatchClose("");
+
+      expect(seen).toHaveLength(1);
+      expect(seen[0].reason).toBe(TRANSPORT_EXHAUSTED_REASON);
+    });
+
+    it("does not rewrite a server-sent reason on exhaustion", () => {
+      // If partysocket happens to deliver a real MARIMO_* reason at the same
+      // moment the retry budget exhausts, the exhausted-state rewrite still
+      // wins because the wrapper keys off retryCount, not the original reason.
+      // Document the behavior so a future change is deliberate.
+      const seen: CloseEvent[] = [];
+      transport.addEventListener("close", (e) => seen.push(e));
+
+      inner.retryCount = MAX_RETRIES;
+      dispatchClose("MARIMO_SHUTDOWN");
+
+      expect(seen[0].reason).toBe(TRANSPORT_EXHAUSTED_REASON);
+    });
+  });
+
+  describe("addEventListener dedupe", () => {
+    it("does not double-fire when the same close listener is added twice", () => {
+      const cb = vi.fn();
+      transport.addEventListener("close", cb);
+      transport.addEventListener("close", cb);
+
+      dispatchClose("");
+
+      expect(cb).toHaveBeenCalledTimes(1);
+    });
+
+    it("a single removeEventListener fully unregisters a duplicated add", () => {
+      const cb = vi.fn();
+      transport.addEventListener("close", cb);
+      transport.addEventListener("close", cb);
+      transport.removeEventListener("close", cb);
+
+      dispatchClose("");
+
+      expect(cb).not.toHaveBeenCalled();
+      // Inner socket has no orphaned wrappers left.
+      expect(innerListeners.close).toHaveLength(0);
+    });
+  });
+
+  describe("removeEventListener", () => {
+    it("unregisters the right wrapper for close listeners", () => {
+      const cb = vi.fn();
+      transport.addEventListener("close", cb);
+      transport.removeEventListener("close", cb);
+
+      inner.retryCount = MAX_RETRIES;
+      dispatchClose("");
+
+      expect(cb).not.toHaveBeenCalled();
+    });
+  });
+});

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

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

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

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

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

@@ -0,0 +1,96 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import ReconnectingWebSocket from "partysocket/ws";
+import type {
+  ConnectionEvent,
+  ConnectionTransportCallback,
+  IConnectionTransport,
+} from "./transport";
+
+// Per-`reconnect()` retry budget. After exhaustion, partysocket stops silently;
+// the wrapper rewrites the close-event reason to surface the give-up.
+export const MAX_RETRIES = 10;
+
+export const TRANSPORT_EXHAUSTED_REASON = "MARIMO_TRANSPORT_EXHAUSTED";
+
+export class WsTransport implements IConnectionTransport {
+  private inner: ReconnectingWebSocket;
+  private closeWrappers = new WeakMap<
+    ConnectionTransportCallback<"close">,
+    ConnectionTransportCallback<"close">
+  >();
+
+  constructor(urlProvider: () => string) {
+    this.inner = new ReconnectingWebSocket(urlProvider, undefined, {
+      maxRetries: MAX_RETRIES,
+      debug: false,
+      startClosed: true,
+      // long timeout — the server can become slow when many notebooks are open.
+      connectionTimeout: 10_000,
+    });
+  }
+
+  get readyState(): WebSocket["readyState"] {
+    return this.inner.readyState as WebSocket["readyState"];
+  }
+
+  reconnect(code?: number, reason?: string): void {
+    this.inner.reconnect(code, reason);
+  }
+
+  close(): void {
+    this.inner.close();
+  }
+
+  send(data: string | ArrayBuffer | Blob | ArrayBufferView): void {
+    this.inner.send(data);
+  }
+
+  addEventListener<T extends ConnectionEvent>(
+    event: T,
+    callback: ConnectionTransportCallback<T>,
+  ): void {
+    if (event === "close") {
+      const userCb = callback as ConnectionTransportCallback<"close">;
+      // Match native EventTarget dedupe: a second addEventListener with the
+      // same listener is a no-op. Without this, repeated adds leak wrappers
+      // on the inner socket and double-fire on close.
+      if (this.closeWrappers.has(userCb)) {
+        return;
+      }
+      const wrapper: ConnectionTransportCallback<"close"> = (e) => {
+        if (this.inner.retryCount >= MAX_RETRIES) {
+          userCb(
+            new CloseEvent("close", {
+              code: e.code,
+              reason: TRANSPORT_EXHAUSTED_REASON,
+              wasClean: e.wasClean,
+            }),
+          );
+        } else {
+          userCb(e);
+        }
+      };
+      this.closeWrappers.set(userCb, wrapper);
+      this.inner.addEventListener("close", wrapper);
+      return;
+    }
+    this.inner.addEventListener(event, callback as never);
+  }
+
+  removeEventListener<T extends ConnectionEvent>(
+    event: T,
+    callback: ConnectionTransportCallback<T>,
+  ): void {
+    if (event === "close") {
+      const userCb = callback as ConnectionTransportCallback<"close">;
+      const wrapper = this.closeWrappers.get(userCb);
+      if (wrapper) {
+        this.closeWrappers.delete(userCb);
+        this.inner.removeEventListener("close", wrapper);
+      }
+      return;
+    }
+    this.inner.removeEventListener(event, callback as never);
+  }
+}

package/src/core/websocket/useMarimoKernelConnection.tsx

@@ -11,10 +11,8 @@ import type {
   NotificationMessageData,
   NotificationPayload,
 } from "@/core/kernel/messages";
-import {
-  MAX_RETRIES,
-  useConnectionTransport,
-} from "@/core/websocket/useWebSocket";
+import { TRANSPORT_EXHAUSTED_REASON } from "@/core/websocket/transports/ws";
+import { useConnectionTransport } from "@/core/websocket/useWebSocket";
 import { renderHTML } from "@/plugins/core/RenderHTML";
 import {
   handleWidgetMessage,
@@ -80,16 +78,27 @@ import {
 
 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.
+export type CloseReason =
+  | "MARIMO_ALREADY_CONNECTED"
+  | "MARIMO_WRONG_KERNEL_ID"
+  | "MARIMO_NO_FILE_KEY"
+  | "MARIMO_NO_SESSION_ID"
+  | "MARIMO_NO_SESSION"
+  | "MARIMO_SHUTDOWN"
+  | "MARIMO_MALFORMED_QUERY"
+  | "MARIMO_KERNEL_STARTUP_ERROR"
+  | 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 },
-  context: { retryCount: number; maxRetries: number },
-): CloseDecision {
-  switch (event.reason) {
+export function classifyCloseEvent(event: { reason?: string }): CloseDecision {
+  switch (event.reason as CloseReason | undefined) {
     case "MARIMO_ALREADY_CONNECTED":
       return {
         kind: "terminal",
@@ -101,6 +110,15 @@ export function classifyCloseEvent(
         },
         closeTransport: true,
       };
+    case TRANSPORT_EXHAUSTED_REASON:
+      return {
+        kind: "gave-up",
+        status: {
+          state: WebSocketState.CLOSED,
+          code: WebSocketClosedReason.KERNEL_DISCONNECTED,
+          reason: "kernel not found",
+        },
+      };
     case "MARIMO_WRONG_KERNEL_ID":
     case "MARIMO_NO_FILE_KEY":
     case "MARIMO_NO_SESSION_ID":
@@ -144,18 +162,7 @@ export function classifyCloseEvent(
         logNever(event.reason as never);
       }
   }
-  // partysocket stops retrying silently once `maxRetries` is hit; surface
-  // CLOSED so callers can detect the give-up.
-  if (context.retryCount >= context.maxRetries) {
-    return {
-      kind: "gave-up",
-      status: {
-        state: WebSocketState.CLOSED,
-        code: WebSocketClosedReason.KERNEL_DISCONNECTED,
-        reason: "kernel not found",
-      },
-    };
-  }
+
   return {
     kind: "retry",
     status: { state: WebSocketState.CONNECTING },
@@ -440,7 +447,7 @@ export function useMarimoKernelConnection(opts: {
     }
     shouldTryReconnecting.current = true;
     setConnection({ state: WebSocketState.CONNECTING });
-    const healthy = await runtimeManager.isHealthy();
+    const healthy = await runtimeManager.reconcileFromHealth();
     if (!healthy) {
       shouldTryReconnecting.current = false;
       setConnection({
@@ -512,10 +519,7 @@ export function useMarimoKernelConnection(opts: {
      */
     onClose: (e) => {
       Logger.warn("WebSocket closed", e.code, e.reason);
-      const decision = classifyCloseEvent(e, {
-        retryCount: ws.retryCount,
-        maxRetries: MAX_RETRIES,
-      });
+      const decision = classifyCloseEvent(e);
       setConnection(decision.status);
       if (decision.kind === "terminal" && decision.closeTransport) {
         ws.close(); // close to prevent reconnecting

package/src/core/websocket/useWebSocket.tsx

@@ -1,11 +1,11 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import ReconnectingWebSocket from "partysocket/ws";
 import { useEffect, useState } from "react";
 import { Logger } from "@/utils/Logger";
 import { createPyodideConnection } from "../wasm/bridge";
 import { isWasm } from "../wasm/utils";
 import { BasicTransport } from "./transports/basic";
+import { WsTransport } from "./transports/ws";
 import type { IConnectionTransport } from "./transports/transport";
 
 interface UseConnectionTransportOptions {
@@ -18,10 +18,6 @@ interface UseConnectionTransportOptions {
   onError: (event: WebSocketEventMap["error"]) => void;
 }
 
-// Per-`reconnect()` retry budget for partysocket. After exhaustion, partysocket
-// stops silently; treat `retryCount >= MAX_RETRIES` as the give-up signal.
-export const MAX_RETRIES = 10;
-
 function createConnectionTransport(
   options: Pick<UseConnectionTransportOptions, "url" | "static">,
 ): IConnectionTransport {
@@ -31,19 +27,8 @@ function createConnectionTransport(
   if (isWasm()) {
     return createPyodideConnection();
   }
-  // Create a connection transport using the ReconnectingWebSocket from partysocket
-  // This handles reconnecting when the connection is lost.
-  const urlProvider = options.url; // We don't call the URL provider now since it may change (i.e. if the runtime redirects)
-  // Cast needed: ReconnectingWebSocket types readyState as `number`
-  // but IConnectionTransport expects `0 | 1 | 2 | 3`
-  return new ReconnectingWebSocket(urlProvider, undefined, {
-    maxRetries: MAX_RETRIES,
-    debug: false,
-    startClosed: true,
-    // long timeout -- the server can become slow when many notebooks
-    // are open.
-    connectionTimeout: 10_000,
-  }) as unknown as IConnectionTransport;
+  // urlProvider is passed lazily; it may change after a runtime redirect.
+  return new WsTransport(options.url);
 }
 
 /**