@marimo-team/islands 0.23.14-dev6 → 0.23.14-dev61

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

@@ -0,0 +1,338 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import { SseTransport } from "../sse";
+import { MAX_RETRIES, TRANSPORT_EXHAUSTED_REASON } from "../ws";
+
+const encoder = new TextEncoder();
+
+function createControllableStream() {
+  let controller!: ReadableStreamDefaultController<Uint8Array>;
+  const stream = new ReadableStream<Uint8Array>({
+    start(c) {
+      controller = c;
+    },
+  });
+  return {
+    stream,
+    push: (text: string) => controller.enqueue(encoder.encode(text)),
+    end: () => controller.close(),
+  };
+}
+
+function sseResponse(body: ReadableStream<Uint8Array> | null, status = 200) {
+  return {
+    ok: status >= 200 && status < 300,
+    status,
+    body,
+  } as Response;
+}
+
+describe("SseTransport", () => {
+  let transport: SseTransport;
+  let fetchMock: ReturnType<typeof vi.fn>;
+  let opens: Event[];
+  let messages: MessageEvent[];
+  let closes: CloseEvent[];
+  let errors: Event[];
+
+  beforeEach(() => {
+    fetchMock = vi.fn();
+    vi.stubGlobal("fetch", fetchMock);
+    transport = new SseTransport(() => "http://example.invalid/sse");
+    opens = [];
+    messages = [];
+    closes = [];
+    errors = [];
+    transport.addEventListener("open", (e) => opens.push(e));
+    transport.addEventListener("message", (e) => messages.push(e));
+    transport.addEventListener("close", (e) => closes.push(e));
+    transport.addEventListener("error", (e) => errors.push(e));
+  });
+
+  afterEach(() => {
+    transport.close();
+    vi.unstubAllGlobals();
+    vi.useRealTimers();
+  });
+
+  it("starts closed and connects on reconnect()", async () => {
+    expect(transport.readyState).toBe(WebSocket.CLOSED);
+    const { stream } = createControllableStream();
+    fetchMock.mockResolvedValueOnce(sseResponse(stream));
+
+    transport.reconnect();
+    expect(transport.readyState).toBe(WebSocket.CONNECTING);
+    await vi.waitFor(() => expect(opens).toHaveLength(1));
+    expect(transport.readyState).toBe(WebSocket.OPEN);
+    expect(fetchMock).toHaveBeenCalledWith(
+      "http://example.invalid/sse",
+      expect.objectContaining({
+        headers: { Accept: "text/event-stream" },
+      }),
+    );
+  });
+
+  it("dispatches message events, joining multi-line data", async () => {
+    const { stream, push } = createControllableStream();
+    fetchMock.mockResolvedValueOnce(sseResponse(stream));
+    transport.reconnect();
+    await vi.waitFor(() => expect(opens).toHaveLength(1));
+
+    push('data: {"op": "alert"}\n\n');
+    await vi.waitFor(() => expect(messages).toHaveLength(1));
+    expect(messages[0].data).toBe('{"op": "alert"}');
+
+    // Multi-line data and records split across chunks
+    push("data: line1\nda");
+    push("ta: line2\n\n");
+    await vi.waitFor(() => expect(messages).toHaveLength(2));
+    expect(messages[1].data).toBe("line1\nline2");
+  });
+
+  it("ignores keep-alive comments", async () => {
+    const { stream, push } = createControllableStream();
+    fetchMock.mockResolvedValueOnce(sseResponse(stream));
+    transport.reconnect();
+    await vi.waitFor(() => expect(opens).toHaveLength(1));
+
+    push(": keep-alive\n\n");
+    push("data: after\n\n");
+    await vi.waitFor(() => expect(messages).toHaveLength(1));
+    expect(messages[0].data).toBe("after");
+  });
+
+  it("surfaces a server close event with its code and reason", async () => {
+    const { stream, push } = createControllableStream();
+    fetchMock.mockResolvedValueOnce(sseResponse(stream));
+    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));
+    expect(closes[0].code).toBe(1000);
+    expect(closes[0].reason).toBe("MARIMO_SHUTDOWN");
+    expect(closes[0].wasClean).toBe(true);
+    expect(transport.readyState).toBe(WebSocket.CLOSED);
+    expect(fetchMock).toHaveBeenCalledTimes(1);
+  });
+
+  it("emits an empty-reason close and retries when the stream ends", async () => {
+    vi.useFakeTimers();
+    const first = createControllableStream();
+    fetchMock.mockResolvedValueOnce(sseResponse(first.stream));
+    transport.reconnect();
+    await vi.waitFor(() => expect(opens).toHaveLength(1));
+
+    first.end();
+    await vi.waitFor(() => expect(closes).toHaveLength(1));
+    expect(closes[0].reason).toBe("");
+    expect(closes[0].wasClean).toBe(false);
+
+    // A retry is scheduled and reconnects
+    const second = createControllableStream();
+    fetchMock.mockResolvedValueOnce(sseResponse(second.stream));
+    await vi.advanceTimersByTimeAsync(10_000);
+    await vi.waitFor(() => expect(opens).toHaveLength(2));
+  });
+
+  it("does not stack its own retry when the consumer reconnects in onClose", async () => {
+    vi.useFakeTimers();
+    const first = createControllableStream();
+    const second = createControllableStream();
+    fetchMock
+      .mockResolvedValueOnce(sseResponse(first.stream))
+      .mockResolvedValueOnce(sseResponse(second.stream));
+    // Mirror useMarimoKernelConnection's retry path: reconnect()
+    // synchronously from inside the close handler.
+    transport.addEventListener("close", () => transport.reconnect());
+
+    transport.reconnect();
+    await vi.waitFor(() => expect(opens).toHaveLength(1));
+
+    first.end();
+    await vi.waitFor(() => expect(opens).toHaveLength(2));
+
+    // The superseded run must not schedule a second connection on top of
+    // the consumer's reconnect.
+    await vi.advanceTimersByTimeAsync(60_000);
+    expect(fetchMock).toHaveBeenCalledTimes(2);
+
+    // The consumer-driven connection is healthy
+    second.push("data: alive\n\n");
+    await vi.waitFor(() => expect(messages).toHaveLength(1));
+    expect(transport.readyState).toBe(WebSocket.OPEN);
+  });
+
+  it("exhausts the retry budget when the server closes every connection", async () => {
+    vi.useFakeTimers();
+    // A server that accepts, immediately sends a close event, and never
+    // delivers a message must not be retried forever.
+    fetchMock.mockImplementation(() => {
+      const { stream, push } = createControllableStream();
+      push('event: close\ndata: {"code": 3000, "reason": "MARIMO_X"}\n\n');
+      return Promise.resolve(sseResponse(stream));
+    });
+
+    transport.reconnect();
+    for (let i = 0; i < MAX_RETRIES; i++) {
+      await vi.advanceTimersByTimeAsync(15_000);
+    }
+
+    expect(closes.length).toBe(MAX_RETRIES);
+    expect(closes.at(-1)?.reason).toBe(TRANSPORT_EXHAUSTED_REASON);
+    const fetchCallsWhenExhausted = fetchMock.mock.calls.length;
+    await vi.advanceTimersByTimeAsync(60_000);
+    expect(fetchMock.mock.calls.length).toBe(fetchCallsWhenExhausted);
+  });
+
+  it("a delivered message resets the retry budget", async () => {
+    vi.useFakeTimers();
+    const first = createControllableStream();
+    fetchMock.mockResolvedValueOnce(sseResponse(first.stream));
+    transport.reconnect();
+    await vi.waitFor(() => expect(opens).toHaveLength(1));
+    first.push("data: healthy\n\n");
+    await vi.waitFor(() => expect(messages).toHaveLength(1));
+
+    first.end();
+    await vi.waitFor(() => expect(closes).toHaveLength(1));
+    // Budget was reset by the message, so this close is retry #1, not
+    // an exhaustion.
+    expect(closes[0].reason).toBe("");
+  });
+
+  it("gives up with TRANSPORT_EXHAUSTED_REASON after MAX_RETRIES failures", async () => {
+    vi.useFakeTimers();
+    fetchMock.mockRejectedValue(new Error("connection refused"));
+
+    transport.reconnect();
+    for (let i = 0; i < MAX_RETRIES; i++) {
+      await vi.advanceTimersByTimeAsync(15_000);
+    }
+
+    expect(closes.length).toBe(MAX_RETRIES);
+    expect(closes.at(-1)?.reason).toBe(TRANSPORT_EXHAUSTED_REASON);
+    expect(errors.length).toBe(MAX_RETRIES);
+    const fetchCallsWhenExhausted = fetchMock.mock.calls.length;
+
+    // No further retries once exhausted
+    await vi.advanceTimersByTimeAsync(60_000);
+    expect(fetchMock.mock.calls.length).toBe(fetchCallsWhenExhausted);
+
+    // A manual reconnect() resets the budget
+    const { stream } = createControllableStream();
+    fetchMock.mockResolvedValueOnce(sseResponse(stream));
+    transport.reconnect();
+    await vi.waitFor(() => expect(opens).toHaveLength(1));
+  });
+
+  it("close() aborts the connection and cancels retries", async () => {
+    vi.useFakeTimers();
+    const { stream } = createControllableStream();
+    fetchMock.mockResolvedValueOnce(sseResponse(stream));
+    transport.reconnect();
+    await vi.waitFor(() => expect(opens).toHaveLength(1));
+
+    transport.close();
+    expect(transport.readyState).toBe(WebSocket.CLOSED);
+    // No close event is dispatched for a user-initiated close, and no
+    // retry is scheduled.
+    await vi.advanceTimersByTimeAsync(60_000);
+    expect(closes).toHaveLength(0);
+    expect(fetchMock).toHaveBeenCalledTimes(1);
+  });
+
+  it("treats a non-2xx response as a connection failure", async () => {
+    vi.useFakeTimers();
+    fetchMock.mockResolvedValueOnce(sseResponse(null, 500));
+    transport.reconnect();
+
+    await vi.waitFor(() => expect(closes).toHaveLength(1));
+    expect(opens).toHaveLength(0);
+    expect(errors).toHaveLength(1);
+    expect(closes[0].reason).toBe("");
+  });
+
+  it("sends provided headers (auth) with the request", async () => {
+    const withHeaders = new SseTransport(
+      () => "http://example.invalid/sse",
+      () => ({ Authorization: "Bearer token123" }),
+    );
+    const { stream } = createControllableStream();
+    fetchMock.mockResolvedValueOnce(sseResponse(stream));
+    withHeaders.reconnect();
+    await vi.waitFor(() => expect(fetchMock).toHaveBeenCalled());
+    withHeaders.close();
+
+    expect(fetchMock).toHaveBeenCalledWith(
+      "http://example.invalid/sse",
+      expect.objectContaining({
+        headers: {
+          Authorization: "Bearer token123",
+          Accept: "text/event-stream",
+        },
+      }),
+    );
+  });
+
+  it("aborts a connection attempt that exceeds the timeout", async () => {
+    vi.useFakeTimers();
+    // fetch that only settles when aborted
+    fetchMock.mockImplementation(
+      (_url: string, init: RequestInit) =>
+        new Promise((_, reject) => {
+          init.signal?.addEventListener("abort", () =>
+            reject(new Error("aborted")),
+          );
+        }),
+    );
+
+    transport.reconnect();
+    await vi.advanceTimersByTimeAsync(11_000);
+    expect(errors).toHaveLength(1);
+    expect(closes).toHaveLength(1);
+    expect(closes[0].reason).toBe("");
+  });
+
+  it("reconnect() mid-stream supersedes the old connection silently", async () => {
+    const first = createControllableStream();
+    const second = createControllableStream();
+    fetchMock
+      .mockResolvedValueOnce(sseResponse(first.stream))
+      .mockResolvedValueOnce(sseResponse(second.stream));
+
+    transport.reconnect();
+    await vi.waitFor(() => expect(opens).toHaveLength(1));
+
+    transport.reconnect();
+    await vi.waitFor(() => expect(opens).toHaveLength(2));
+
+    // The aborted first connection must not surface close/error events
+    second.push("data: alive\n\n");
+    await vi.waitFor(() => expect(messages).toHaveLength(1));
+    expect(closes).toHaveLength(0);
+    expect(errors).toHaveLength(0);
+  });
+
+  it("send() is a no-op", () => {
+    expect(() => transport.send("data")).not.toThrow();
+  });
+
+  it("dedupes repeated addEventListener calls", async () => {
+    const { stream, push } = createControllableStream();
+    fetchMock.mockResolvedValueOnce(sseResponse(stream));
+    const onMessage = (e: MessageEvent) => messages.push(e);
+    transport.addEventListener("message", onMessage);
+    transport.addEventListener("message", onMessage);
+    transport.reconnect();
+    await vi.waitFor(() => expect(opens).toHaveLength(1));
+
+    push("data: once\n\n");
+    // The beforeEach listener plus the deduped listener = 2 events
+    await vi.waitFor(() => expect(messages).toHaveLength(2));
+  });
+});

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

