@marimo-team/islands 0.19.5-dev12 → 0.19.5-dev13

package/dist/main.js

@@ -32662,6 +32662,7 @@ ${c.sqlString}
       super();
       __publicField(this, "isClosed", false);
       __publicField(this, "hasConnectedBefore", false);
+      __publicField(this, "pendingSubscriptions", []);
       this.options = e, this.delegate = void 0;
     }
     createDelegate() {
@@ -32670,7 +32671,9 @@ ${c.sqlString}
       } catch (e) {
         Logger.warn("Error closing old WebSocket delegate", e);
       }
-      return this.delegate = new import_build.WebSocketTransport(this.options.getWsUrl()), this.delegate;
+      this.delegate = new import_build.WebSocketTransport(this.options.getWsUrl());
+      for (let { event: e, handler: r } of this.pendingSubscriptions) this.delegate.subscribe(e, r);
+      return this.delegate;
     }
     isDelegateClosedOrClosing() {
       if (!this.delegate) return true;
@@ -32697,6 +32700,18 @@ ${c.sqlString}
       var _a2;
       this.isClosed = true, (_a2 = this.delegate) == null ? void 0 : _a2.close(), this.delegate = void 0, this.connectionPromise = void 0;
     }
+    subscribe(...e) {
+      super.subscribe(...e);
+      let [r, c] = e;
+      this.pendingSubscriptions.push({
+        event: r,
+        handler: c
+      }), this.delegate && this.delegate.subscribe(r, c);
+    }
+    unsubscribe(...e) {
+      let r = super.unsubscribe(...e), [c, d] = e;
+      return this.pendingSubscriptions = this.pendingSubscriptions.filter((e2) => !(e2.event === c && e2.handler === d)), this.delegate && this.delegate.unsubscribe(c, d), r;
+    }
     async sendData(e, r) {
       var _a2;
       if (this.isDelegateClosedOrClosing()) {
@@ -101063,7 +101078,7 @@ Defaulting to \`null\`.`;
       return Logger.warn("Failed to get version from mount config"), null;
     }
   }
-  const marimoVersionAtom = atom(getVersionFromMountConfig() || "0.19.5-dev12"), showCodeInRunModeAtom = atom(true);
+  const marimoVersionAtom = atom(getVersionFromMountConfig() || "0.19.5-dev13"), showCodeInRunModeAtom = atom(true);
   atom(null);
   var VIRTUAL_FILE_REGEX = /\/@file\/([^\s"&'/]+)\.([\dA-Za-z]+)/g, VirtualFileTracker = class e {
     constructor() {

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@marimo-team/islands",
-  "version": "0.19.5-dev12",
+  "version": "0.19.5-dev13",
   "main": "dist/main.js",
   "types": "dist/index.d.ts",
   "type": "module",

package/src/core/lsp/__tests__/transport.test.ts

@@ -41,6 +41,8 @@ describe("ReconnectingWebSocketTransport", () => {
       this.connect = vi.fn().mockResolvedValue(undefined);
       this.close = vi.fn();
       this.sendData = vi.fn().mockResolvedValue({ result: "success" });
+      this.subscribe = vi.fn();
+      this.unsubscribe = vi.fn();
     });
   });
 
@@ -287,4 +289,151 @@ describe("ReconnectingWebSocketTransport", () => {
       "Reconnect callback failed",
     );
   });
