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

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

@@ -0,0 +1,70 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import type { CloseDecision } from "../host/connection";
+import { Logger } from "@/utils/Logger";
+import { TRANSPORT_EXHAUSTED_REASON } from "./transports/ws";
+import { ConnectionState, WebSocketClosedReason } from "./types";
+
+// Keep server reasons in sync with ws_endpoint.py and ws/*.py.
+// SSE carries the same reasons in close events. Unclassified reasons retry.
+export function classifyCloseEvent(event: { reason?: string }): CloseDecision {
+  switch (event.reason) {
+    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",
+        },
+      };
+    case "MARIMO_UNAUTHORIZED":
+      return {
+        kind: "terminal",
+        status: {
+          state: ConnectionState.CLOSED,
+          code: WebSocketClosedReason.KERNEL_DISCONNECTED,
+          reason: "not authorized",
+        },
+      };
+    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",
+        },
+      };
+    case "MARIMO_KERNEL_STARTUP_ERROR":
+      return {
+        kind: "terminal",
+        status: {
+          state: ConnectionState.CLOSED,
+          code: WebSocketClosedReason.KERNEL_STARTUP_ERROR,
+          reason: "Failed to start kernel sandbox",
+        },
+      };
+    default:
+      if (event.reason) {
+        Logger.warn("Unrecognized connection close reason", event.reason);
+      }
+  }
+
+  return {
+    kind: "retry",
+    status: { state: ConnectionState.CONNECTING },
+  };
+}

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

