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

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);
 }
 
 /**