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

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

@@ -1,125 +0,0 @@
-/* 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 {
-    public retryCount = 0;
-    public readyState = WebSocket.CONNECTING;
-    public constructor() {
-      innerListeners = { open: [], close: [], message: [], error: [] };
-    }
-    public addEventListener(event: string, cb: (e: unknown) => void) {
-      innerListeners[event].push(cb);
-    }
-    public removeEventListener(event: string, cb: (e: unknown) => void) {
-      innerListeners[event] = innerListeners[event].filter((c) => c !== cb);
-    }
-    public reconnect() {}
-    public close() {}
-    public 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

@@ -1,102 +0,0 @@
-/* Copyright 2026 Marimo. All rights reserved. */
-
-import {
-  type ConnectionEvent,
-  ConnectionSubscriptions,
-  type ConnectionTransportCallback,
-  type IConnectionTransport,
-} from "./transport";
-
-type MessageProducer = (callback: (message: MessageEvent) => void) => void;
-
-export class BasicTransport implements IConnectionTransport {
-  private subscriptions = new ConnectionSubscriptions();
-  private producer?: MessageProducer;
-
-  public static withProducerCallback(
-    producer: MessageProducer,
-  ): IConnectionTransport {
-    return new BasicTransport(producer);
-  }
-
-  public static empty(): IConnectionTransport {
-    return new BasicTransport();
-  }
-
-  private constructor(producer?: MessageProducer) {
-    this.producer = producer;
-  }
-
-  private startProducer() {
-    if (this.producer) {
-      this.producer((message) => {
-        this.subscriptions.notify("message", message);
-      });
-    }
-  }
-
-  private connect(): Promise<void> {
-    return new Promise<void>((resolve) => setTimeout(resolve, 0)).then(() => {
-      this.subscriptions.notify("open", new Event("open"));
-    });
-  }
-
-  public get readyState(): WebSocket["readyState"] {
-    return WebSocket.OPEN;
-  }
-
-  public reconnect(
-    _code?: number | undefined,
-    _reason?: string | undefined,
-  ): void {
-    this.close();
-    this.connect();
-    return;
-  }
-
-  public close(): void {
-    this.subscriptions.notify("close", new Event("close"));
-  }
-
-  public send(
-    data: string | ArrayBuffer | Blob | ArrayBufferView,
-  ): Promise<void> {
-    this.subscriptions.notify(
-      "message",
-      new MessageEvent("message", {
-        data,
-      }),
-    );
-    return Promise.resolve();
-  }
-
-  public addEventListener<T extends ConnectionEvent>(
-    event: T,
-    callback: ConnectionTransportCallback<T>,
-  ): void {
-    this.subscriptions.addSubscription(
-      event,
-      callback as ConnectionTransportCallback<ConnectionEvent>,
-    );
-
-    // Call open right away
-    if (event === "open") {
-      (callback as ConnectionTransportCallback<"open">)(new Event("open"));
-    }
-
-    // Start the producer when we have one consumer
-    if (event === "message") {
-      this.startProducer();
-    }
-  }
-
-  public removeEventListener<T extends ConnectionEvent>(
-    event: T,
-    callback: ConnectionTransportCallback<T>,
-  ): void {
-    this.subscriptions.removeSubscription(
-      event,
-      callback as ConnectionTransportCallback<ConnectionEvent>,
-    );
-  }
-}

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

@@ -1,31 +0,0 @@
-/* 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/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);
-    }
-  }
-}