@marimo-team/frontend 0.25.3-dev81 → 0.25.3-dev82

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

@@ -1,320 +0,0 @@
-/* Copyright 2026 Marimo. All rights reserved. */
-
-import { Logger } from "@/utils/Logger";
-import {
-  type ConnectionEvent,
-  ConnectionSubscriptions,
-  type ConnectionTransportCallback,
-  type IConnectionTransport,
-} from "./transport";
-import { MAX_RETRIES, TRANSPORT_EXHAUSTED_REASON } from "./ws";
-
-// Mirrors WsTransport/partysocket retry behavior.
-const CONNECTION_TIMEOUT_MS = 10_000;
-const MIN_RETRY_DELAY_MS = 1000;
-const MAX_RETRY_DELAY_MS = 10_000;
-const RETRY_GROWTH_FACTOR = 1.3;
-
-// Stream ended without a server close event (transient disconnect); matches
-// the WebSocket abnormal-closure code so close handling stays uniform.
-const ABNORMAL_CLOSURE = 1006;
-
-interface ServerClose {
-  code: number;
-  reason: string;
-}
-
-/**
- * A kernel-connection transport over server-sent events, used when
- * `server.transport` is `"sse"` (experimental).
- *
- * The kernel connection only carries messages server to client (control
- * requests go over HTTP POST), so an SSE stream can fully replace the
- * WebSocket. The server mirrors WebSocket close frames with a `close` SSE
- * event carrying `{code, reason}`; a stream that ends without one is a
- * transient disconnect.
- *
- * Uses `fetch` + stream parsing rather than `EventSource`: we need to
- * surface server close reasons, send auth headers, and match partysocket's
- * retry semantics (per-`reconnect()` retry budget, exhaustion surfaced via
- * `TRANSPORT_EXHAUSTED_REASON`), none of which `EventSource` supports.
- *
- * Like `WsTransport`, starts closed; the first `reconnect()` connects.
- * `close()` stops retrying and does not dispatch a close event.
- */
-export class SseTransport implements IConnectionTransport {
-  private subscriptions = new ConnectionSubscriptions();
-  private urlProvider: () => string;
-  private headersProvider: () => Record<string, string>;
-  private abortController: AbortController | null = null;
-  private retryTimer: ReturnType<typeof setTimeout> | null = null;
-  private retryCount = 0;
-  private userClosed = false;
-  private state: WebSocket["readyState"] = WebSocket.CLOSED;
-
-  public constructor(
-    urlProvider: () => string,
-    headersProvider: () => Record<string, string> = () => ({}),
-  ) {
-    this.urlProvider = urlProvider;
-    this.headersProvider = headersProvider;
-  }
-
-  public get readyState(): WebSocket["readyState"] {
-    return this.state;
-  }
-
-  public reconnect(_code?: number, _reason?: string): void {
-    this.userClosed = false;
-    this.retryCount = 0;
-    this.clearRetryTimer();
-    this.abortController?.abort();
-    this.connect();
-  }
-
-  public close(): void {
-    this.userClosed = true;
-    this.clearRetryTimer();
-    this.state = WebSocket.CLOSED;
-    this.abortController?.abort();
-    this.abortController = null;
-  }
-
-  public send(_data: string | ArrayBuffer | Blob | ArrayBufferView): void {
-    // The kernel connection is receive-only; all client requests go over
-    // HTTP POST endpoints.
-    Logger.warn("SseTransport does not support send(); dropping message");
-  }
-
-  public addEventListener<T extends ConnectionEvent>(
-    event: T,
-    callback: ConnectionTransportCallback<T>,
-  ): void {
-    this.subscriptions.addSubscription(
-      event,
-      callback as ConnectionTransportCallback<ConnectionEvent>,
-    );
-  }
-
-  public removeEventListener<T extends ConnectionEvent>(
-    event: T,
-    callback: ConnectionTransportCallback<T>,
-  ): void {
-    this.subscriptions.removeSubscription(
-      event,
-      callback as ConnectionTransportCallback<ConnectionEvent>,
-    );
-  }
-
-  private connect(): void {
-    this.state = WebSocket.CONNECTING;
-    void this.run();
-  }
-
-  private async run(): Promise<void> {
-    const controller = new AbortController();
-    this.abortController = controller;
-    // Long timeout — the server can become slow when many notebooks are open.
-    const connectTimeout = setTimeout(
-      () => controller.abort(),
-      CONNECTION_TIMEOUT_MS,
-    );
-
-    let serverClose: ServerClose | null = null;
-    try {
-      const response = await fetch(this.urlProvider(), {
-        headers: {
-          ...this.headersProvider(),
-          Accept: "text/event-stream",
-        },
-        signal: controller.signal,
-      });
-      clearTimeout(connectTimeout);
-      if (!response.ok || response.body === null) {
-        throw new Error(`Unexpected SSE response: ${response.status}`);
-      }
-
-      this.state = WebSocket.OPEN;
-      this.subscriptions.notify("open", new Event("open"));
-
-      serverClose = await this.readEvents(response.body);
-    } catch (error) {
-      clearTimeout(connectTimeout);
-      if (this.isSuperseded(controller)) {
-        return;
-      }
-      Logger.warn("SSE connection failed", error);
-      this.subscriptions.notify("error", new Event("error"));
-    }
-
-    if (this.isSuperseded(controller)) {
-      return;
-    }
-    controller.abort();
-    this.state = WebSocket.CLOSED;
-
-    // The retry budget counts consecutive attempts without a healthy
-    // stream (readEvents resets it on the first message), so a server
-    // that accepts and immediately closes still exhausts the budget.
-    this.retryCount += 1;
-    const exhausted = this.retryCount >= MAX_RETRIES;
-    this.subscriptions.notify(
-      "close",
-      new CloseEvent("close", {
-        code: serverClose?.code ?? ABNORMAL_CLOSURE,
-        // On exhaustion the reason is rewritten so the consumer stops
-        // retrying, matching WsTransport.
-        reason: exhausted
-          ? TRANSPORT_EXHAUSTED_REASON
-          : (serverClose?.reason ?? ""),
-        wasClean: serverClose !== null,
-      }),
-    );
-    // The consumer's close handler runs synchronously inside notify and
-    // may have called reconnect() (starting a fresh connection) or
-    // close(); re-check before scheduling our own retry so we never
-    // stack a second connection on top of theirs.
-    if (exhausted || this.isSuperseded(controller)) {
-      return;
-    }
-    this.scheduleRetry();
-  }
-
-  /** Whether this run was replaced by a newer reconnect() or close(). */
-  private isSuperseded(controller: AbortController): boolean {
-    return this.userClosed || this.abortController !== controller;
-  }
-
-  private scheduleRetry(): void {
-    this.clearRetryTimer();
-    this.state = WebSocket.CONNECTING;
-    const delay = Math.min(
-      MIN_RETRY_DELAY_MS * RETRY_GROWTH_FACTOR ** this.retryCount,
-      MAX_RETRY_DELAY_MS,
-    );
-    this.retryTimer = setTimeout(() => {
-      this.retryTimer = null;
-      this.connect();
-    }, delay);
-  }
-
-  private clearRetryTimer(): void {
-    if (this.retryTimer !== null) {
-      clearTimeout(this.retryTimer);
-      this.retryTimer = null;
-    }
-  }
-
-  /**
-   * Read and dispatch SSE events until the stream ends.
-   *
-   * Returns the server close payload if a `close` event arrives, or null
-   * if the stream ends without one.
-   */
-  private async readEvents(
-    body: ReadableStream<Uint8Array>,
-  ): Promise<ServerClose | null> {
-    const reader = body.getReader();
-    const decoder = new TextDecoder();
-    let buffer = "";
-    let eventName: string | null = null;
-    let eventId: string | null = null;
-    let dataLines: string[] = [];
-
-    const dispatch = (): ServerClose | null => {
-      if (dataLines.length === 0) {
-        eventName = null;
-        eventId = null;
-        return null;
-      }
-      const data = dataLines.join("\n");
-      const name = eventName;
-      const id = eventId;
-      eventName = null;
-      eventId = null;
-      dataLines = [];
-      if (name === "close") {
-        return parseServerClose(data);
-      }
-      // A message means the stream is healthy: reset the retry budget
-      // (rather than on open, so a server that accepts and immediately
-      // closes cannot retry forever).
-      this.retryCount = 0;
-      // Normalize SSE fields into the same envelope as host WebSockets.
-      // Leave the payload text intact for the client's JSON decoder.
-      const envelope =
-        name === null
-          ? data
-          : `{"event":${JSON.stringify(name)},"data":${data}${id === null ? "" : `,"id":${JSON.stringify(id)}`}}`;
-      this.subscriptions.notify(
-        "message",
-        new MessageEvent("message", { data: envelope }),
-      );
-      return null;
-    };
-
-    while (true) {
-      const { done, value } = await reader.read();
-      if (done) {
-        return null;
-      }
-      buffer += decoder.decode(value, { stream: true });
-
-      // Consume complete lines with a cursor; a single tail slice per
-      // read keeps parsing linear even for many buffered lines.
-      let cursor = 0;
-      let newlineIndex = buffer.indexOf("\n", cursor);
-      while (newlineIndex !== -1) {
-        let line = buffer.slice(cursor, newlineIndex);
-        cursor = newlineIndex + 1;
-        newlineIndex = buffer.indexOf("\n", cursor);
-        if (line.endsWith("\r")) {
-          line = line.slice(0, -1);
-        }
-
-        if (line === "") {
-          const serverClose = dispatch();
-          if (serverClose !== null) {
-            return serverClose;
-          }
-          continue;
-        }
-        if (line.startsWith(":")) {
-          // Comment (keep-alive)
-          continue;
-        }
-
-        const colonIndex = line.indexOf(":");
-        const field = colonIndex === -1 ? line : line.slice(0, colonIndex);
-        let fieldValue = colonIndex === -1 ? "" : line.slice(colonIndex + 1);
-        if (fieldValue.startsWith(" ")) {
-          fieldValue = fieldValue.slice(1);
-        }
-        if (field === "data") {
-          dataLines.push(fieldValue);
-        } else if (field === "event") {
-          eventName = fieldValue;
-        } else if (field === "id") {
-          eventId = fieldValue;
-        }
-        // The host client owns the cursor; `retry` does not override our policy.
-      }
-      if (cursor > 0) {
-        buffer = buffer.slice(cursor);
-      }
-    }
-  }
-}
-
-function parseServerClose(data: string): ServerClose {
-  try {
-    const parsed = JSON.parse(data) as Partial<ServerClose>;
-    return {
-      code: typeof parsed.code === "number" ? parsed.code : 1000,
-      reason: typeof parsed.reason === "string" ? parsed.reason : "",
-    };
-  } catch {
-    Logger.warn("Failed to parse SSE close event", data);
-    return { code: 1000, reason: "" };
-  }
-}

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