@@ -0,0 +1,308 @@
+/* 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;
+
+  constructor(
+    urlProvider: () => string,
+    headersProvider: () => Record<string, string> = () => ({}),
+  ) {
+    this.urlProvider = urlProvider;
+    this.headersProvider = headersProvider;
+  }
+
+  get readyState(): WebSocket["readyState"] {
+    return this.state;
+  }
+
+  reconnect(_code?: number, _reason?: string): void {
+    this.userClosed = false;
+    this.retryCount = 0;
+    this.clearRetryTimer();
+    this.abortController?.abort();
+    this.connect();
+  }
+
+  close(): void {
+    this.userClosed = true;
+    this.clearRetryTimer();
+    this.state = WebSocket.CLOSED;
+    this.abortController?.abort();
+    this.abortController = null;
+  }
+
+  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");
+  }
+
+  addEventListener<T extends ConnectionEvent>(
+    event: T,
+    callback: ConnectionTransportCallback<T>,
+  ): void {
+    this.subscriptions.addSubscription(
+      event,
+      callback as ConnectionTransportCallback<ConnectionEvent>,
+    );
+  }
+
+  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 dataLines: string[] = [];
+
+    const dispatch = (): ServerClose | null => {
+      if (dataLines.length === 0) {
+        eventName = null;
+        return null;
+      }
+      const data = dataLines.join("\n");
+      const name = eventName;
+      eventName = 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;
+      this.subscriptions.notify(
+        "message",
+        new MessageEvent("message", { data }),
+      );
+      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;
+        }
+        // `id` and `retry` fields are unused.
+      }
+      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: "" };
+  }
+}