@marimo-team/frontend 0.23.16-dev5 → 0.23.16-dev50

package/src/core/lsp/transport.ts

@@ -1,9 +1,25 @@
 /* Copyright 2026 Marimo. All rights reserved. */
-import { WebSocketTransport } from "@open-rpc/client-js";
-import type { JSONRPCRequestData } from "@open-rpc/client-js/build/Request";
-import { Transport } from "@open-rpc/client-js/build/transports/Transport";
+
+import {
+  type JSONRPCMessage,
+  type Transport,
+  WebSocketTransport,
+} from "@marimo-team/codemirror-languageserver";
 import { Logger } from "@/utils/Logger";
 
+const MAX_PENDING_MESSAGES = 1000;
+const RESYNCED_NOTIFICATION_METHODS = new Set([
+  "textDocument/didOpen",
+  "textDocument/didChange",
+  "textDocument/didClose",
+]);
+const RESTORATION_METHODS = new Set([
+  "initialize",
+  "initialized",
+  "workspace/didChangeConfiguration",
+  "textDocument/didOpen",
+]);
+
 export interface ReconnectingWebSocketTransportOptions {
   /**
    * Function that returns the WebSocket URL to connect to.
@@ -21,183 +37,339 @@ export interface ReconnectingWebSocketTransportOptions {
    * This allows the LSP client to re-synchronize state (e.g., re-send document open notifications).
    */
   onReconnect?: () => Promise<void>;
-}
 
-interface Subscription {
-  event: "pending" | "notification" | "response" | "error";
-  handler: Parameters<Transport["subscribe"]>[1];
+  /**
+   * Number of connection attempts in each retry cycle.
+   * @default 1
+   */
+  retries?: number;
+
+  /**
+   * Delay between connection attempts.
+   * @default 1000
+   */
+  retryDelayMs?: number;
+
+  /**
+   * Called when every connection attempt in a retry cycle fails.
+   */
+  onConnectionFailure?: (error: Error) => void;
 }
 
 /**
- * A WebSocket transport that automatically reconnects when the connection is lost.
- * This handles cases like computer sleep/wake or network interruptions.
+ * A WebSocket transport that automatically reconnects when the connection is
+ * lost. Messages produced while reconnecting are replayed only after the LSP
+ * client has re-initialized and restored its document state.
+ *
+ * Deliberately does not implement the optional `Transport.onClose`: JSONRPCClient
+ * latches the first close error and never clears it, so forwarding drops would
+ * fail every later request instead of surviving a reconnect.
  */