@@ -1,50 +0,0 @@
-/* Copyright 2026 Marimo. All rights reserved. */
-
-import type {
-  NotebookTransport,
-  ConnectionEvent,
-  ConnectionTransportCallback,
-  IConnectionTransportMap,
-} from "../../host/connection";
-export type {
-  ConnectionEvent,
-  ConnectionTransportCallback,
-  IConnectionTransportMap,
-} from "../../host/connection";
-
-export interface IConnectionTransport extends NotebookTransport {
-  send(data: string | ArrayBuffer | Blob | ArrayBufferView<ArrayBuffer>): void;
-}
-
-export class ConnectionSubscriptions {
-  private subscriptions = new Map<
-    ConnectionEvent,
-    Set<ConnectionTransportCallback<ConnectionEvent>>
-  >();
-
-  public addSubscription(
-    event: ConnectionEvent,
-    callback: ConnectionTransportCallback<ConnectionEvent>,
-  ): void {
-    if (!this.subscriptions.has(event)) {
-      this.subscriptions.set(event, new Set());
-    }
-    this.subscriptions.get(event)?.add(callback);
-  }
-
-  public removeSubscription(
-    event: ConnectionEvent,
-    callback: ConnectionTransportCallback<ConnectionEvent>,
-  ): void {
-    this.subscriptions.get(event)?.delete(callback);
-  }
-
-  public notify(
-    event: ConnectionEvent,
-    data: IConnectionTransportMap[ConnectionEvent],
-  ): void {
-    for (const callback of this.subscriptions.get(event) ?? []) {
-      callback(data);
-    }
-  }
-}

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

@@ -1,98 +0,0 @@
-/* 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">
-  >();
-
-  public 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,
-    });
-  }
-
-  public get readyState(): WebSocket["readyState"] {
-    return this.inner.readyState as WebSocket["readyState"];
-  }
-
-  public reconnect(code?: number, reason?: string): void {
-    this.inner.reconnect(code, reason);
-  }
-
-  public close(): void {
-    this.inner.close();
-  }
-
-  public send(
-    data: string | ArrayBuffer | Blob | ArrayBufferView<ArrayBuffer>,
-  ): void {
-    this.inner.send(data);
-  }
-
-  public 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);
-  }
-
-  public 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);
-  }
-}