@@ -108,8 +108,6 @@ describe("SseTransport", () => {
     transport.reconnect();
     await vi.waitFor(() => expect(opens).toHaveLength(1));
 
-    // Stop retrying when the consumer treats the close as terminal,
-    // mirroring useMarimoKernelConnection's terminal path.
     transport.addEventListener("close", () => transport.close());
     push('event: close\ndata: {"code": 1000, "reason": "MARIMO_SHUTDOWN"}\n\n');
     await vi.waitFor(() => expect(closes).toHaveLength(1));
@@ -146,8 +144,7 @@ describe("SseTransport", () => {
     fetchMock
       .mockResolvedValueOnce(sseResponse(first.stream))
       .mockResolvedValueOnce(sseResponse(second.stream));
-    // Mirror useMarimoKernelConnection's retry path: reconnect()
-    // synchronously from inside the close handler.
+    // Reconnect synchronously from the close handler.
     transport.addEventListener("close", () => transport.reconnect());
 
     transport.reconnect();

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

@@ -0,0 +1,31 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { createPyodideConnection } from "../../wasm/bridge";
+import { isWasm } from "../../wasm/utils";
+import { BasicTransport } from "./basic";
+import { SseTransport } from "./sse";
+import { WsTransport } from "./ws";
+import type { IConnectionTransport } from "./transport";
+
+export interface ConnectionTransportOptions {
+  url: () => string;
+  static: boolean;
+  transportType: "websocket" | "sse";
+  headers: () => Record<string, string>;
+}
+
+export function createConnectionTransport(
+  options: ConnectionTransportOptions,
+): IConnectionTransport {
+  if (options.static) {
+    return BasicTransport.empty();
+  }
+  if (isWasm()) {
+    return createPyodideConnection();
+  }
+  // Read the URL lazily to follow runtime redirects.
+  if (options.transportType === "sse") {
+    return new SseTransport(options.url, options.headers);
+  }
+  return new WsTransport(options.url);
+}

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

@@ -1,30 +1,19 @@
 /* Copyright 2026 Marimo. All rights reserved. */
-export type ConnectionEvent = "open" | "close" | "message" | "error";
 
-export interface IConnectionTransportMap {
-  open: WebSocketEventMap["open"];
-  message: WebSocketEventMap["message"];
-  close: WebSocketEventMap["close"];
-  error: WebSocketEventMap["error"];
-}
-
-export type ConnectionTransportCallback<T extends ConnectionEvent> = (
-  data: IConnectionTransportMap[T],
-) => void;
+import type {
+  NotebookTransport,
+  ConnectionEvent,
+  ConnectionTransportCallback,
+  IConnectionTransportMap,
+} from "../../host/connection";
+export type {
+  ConnectionEvent,
+  ConnectionTransportCallback,
+  IConnectionTransportMap,
+} from "../../host/connection";
 
-export interface IConnectionTransport {
-  reconnect(code?: number | undefined, reason?: string | undefined): void;
-  close(): void;
+export interface IConnectionTransport extends NotebookTransport {
   send(data: string | ArrayBuffer | Blob | ArrayBufferView<ArrayBuffer>): void;
-  addEventListener<T extends ConnectionEvent>(
-    event: T,
-    callback: ConnectionTransportCallback<T>,
-  ): void;
-  removeEventListener<T extends ConnectionEvent>(
-    event: T,
-    callback: ConnectionTransportCallback<T>,
-  ): void;
-  readyState: WebSocket["readyState"];
 }
 
 export class ConnectionSubscriptions {

package/src/core/websocket/types.ts

@@ -1,47 +1,17 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
+import type { ConnectionStatus as HostConnectionStatus } from "../host/connection";
 import type { NotificationMessageData } from "../kernel/messages";
 
+export {
+  ConnectionState,
+  ConnectionClosedReason as WebSocketClosedReason,
+} from "../host/connection";
+
 export type ConnectionPhase =
   | NotificationMessageData<"startup-progress">["phase"]
   | "reconnecting";
 
-export const ConnectionState = {
-  NOT_STARTED: "NOT_STARTED",
-  CONNECTING: "CONNECTING",
-  OPEN: "OPEN",
-  CLOSING: "CLOSING",
-  CLOSED: "CLOSED",
-} as const;
-
-export type ConnectionState =
-  (typeof ConnectionState)[keyof typeof ConnectionState];
-
-export const WebSocketClosedReason = {
-  KERNEL_DISCONNECTED: "KERNEL_DISCONNECTED",
-  KERNEL_STARTUP_ERROR: "KERNEL_STARTUP_ERROR",
-} as const;
-
-export type WebSocketClosedReason =
-  (typeof WebSocketClosedReason)[keyof typeof WebSocketClosedReason];
-
-export type ConnectionStatus =
-  | {
-      state: typeof ConnectionState.CLOSED;
-      code: WebSocketClosedReason;
-      /**
-       * Human-readable reason for closing the connection.
-       */
-      reason: string;
-      phase?: ConnectionPhase;
-    }
-  | {
-      state: typeof ConnectionState.CONNECTING;
-      phase?: ConnectionPhase;
-    }
-  | {
-      state:
-        | typeof ConnectionState.OPEN
-        | typeof ConnectionState.CLOSING
-        | typeof ConnectionState.NOT_STARTED;
-    };
+export type ConnectionStatus = HostConnectionStatus & {
+  phase?: ConnectionPhase;
+};

package/src/core/websocket/useNotebookClient.ts

@@ -0,0 +1,80 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { useStore } from "jotai";
+import { useEffect, useMemo } from "react";
+import useEvent from "react-use-event-hook";
+import { jsonParseWithSpecialChar } from "@/utils/json/json-parser";
+import type {
+  NotificationMessage,
+  NotificationPayload,
+} from "../kernel/messages";
+import {
+  createNotebookClient,
+  type NotebookClient,
+  type NotebookConnection,
+} from "../host/notebook-client";
+import { classifyCloseEvent } from "./connection-state";
+import {
+  createConnectionTransport,
+  type ConnectionTransportOptions,
+} from "./transports/create";
+
+type Options = ConnectionTransportOptions & {
+  sessionId: string;
+  messageFormat: NotebookConnection["messageFormat"];
+  waitUntilReady: () => Promise<void>;
+  checkHealth: () => Promise<boolean>;
+  onNotification: (notification: NotificationPayload) => void;
+  onMessageError: NotebookConnection["onMessageError"];
+};
+
+export function useNotebookClient(options: Options): NotebookClient {
+  const store = useStore();
+  const { sessionId, static: isStatic, transportType, messageFormat } = options;
+  const controller = useMemo(() => createNotebookClient(store), [store]);
+  const url = useEvent(options.url);
+  const headers = useEvent(options.headers);
+  const waitUntilReady = useEvent(options.waitUntilReady);
+  const checkHealth = useEvent(options.checkHealth);
+  const onNotification = useEvent(options.onNotification);
+  const onMessageError = useEvent(options.onMessageError);
+
+  useEffect(() => {
+    const lifetime = new AbortController();
+    controller.connect({
+      createTransport: (url) =>
+        createConnectionTransport({
+          url,
+          headers,
+          static: isStatic,
+          transportType,
+        }),
+      url,
+      messageFormat,
+      signal: lifetime.signal,
+      waitUntilReady,
+      checkHealth,
+      onMessageError,
+      // The host validates the envelope, not the native notification contents.
+      onNotification: (notification) =>
+        onNotification({ data: notification as NotificationMessage }),
+      decodeJson: jsonParseWithSpecialChar,
+      classifyClose: classifyCloseEvent,
+    });
+    return () => lifetime.abort();
+  }, [
+    controller,
+    sessionId,
+    isStatic,
+    transportType,
+    messageFormat,
+    url,
+    headers,
+    waitUntilReady,
+    checkHealth,
+    onNotification,
+    onMessageError,
+  ]);
+
+  return controller.client;
+}