+
+  describe("subscribe", () => {
+    it("should track subscriptions", () => {
+      const getWsUrl = vi.fn(() => mockWsUrl);
+      const transport = new ReconnectingWebSocketTransport({ getWsUrl });
+
+      const handler = vi.fn();
+      transport.subscribe("notification", handler);
+
+      expect((transport as any).pendingSubscriptions).toHaveLength(1);
+      expect((transport as any).pendingSubscriptions[0]).toEqual({
+        event: "notification",
+        handler,
+      });
+    });
+
+    it("should register handler on delegate if it exists", async () => {
+      const getWsUrl = vi.fn(() => mockWsUrl);
+      const transport = new ReconnectingWebSocketTransport({ getWsUrl });
+
+      await transport.connect();
+
+      const handler = vi.fn();
+      transport.subscribe("notification", handler);
+
+      const delegate = (transport as any).delegate;
+      expect(delegate.subscribe).toHaveBeenCalledWith("notification", handler);
+    });
+
+    it("should register pending subscriptions when delegate is created", async () => {
+      const getWsUrl = vi.fn(() => mockWsUrl);
+      const transport = new ReconnectingWebSocketTransport({ getWsUrl });
+
+      const handler1 = vi.fn();
+      const handler2 = vi.fn();
+      transport.subscribe("notification", handler1);
+      transport.subscribe("response", handler2);
+
+      await transport.connect();
+
+      const delegate = (transport as any).delegate;
+      expect(delegate.subscribe).toHaveBeenCalledWith("notification", handler1);
+      expect(delegate.subscribe).toHaveBeenCalledWith("response", handler2);
+    });
+
+    it("should re-register subscriptions on reconnection", async () => {
+      const getWsUrl = vi.fn(() => mockWsUrl);
+      const transport = new ReconnectingWebSocketTransport({ getWsUrl });
+
+      // Add subscription before connection
+      const handler = vi.fn();
+      transport.subscribe("notification", handler);
+
+      // First connection
+      await transport.connect();
+      const firstDelegate = (transport as any).delegate;
+      expect(firstDelegate.subscribe).toHaveBeenCalledWith(
+        "notification",
+        handler,
+      );
+
+      // Clear mock calls
+      firstDelegate.subscribe.mockClear();
+
+      // Simulate connection loss
+      mockConnection.readyState = WebSocket.CLOSED;
+
+      // Reconnect by sending data
+      const data: any = { method: "test", params: [] };
+      await transport.sendData(data, 5000);
+
+      // New delegate should have been created
+      const secondDelegate = (transport as any).delegate;
+      expect(secondDelegate).not.toBe(firstDelegate);
+
+      // Subscription should be re-registered on new delegate
+      expect(secondDelegate.subscribe).toHaveBeenCalledWith(
+        "notification",
+        handler,
+      );
+    });
+  });
+
+  describe("unsubscribe", () => {
+    it("should remove subscription from tracking", () => {
+      const getWsUrl = vi.fn(() => mockWsUrl);
+      const transport = new ReconnectingWebSocketTransport({ getWsUrl });
+
+      const handler = vi.fn();
+      transport.subscribe("notification", handler);
+      expect((transport as any).pendingSubscriptions).toHaveLength(1);
+
+      transport.unsubscribe("notification", handler);
+      expect((transport as any).pendingSubscriptions).toHaveLength(0);
+    });
+
+    it("should unregister from delegate if it exists", async () => {
+      const getWsUrl = vi.fn(() => mockWsUrl);
+      const transport = new ReconnectingWebSocketTransport({ getWsUrl });
+
+      await transport.connect();
+
+      const handler = vi.fn();
+      transport.subscribe("notification", handler);
+
+      const delegate = (transport as any).delegate;
+      delegate.unsubscribe.mockClear();
+
+      transport.unsubscribe("notification", handler);
+
+      expect(delegate.unsubscribe).toHaveBeenCalledWith(
+        "notification",
+        handler,
+      );
+    });
+
+    it("should not re-register unsubscribed handlers on reconnection", async () => {
+      const getWsUrl = vi.fn(() => mockWsUrl);
+      const transport = new ReconnectingWebSocketTransport({ getWsUrl });
+
+      const handler1 = vi.fn();
+      const handler2 = vi.fn();
+      transport.subscribe("notification", handler1);
+      transport.subscribe("response", handler2);
+
+      await transport.connect();
+
+      // Unsubscribe handler1
+      transport.unsubscribe("notification", handler1);
+
+      // Simulate connection loss
+      mockConnection.readyState = WebSocket.CLOSED;
+
+      // Reconnect by sending data
+      const data: any = { method: "test", params: [] };
+      await transport.sendData(data, 5000);
+
+      const newDelegate = (transport as any).delegate;
+
+      // Only handler2 should be registered on the new delegate
+      expect(newDelegate.subscribe).not.toHaveBeenCalledWith(
+        "notification",
+        handler1,
+      );
+      expect(newDelegate.subscribe).toHaveBeenCalledWith("response", handler2);
+    });
+  });
 });

package/src/core/lsp/transport.ts

@@ -23,6 +23,11 @@ export interface ReconnectingWebSocketTransportOptions {
   onReconnect?: () => Promise<void>;
 }
 
+interface Subscription {
+  event: "pending" | "notification" | "response" | "error";
+  handler: Parameters<Transport["subscribe"]>[1];
+}
+
 /**
  * A WebSocket transport that automatically reconnects when the connection is lost.
  * This handles cases like computer sleep/wake or network interruptions.
@@ -33,6 +38,7 @@ export class ReconnectingWebSocketTransport extends Transport {
   private connectionPromise: Promise<void> | undefined;
   private isClosed = false;
   private hasConnectedBefore = false;
+  private pendingSubscriptions: Subscription[] = [];
 
   constructor(options: ReconnectingWebSocketTransportOptions) {
     super();
@@ -55,6 +61,12 @@ export class ReconnectingWebSocketTransport extends Transport {
 
     // 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);
+    }
+
     return this.delegate;
   }
 
@@ -134,6 +146,42 @@ export class ReconnectingWebSocketTransport extends Transport {
     this.connectionPromise = undefined;
   }
 
+  override subscribe(...args: Parameters<Transport["subscribe"]>): void {
+    // Register handler on parent Transport
+    super.subscribe(...args);
+
+    const [event, handler] = args;
+
+    // Track the subscription
+    this.pendingSubscriptions.push({ event, handler });
+
+    // Also register on delegate if it exists
+    if (this.delegate) {
+      this.delegate.subscribe(event, handler);
+    }
+  }
+
+  override unsubscribe(
+    ...args: Parameters<Transport["unsubscribe"]>
+  ): import("events").EventEmitter | undefined {
+    // Unregister from parent
+    const result = super.unsubscribe(...args);
+
+    const [event, handler] = args;
+
+    // Remove from pending subscriptions
+    this.pendingSubscriptions = this.pendingSubscriptions.filter(
+      (sub) => !(sub.event === event && sub.handler === handler),
+    );
+
+    // Also unregister from delegate if it exists
+    if (this.delegate) {
+      this.delegate.unsubscribe(event, handler);
+    }
+
+    return result;
+  }
+
   override async sendData(
     data: JSONRPCRequestData,
     timeout: number | null | undefined,