-export class ReconnectingWebSocketTransport extends Transport {
+export class ReconnectingWebSocketTransport implements Transport {
   private delegate: WebSocketTransport | undefined;
-  private readonly options: ReconnectingWebSocketTransportOptions;
   private connectionPromise: Promise<void> | undefined;
+  private isDelegateConnected = false;
+  /**
+   * Identifies the current connection attempt. Bumping it abandons whatever is
+   * in flight: an attempt whose socket has already died may stay parked on an
+   * LSP request timeout, and it must neither delay recovery nor mutate state
+   * belonging to the attempt that replaced it.
+   */
+  private connectionGeneration = 0;
+  /** Generation currently restoring LSP state, if any. */
+  private restoringGeneration: number | undefined;
   private isClosed = false;
   private hasConnectedBefore = false;
-  private pendingSubscriptions: Subscription[] = [];
+  private readonly messageHandlers = new Set<
+    (message: JSONRPCMessage) => void
+  >();
+  private pendingMessages: JSONRPCMessage[] = [];
+  /**
+   * Backlog detached by the in-flight attempt. Held here rather than in a local
+   * so that abandoning that attempt can hand the messages straight back to the
+   * queue instead of stranding them until it eventually settles.
+   */
+  private inFlightBacklog: JSONRPCMessage[] | undefined;
+  private hasLoggedQueueOverflow = false;
+  private disposeDelegateMessage: (() => void) | undefined;
+  private disposeDelegateClose: (() => void) | undefined;
+  private readonly options: Required<
+    Pick<ReconnectingWebSocketTransportOptions, "retries" | "retryDelayMs">
+  > &
+    Omit<ReconnectingWebSocketTransportOptions, "retries" | "retryDelayMs">;
+
+  public onReconnect: (() => Promise<void>) | undefined;
 
   constructor(options: ReconnectingWebSocketTransportOptions) {
-    super();
-    this.options = options;
-    this.delegate = undefined;
+    this.options = {
+      ...options,
+      retries: options.retries ?? 1,
+      retryDelayMs: options.retryDelayMs ?? 1000,
+    };
+    this.onReconnect = options.onReconnect;
   }
 
-  /**
-   * Create a new WebSocket delegate, replacing any existing one.
-   */
-  private createDelegate(): WebSocketTransport {
-    // Close the old delegate if it exists
-    if (this.delegate) {
-      try {
-        this.delegate.close();
-      } catch (error) {
-        Logger.warn("Error closing old WebSocket delegate", error);
-      }
+  async connect(): Promise<void> {
+    if (this.isClosed) {
+      throw new Error("Transport is closed");
     }
-
-    // Create a new delegate
-    this.delegate = new WebSocketTransport(this.options.getWsUrl());
-
-    // Re-register all pending subscriptions on the new delegate
-    for (const { event, handler } of this.pendingSubscriptions) {
-      this.delegate.subscribe(event, handler);
+    if (this.isDelegateConnected) {
+      return;
+    }
+    if (this.connectionPromise) {
+      return this.connectionPromise;
     }
 
-    return this.delegate;
+    const generation = ++this.connectionGeneration;
+    const promise = this.connectWithRetries(generation).finally(() => {
+      // An abandoned attempt must not clear its successor's promise.
+      if (this.connectionPromise === promise) {
+        this.connectionPromise = undefined;
+      }
+    });
+    this.connectionPromise = promise;
+    return promise;
   }
 
-  /**
-   * Check if the current delegate's WebSocket is in a closed or closing state.
-   */
-  private isDelegateClosedOrClosing(): boolean {
-    if (!this.delegate) {
-      return true;
+  /** Whether `generation` has been superseded or the transport was closed. */
+  private isStale(generation: number): boolean {
+    return this.isClosed || generation !== this.connectionGeneration;
+  }
+
+  send(message: JSONRPCMessage): void {
+    if (this.isClosed) {
+      return;
+    }
+    if (
+      this.delegate &&
+      this.isDelegateConnected &&
+      (this.restoringGeneration === undefined ||
+        this.canSendWhileRestoring(message))
+    ) {
+      this.delegate.send(message);
+      return;
     }
 
-    // Access the internal connection to check its readyState
-    const ws = this.delegate.connection;
-    if (!ws) {
-      return true;
+    if (this.pendingMessages.length < MAX_PENDING_MESSAGES) {
+      this.pendingMessages.push(message);
+    } else if (!this.hasLoggedQueueOverflow) {
+      // Log once per episode; the queue fills one keystroke at a time.
+      this.hasLoggedQueueOverflow = true;
+      Logger.error(
+        `LSP reconnect queue is full (${MAX_PENDING_MESSAGES}); dropping messages until it drains`,
+      );
     }
+    this.reconnect();
+  }
 
-    // WebSocket.CLOSING = 2, WebSocket.CLOSED = 3
-    return (
-      ws.readyState === WebSocket.CLOSING || ws.readyState === WebSocket.CLOSED
-    );
+  onMessage(handler: (message: JSONRPCMessage) => void): () => void {
+    this.messageHandlers.add(handler);
+    return () => {
+      this.messageHandlers.delete(handler);
+    };
   }
 
-  override async connect() {
-    // Don't reconnect if explicitly closed
+  close(): void {
     if (this.isClosed) {
-      throw new Error("Transport is closed");
+      return;
     }
+    this.isClosed = true;
+    this.isDelegateConnected = false;
+    this.restoringGeneration = undefined;
+    // Abandon anything in flight so it cannot touch state after teardown.
+    this.connectionGeneration++;
+    this.connectionPromise = undefined;
+    this.pendingMessages = [];
+    this.inFlightBacklog = undefined;
+    this.messageHandlers.clear();
+    this.disposeDelegate();
+    this.delegate?.close();
+    this.delegate = undefined;
+  }
 
-    // If already connecting, wait for that connection
-    if (this.connectionPromise) {
-      return this.connectionPromise;
-    }
+  private async connectWithRetries(generation: number): Promise<void> {
+    let lastError: Error | undefined;
 
-    this.connectionPromise = (async () => {
+    for (let attempt = 1; attempt <= this.options.retries; attempt++) {
       try {
-        // Wait for dependencies to be ready (e.g., runtime connection)
-        if (this.options.waitForConnection) {
-          await this.options.waitForConnection();
-        }
-
-        // Create a new delegate if needed
-        let delegate = this.delegate;
-        if (!delegate || this.isDelegateClosedOrClosing()) {
-          delegate = this.createDelegate();
+        await this.connectDelegate(generation);
+        return;
+      } catch (error) {
+        lastError = error instanceof Error ? error : new Error(String(error));
+        if (this.isStale(generation)) {
+          // Superseded while we were failing; the live attempt owns the
+          // delegate now, so tear down nothing and retry nothing.
+          throw lastError;
         }
-
-        // Connect the delegate
-        await delegate.connect();
-        Logger.log("WebSocket transport connected successfully");
-
-        // If this is a reconnection, call the onReconnect callback
-        const isReconnection = this.hasConnectedBefore;
-        this.hasConnectedBefore = true;
-
-        if (isReconnection && this.options.onReconnect) {
-          Logger.log("Calling onReconnect callback to re-synchronize state");
-          await this.options.onReconnect();
+        this.invalidateDelegate();
+        Logger.warn(
+          `WebSocket connection attempt ${attempt}/${this.options.retries} failed`,
+          lastError,
+        );
+
+        if (attempt < this.options.retries) {
+          await new Promise((resolve) =>
+            setTimeout(resolve, this.options.retryDelayMs),
+          );
         }
-      } catch (error) {
-        Logger.error("WebSocket transport connection failed", error);
-        // Clear the delegate on failure so we create a new one on retry
-        this.delegate = undefined;
-        throw error;
-      } finally {
-        this.connectionPromise = undefined;
       }
-    })();
+    }
 
-    return this.connectionPromise;
+    const error = lastError ?? new Error("WebSocket connection failed");
+    this.options.onConnectionFailure?.(error);
+    throw error;
   }
 
-  override close() {
-    this.isClosed = true;
-    this.delegate?.close();
-    this.delegate = undefined;
-    this.connectionPromise = undefined;
-  }
-
-  override subscribe(...args: Parameters<Transport["subscribe"]>): void {
-    // Register handler on parent Transport
-    super.subscribe(...args);
+  private async connectDelegate(generation: number): Promise<void> {
+    await this.options.waitForConnection?.();
+    if (this.isStale(generation)) {
+      throw new Error("Connection attempt abandoned");
+    }
 
-    const [event, handler] = args;
+    const delegate = this.createDelegate();
+    await delegate.connect();
+    if (this.isStale(generation) || this.delegate !== delegate) {
+      delegate.close();
+      throw new Error("Transport closed while connecting");
+    }
 
-    // Track the subscription
-    this.pendingSubscriptions.push({ event, handler });
+    this.isDelegateConnected = true;
+    const isReconnection = this.hasConnectedBefore;
+    this.hasConnectedBefore = true;
+
+    // Detach before restoring: notifications queued while the socket was down
+    // are superseded by what `onReconnect` re-sends, but ones produced *during*
+    // restoration are not — they carry edits the restore snapshot never saw,
+    // and the notebook snapshotter will not re-emit them.
+    const queuedWhileDisconnected = this.detachPendingMessages();
+    this.inFlightBacklog = queuedWhileDisconnected;
+
+    try {
+      if (isReconnection) {
+        this.restoringGeneration = generation;
+        try {
+          await this.onReconnect?.();
+        } finally {
+          // Only clear the gate if a newer attempt hasn't taken it over.
+          if (this.restoringGeneration === generation) {
+            this.restoringGeneration = undefined;
+          }
+        }
+      }
 
-    // Also register on delegate if it exists
-    if (this.delegate) {
-      this.delegate.subscribe(event, handler);
+      if (
+        this.isStale(generation) ||
+        this.delegate !== delegate ||
+        !this.isDelegateConnected
+      ) {
+        throw new Error("WebSocket closed while restoring LSP state");
+      }
+    } catch (error) {
+      // Nothing superseded the backlog, so restore it ahead of anything queued
+      // since. If this attempt was abandoned, whoever abandoned it already
+      // reclaimed the backlog and owns the queue now.
+      if (!this.isStale(generation)) {
+        this.reclaimInFlightBacklog();
+      }
+      throw error;
     }
-  }
 
-  override unsubscribe(
-    ...args: Parameters<Transport["unsubscribe"]>
-  ): import("events").EventEmitter | undefined {
-    // Unregister from parent
-    const result = super.unsubscribe(...args);
+    const replayable =
+      isReconnection && this.onReconnect
+        ? queuedWhileDisconnected.filter(
+            (message) =>
+              !(
+                "method" in message &&
+                RESYNCED_NOTIFICATION_METHODS.has(message.method)
+              ),
+          )
+        : queuedWhileDisconnected;
+
+    this.inFlightBacklog = undefined;
+    for (const message of [...replayable, ...this.detachPendingMessages()]) {
+      delegate.send(message);
+    }
 
-    const [event, handler] = args;
+    Logger.log("WebSocket transport connected successfully");
+  }
 
-    // Remove from pending subscriptions
-    this.pendingSubscriptions = this.pendingSubscriptions.filter(
-      (sub) => !(sub.event === event && sub.handler === handler),
-    );
+  private detachPendingMessages(): JSONRPCMessage[] {
+    const messages = this.pendingMessages;
+    this.pendingMessages = [];
+    this.hasLoggedQueueOverflow = false;
+    return messages;
+  }
 
-    // Also unregister from delegate if it exists
-    if (this.delegate) {
-      this.delegate.unsubscribe(event, handler);
+  /** Return the in-flight attempt's backlog to the front of the queue. */
+  private reclaimInFlightBacklog(): void {
+    if (!this.inFlightBacklog) {
+      return;
     }
+    this.pendingMessages = [
+      ...this.inFlightBacklog,
+      ...this.pendingMessages,
+    ].slice(0, MAX_PENDING_MESSAGES);
+    this.inFlightBacklog = undefined;
+  }
 
-    return result;
+  private canSendWhileRestoring(message: JSONRPCMessage): boolean {
+    // Responses to server-initiated requests must not be held or the
+    // initialization handshake could deadlock.
+    return !("method" in message) || RESTORATION_METHODS.has(message.method);
   }
 
-  override async sendData(
-    data: JSONRPCRequestData,
-    timeout: number | null | undefined,
-  ) {
-    // If the delegate is closed or closing, try to reconnect
-    if (this.isDelegateClosedOrClosing()) {
-      Logger.warn("WebSocket is closed or closing, attempting to reconnect");
-      try {
-        await this.connect();
-      } catch (error) {
-        Logger.error("Failed to reconnect WebSocket", error);
-        throw error;
+  private createDelegate(): WebSocketTransport {
+    this.invalidateDelegate();
+
+    const delegate = new WebSocketTransport(this.options.getWsUrl());
+    this.delegate = delegate;
+    this.disposeDelegateMessage = delegate.onMessage((message) => {
+      for (const handler of this.messageHandlers) {
+        try {
+          handler(message);
+        } catch (error) {
+          Logger.error("LSP message handler failed", error);
+        }
       }
-    }
+    });
+    this.disposeDelegateClose = delegate.onClose((error) => {
+      if (this.delegate !== delegate || this.isClosed) {
+        return;
+      }
+      Logger.warn("WebSocket transport connection closed", error);
+      this.isDelegateConnected = false;
+      this.disposeDelegate();
+      this.delegate = undefined;
+      // Abandon any attempt still in flight: it is bound to the socket that
+      // just died and may sit parked on an LSP request timeout. Take its
+      // backlog and drop the promise so the next connect() starts immediately
+      // rather than joining a doomed one.
+      this.restoringGeneration = undefined;
+      this.reclaimInFlightBacklog();
+      this.connectionPromise = undefined;
+      queueMicrotask(() => this.reconnect());
+    });
+    return delegate;
+  }
+
+  private reconnect(): void {
+    void this.connect().catch((error) => {
+      Logger.error("WebSocket transport reconnection failed", error);
+    });
+  }
+
+  private invalidateDelegate(): void {
+    this.isDelegateConnected = false;
+    this.restoringGeneration = undefined;
+    this.disposeDelegate();
+    this.delegate?.close();
+    this.delegate = undefined;
+  }
 
-    // Send the data using the delegate
-    return this.delegate?.sendData(data, timeout);
+  private disposeDelegate(): void {
+    this.disposeDelegateMessage?.();
+    this.disposeDelegateMessage = undefined;
+    this.disposeDelegateClose?.();
+    this.disposeDelegateClose = undefined;
   }
 }

package/src/core/mode.ts

@@ -78,6 +78,16 @@ export const viewStateAtom = atom<ViewState>({
 
 export const initialModeAtom = atom<AppMode | undefined>(undefined);
 
+/**
+ * Whether the current page hosts a notebook, and so connects to a kernel.
+ * False for non-notebook pages (home, gallery), which are served without a
+ * session.
+ */
+export function isNotebookPage(): boolean {
+  const mode = store.get(initialModeAtom);
+  return mode !== "home" && mode !== "gallery";
+}
+
 export const kioskModeAtom = atom<boolean>(false);
 
 /**

package/src/core/network/__tests__/api.test.ts

@@ -94,4 +94,46 @@ describe("API", () => {
       }),
     ).rejects.toBeInstanceOf(HTTPError);
   });
+
+  it("preserves export metadata from response headers", async () => {
+    const response = new Response("markdown", {
+      headers: {
+        "Content-Disposition":
+          "attachment; filename*=UTF-8''r%C3%A9sum%C3%A9.qmd",
+        "Content-Type": "text/plain; charset=utf-8",
+      },
+    });
+
+    await expect(
+      API.handleExportResponse({ data: "markdown", response }),
+    ).resolves.toEqual({
+      contents: "markdown",
+      filename: "résumé.qmd",
+      mediaType: "text/plain; charset=utf-8",
+    });
+  });
+
+  it("rejects export responses without a filename", async () => {
+    await expect(
+      API.handleExportResponse({
+        data: "markdown",
+        response: new Response(null, {
+          headers: { "Content-Type": "text/plain" },
+        }),
+      }),
+    ).rejects.toThrow("Export response is missing a filename");
+  });
+
+  it("rejects export responses without a media type", async () => {
+    await expect(
+      API.handleExportResponse({
+        data: "markdown",
+        response: new Response(null, {
+          headers: {
+            "Content-Disposition": "attachment; filename*=UTF-8''notebook.md",
+          },
+        }),
+      }),
+    ).rejects.toThrow("Export response is missing a media type");
+  });
 });

package/src/core/network/__tests__/connection.test.ts

@@ -0,0 +1,45 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { afterEach, describe, expect, it } from "vitest";
+import { initialModeAtom } from "@/core/mode";
+import { store } from "@/core/state/jotai";
+import { WebSocketState } from "../../websocket/types";
+import { connectionAtom, waitForConnectionOpenIfNotebook } from "../connection";
+
+const NEVER = Symbol("never");
+
+/**
+ * Resolves to NEVER if the promise does not settle within a macrotask.
+ */
+function settledOrNever(promise: Promise<unknown>) {
+  return Promise.race([
+    promise,
+    new Promise((resolve) => setTimeout(() => resolve(NEVER), 0)),
+  ]);
+}
+
+describe("waitForConnectionOpenIfNotebook", () => {
+  afterEach(() => {
+    store.set(connectionAtom, { state: WebSocketState.NOT_STARTED });
+    store.set(initialModeAtom, undefined);
+  });
+
+  it.each(["home", "gallery"] as const)(
+    "resolves without a connection in %s mode",
+    async (mode) => {
+      store.set(initialModeAtom, mode);
+      await expect(
+        settledOrNever(waitForConnectionOpenIfNotebook()),
+      ).resolves.not.toBe(NEVER);
+    },
+  );
+
+  it("waits for the connection to open in edit mode", async () => {
+    store.set(initialModeAtom, "edit");
+    const promise = waitForConnectionOpenIfNotebook();
+    await expect(settledOrNever(promise)).resolves.toBe(NEVER);
+
+    store.set(connectionAtom, { state: WebSocketState.OPEN });
+    await expect(settledOrNever(promise)).resolves.not.toBe(NEVER);
+  });
+});

package/src/core/network/__tests__/requests-lazy.test.ts

@@ -4,7 +4,11 @@ import { beforeEach, describe, expect, it, vi } from "vitest";
 import { cellId, requestId, uiElementId } from "@/__tests__/branded";
 import type { RuntimeManager } from "../../runtime/runtime";
 import { createLazyRequests } from "../requests-lazy";
-import type { EditRequests, RunRequests } from "../types";
+import type {
+  EditRequests,
+  ExportAvailabilityResponse,
+  RunRequests,
+} from "../types";
 
 // Mock the connection module
 vi.mock("../connection", async () => {
@@ -72,6 +76,33 @@ describe("createLazyRequests", () => {
     expect(mockDelegate.getEnvironmentInfo).toHaveBeenCalledOnce();
   });
 
+  it("gets export availability without starting a kernel", async () => {
+    const availability: ExportAvailabilityResponse = {
+      source: "server",
+      formats: [
+        {
+          format: "ipynb",
+          dependenciesAvailable: false,
+          missingPackages: ["nbformat"],
+        },
+      ],
+    };
+    mockDelegate.getExportAvailability = vi
+      .fn()
+      .mockResolvedValue(availability) as EditRequests["getExportAvailability"];
+
+    const lazyRequests = createLazyRequests(
+      mockDelegate,
+      mockGetRuntimeManager,
+    );
+
+    await expect(lazyRequests.getExportAvailability()).resolves.toEqual(
+      availability,
+    );
+    expect(mockInit).not.toHaveBeenCalled();
+    expect(mockDelegate.getExportAvailability).toHaveBeenCalledOnce();
+  });
+
   it("should call init once before first request", async () => {
     const lazyRequests = createLazyRequests(
       mockDelegate,