@marimo-team/islands 0.23.15-dev32 → 0.23.15-dev4

package/src/core/islands/__tests__/bridge.test.ts

@@ -45,22 +45,14 @@ vi.stubGlobal("URL", MockURL);
 const {
   mockBridge,
   mockLoadPackages,
-  mockReplaceSessionRequest,
-  mockStopSessionRequest,
   mockStartSessionRequest,
-  mockMessageListeners,
-  mockStoreSet,
   mockParseMarimoIslandApps,
   mockCreateMarimoFile,
   mockGetMarimoExportContext,
 } = vi.hoisted(() => ({
   mockBridge: vi.fn(),
   mockLoadPackages: vi.fn(),
-  mockReplaceSessionRequest: vi.fn(),
-  mockStopSessionRequest: vi.fn(),
   mockStartSessionRequest: vi.fn(),
-  mockMessageListeners: new Map<string, (payload: never) => void>(),
-  mockStoreSet: vi.fn(),
   mockParseMarimoIslandApps: vi.fn<() => TestIslandApp[]>(() => []),
   mockCreateMarimoFile: vi.fn(),
   mockGetMarimoExportContext: vi.fn<() => TestExportContext | undefined>(
@@ -74,19 +66,13 @@ vi.mock("@/core/wasm/rpc", () => ({
       request: {
         bridge: mockBridge,
         loadPackages: mockLoadPackages,
-        replaceSession: mockReplaceSessionRequest,
         startSession: mockStartSessionRequest,
-        stopSession: mockStopSessionRequest,
       },
       send: {
         consumerReady: vi.fn(),
       },
     },
-    addMessageListener: vi.fn(
-      (name: string, listener: (payload: never) => void) => {
-        mockMessageListeners.set(name, listener);
-      },
-    ),
+    addMessageListener: vi.fn(),
   }),
 }));
 
@@ -111,11 +97,10 @@ vi.mock("@/core/meta/globals", () => ({
 }));
 
 // Mock the jotai store
-vi.mock("@/core/state/jotai", async (importOriginal) => ({
-  ...(await importOriginal<typeof import("@/core/state/jotai")>()),
+vi.mock("@/core/state/jotai", () => ({
   store: {
     get: vi.fn(),
-    set: mockStoreSet,
+    set: vi.fn(),
   },
 }));
 
@@ -130,69 +115,40 @@ describe("IslandsPyodideBridge", () => {
     mockParseMarimoIslandApps.mockReturnValue([]);
     mockCreateMarimoFile.mockReset();
     mockGetMarimoExportContext.mockReturnValue(undefined);
-    mockMessageListeners.clear();
-    bridge = new IslandsPyodideBridge();
+    bridge = new IslandsPyodideBridge({ autoStartSessions: false });
   });
 
-  function signalWorkerReady() {
-    const listener = mockMessageListeners.get("ready");
-    if (!listener) {
-      throw new Error("Missing worker ready listener");
-    }
-    listener({} as never);
-  }
-
-  function app(id = "app-1"): TestIslandApp {
-    return {
-      id,
-      cells: [{ code: "x = 1", idx: 0, output: "<div>1</div>" }],
-    };
-  }
-
-  function setSingleApp(file = "generated app 1", appId = "app-1") {
-    mockParseMarimoIslandApps.mockReturnValue([app(appId)]);
-    mockCreateMarimoFile.mockReturnValue(file);
-  }
-
-  function mockSingleApp(file = "generated app 1", appId = "app-1") {
-    setSingleApp(file, appId);
-    signalWorkerReady();
-  }
-
-  function sendSlider() {
-    return bridge.sendComponentValues({
-      objectIds: [uiElementId("slider-1")],
-      values: [2],
-    });
-  }
-
-  async function initializeSingleApp() {
-    mockSingleApp();
-    await bridge.initializeApps();
-  }
-
-  describe("app initialization", () => {
+  describe("startSessionsForAllApps", () => {
     it("should prefer trusted export notebook code when there is exactly one reactive app", async () => {
-      mockParseMarimoIslandApps.mockReturnValue([app()]);
+      mockParseMarimoIslandApps.mockReturnValue([
+        {
+          id: "app-1",
+          cells: [{ code: "x = 1", idx: 0, output: "<div>1</div>" }],
+        },
+      ]);
       mockGetMarimoExportContext.mockReturnValue({
         trusted: true,
         notebookCode:
           "import marimo\napp = marimo.App()\n@app.cell\ndef __():\n    x = 1\n    return",
       });
 
-      signalWorkerReady();
-      await bridge.initializeApps();
+      await (
+        bridge as unknown as { startSessionsForAllApps(): Promise<void> }
+      ).startSessionsForAllApps();
 
       expect(mockCreateMarimoFile).not.toHaveBeenCalled();
       expect(mockStartSessionRequest).toHaveBeenCalledWith({
         appId: "app-1",
         code: "import marimo\napp = marimo.App()\n@app.cell\ndef __():\n    x = 1\n    return",
-        sessionGeneration: 1,
       });
     });
 
     it("should ignore trusted export notebook code for a payload-backed app", async () => {
-      const payloadApp = { ...app(), payloadBacked: true };
+      const payloadApp = {
+        id: "app-1",
+        payloadBacked: true,
+        cells: [{ code: "x = 1", idx: 0, output: "<div>1</div>" }],
+      };
       mockParseMarimoIslandApps.mockReturnValue([payloadApp]);
       mockGetMarimoExportContext.mockReturnValue({
         trusted: true,
@@ -200,19 +156,28 @@ describe("IslandsPyodideBridge", () => {
       });
       mockCreateMarimoFile.mockReturnValue("generated payload app");
 
-      signalWorkerReady();
-      await bridge.initializeApps();
+      await (
+        bridge as unknown as { startSessionsForAllApps(): Promise<void> }
+      ).startSessionsForAllApps();
 
       expect(mockCreateMarimoFile).toHaveBeenCalledWith(payloadApp);
       expect(mockStartSessionRequest).toHaveBeenCalledWith({
         appId: "app-1",
         code: "generated payload app",
-        sessionGeneration: 1,
       });
     });
 
     it("should keep synthesized per-app files for multiple reactive apps even when export context exists", async () => {
-      mockParseMarimoIslandApps.mockReturnValue([app(), app("app-2")]);
+      mockParseMarimoIslandApps.mockReturnValue([
+        {
+          id: "app-1",
+          cells: [{ code: "x = 1", idx: 0, output: "<div>1</div>" }],
+        },
+        {
+          id: "app-2",
+          cells: [{ code: "y = 2", idx: 0, output: "<div>2</div>" }],
+        },
+      ]);
       mockGetMarimoExportContext.mockReturnValue({
         trusted: true,
         notebookCode: "full notebook should be ignored",
@@ -221,281 +186,43 @@ describe("IslandsPyodideBridge", () => {
         .mockReturnValueOnce("generated app 1")
         .mockReturnValueOnce("generated app 2");
 
-      signalWorkerReady();
-      await bridge.initializeApps();
+      await (
+        bridge as unknown as { startSessionsForAllApps(): Promise<void> }
+      ).startSessionsForAllApps();
 
       expect(mockCreateMarimoFile).toHaveBeenCalledTimes(2);
       expect(mockStartSessionRequest).toHaveBeenNthCalledWith(1, {
         appId: "app-1",
         code: "generated app 1",
-        sessionGeneration: 1,
       });
       expect(mockStartSessionRequest).toHaveBeenNthCalledWith(2, {
         appId: "app-2",
         code: "generated app 2",
-        sessionGeneration: 2,
-      });
-
-      await sendSlider();
-      await bridge.stopSession();
-
-      expect(mockReplaceSessionRequest).not.toHaveBeenCalled();
-      expect(mockStopSessionRequest).not.toHaveBeenCalled();
-      expect(mockBridge).toHaveBeenCalledWith({
-        appId: "app-1",
-        functionName: "put_control_request",
-        payload: {
-          objectIds: ["slider-1"],
-          token: "test-uuid-12345",
-          type: "update-ui-element",
-          values: [2],
-        },
-        sessionGeneration: 1,
       });
     });
 
     it("should synthesize a file for a single app when no trusted export context is present", async () => {
-      setSingleApp();
+      mockParseMarimoIslandApps.mockReturnValue([
+        {
+          id: "app-1",
+          cells: [{ code: "x = 1", idx: 0, output: "<div>1</div>" }],
+        },
+      ]);
+      mockCreateMarimoFile.mockReturnValue("generated app 1");
 
-      signalWorkerReady();
-      await bridge.initializeApps();
+      await (
+        bridge as unknown as { startSessionsForAllApps(): Promise<void> }
+      ).startSessionsForAllApps();
 
       expect(mockCreateMarimoFile).toHaveBeenCalledTimes(1);
       expect(mockStartSessionRequest).toHaveBeenCalledWith({
         appId: "app-1",
         code: "generated app 1",
-        sessionGeneration: 1,
-      });
-    });
-  });
-
-  describe("app lifecycle", () => {
-    it("holds controls until the first app session is ready", async () => {
-      setSingleApp();
-
-      const initialization = bridge.initializeApps();
-      const control = sendSlider();
-      await Promise.resolve();
-
-      expect(mockBridge).not.toHaveBeenCalled();
-
-      signalWorkerReady();
-      await Promise.all([initialization, control]);
-      expect(mockBridge).toHaveBeenCalledWith(
-        expect.objectContaining({
-          appId: "app-1",
-          sessionGeneration: 1,
-        }),
-      );
-    });
-
-    it("starts the first app, skips duplicates, and replaces changed source", async () => {
-      mockSingleApp();
-
-      await bridge.initializeApps();
-      await bridge.initializeApps();
-      await sendSlider();
-
-      expect(mockStartSessionRequest).toHaveBeenCalledOnce();
-      expect(mockStartSessionRequest).toHaveBeenCalledWith({
-        appId: "app-1",
-        code: "generated app 1",
-        sessionGeneration: 1,
-      });
-      expect(mockReplaceSessionRequest).not.toHaveBeenCalled();
-      expect(mockBridge).toHaveBeenCalledWith(
-        expect.objectContaining({
-          appId: "app-1",
-          sessionGeneration: 1,
-        }),
-      );
-      mockCreateMarimoFile.mockReturnValue("generated app 2");
-      await bridge.initializeApps();
-
-      expect(mockReplaceSessionRequest).toHaveBeenCalledOnce();
-      expect(mockReplaceSessionRequest).toHaveBeenCalledWith({
-        appId: "app-1",
-        code: "generated app 2",
-        sessionGeneration: 2,
-      });
-      expect(mockStoreSet).toHaveBeenCalledOnce();
-    });
-
-    it("stops the matching active app", async () => {
-      mockSingleApp();
-      await bridge.initializeApps();
-
-      await bridge.stopSession("other-app");
-      await bridge.stopSession("app-1");
-
-      expect(mockStopSessionRequest).toHaveBeenCalledOnce();
-      expect(mockStopSessionRequest).toHaveBeenCalledWith({
-        appId: "app-1",
-        sessionGeneration: 1,
-      });
-
-      mockBridge.mockClear();
-      const control = sendSlider();
-      await Promise.resolve();
-      expect(mockBridge).not.toHaveBeenCalled();
-
-      mockSingleApp("generated app 2");
-      await bridge.initializeApps();
-      await control;
-
-      expect(mockBridge).toHaveBeenCalledWith(
-        expect.objectContaining({ sessionGeneration: 2 }),
-      );
-    });
-
-    it("replaces the active session after stop fails", async () => {
-      mockSingleApp();
-      await bridge.initializeApps();
-      mockStopSessionRequest.mockRejectedValueOnce(new Error("stop failed"));
-
-      await expect(bridge.stopSession("app-1")).rejects.toThrow("stop failed");
-
-      mockSingleApp("generated app 2", "app-2");
-
-      await bridge.initializeApps();
-
-      expect(mockStartSessionRequest).toHaveBeenCalledOnce();
-      expect(mockReplaceSessionRequest).toHaveBeenCalledWith({
-        appId: "app-2",
-        code: "generated app 2",
-        sessionGeneration: 2,
-      });
-    });
-
-    it("keeps controls sent during stop scoped to the stopped app", async () => {
-      mockSingleApp();
-      await bridge.initializeApps();
-      mockBridge.mockClear();
-      let finishStop!: () => void;
-      mockStopSessionRequest.mockReturnValueOnce(
-        new Promise<void>((resolve) => {
-          finishStop = resolve;
-        }),
-      );
-
-      const stop = bridge.stopSession("app-1");
-      await vi.waitFor(() =>
-        expect(mockStopSessionRequest).toHaveBeenCalledOnce(),
-      );
-      const control = sendSlider();
-      mockSingleApp("generated app 2", "app-2");
-      const nextInitialization = bridge.initializeApps();
-
-      finishStop();
-      await Promise.all([stop, control, nextInitialization]);
-
-      expect(mockBridge).toHaveBeenCalledWith(
-        expect.objectContaining({
-          appId: "app-1",
-          sessionGeneration: 1,
-        }),
-      );
-    });
-
-    it("allows initialization to retry after replacement fails", async () => {
-      mockSingleApp();
-      await bridge.initializeApps();
-      mockCreateMarimoFile.mockReturnValue("generated app 2");
-      mockReplaceSessionRequest.mockRejectedValueOnce(
-        new Error("replacement failed"),
-      );
-
-      await expect(bridge.initializeApps()).rejects.toThrow(
-        "replacement failed",
-      );
-      await bridge.initializeApps();
-
-      expect(mockStartSessionRequest).toHaveBeenCalledOnce();
-      expect(mockReplaceSessionRequest).toHaveBeenCalledTimes(2);
-      expect(mockReplaceSessionRequest).toHaveBeenLastCalledWith({
-        appId: "app-1",
-        code: "generated app 2",
-        sessionGeneration: 3,
       });
     });
-
-    it("scopes controls to the replacement active when they were sent", async () => {
-      mockSingleApp();
-      await bridge.initializeApps();
-      mockBridge.mockClear();
-      mockCreateMarimoFile.mockReturnValue("generated app 2");
-      let finishSecondApp!: () => void;
-      let finishThirdApp!: () => void;
-      mockReplaceSessionRequest
-        .mockReturnValueOnce(
-          new Promise<void>((resolve) => {
-            finishSecondApp = resolve;
-          }),
-        )
-        .mockReturnValueOnce(
-          new Promise<void>((resolve) => {
-            finishThirdApp = resolve;
-          }),
-        );
-
-      const secondInitialization = bridge.initializeApps();
-      await vi.waitFor(() =>
-        expect(mockReplaceSessionRequest).toHaveBeenCalledOnce(),
-      );
-      const control = sendSlider();
-      mockSingleApp("generated app 3", "app-2");
-      const thirdInitialization = bridge.initializeApps();
-
-      expect(mockBridge).not.toHaveBeenCalled();
-
-      finishSecondApp();
-      await vi.waitFor(() => expect(mockBridge).toHaveBeenCalledOnce());
-      expect(mockBridge).toHaveBeenCalledWith(
-        expect.objectContaining({
-          appId: "app-1",
-          sessionGeneration: 2,
-        }),
-      );
-
-      finishThirdApp();
-      await Promise.all([secondInitialization, thirdInitialization, control]);
-    });
-
-    it("keeps controls from a failed replacement out of the next app", async () => {
-      mockSingleApp();
-      await bridge.initializeApps();
-      mockBridge.mockClear();
-      mockCreateMarimoFile.mockReturnValue("generated app 2");
-      let failReplacement!: (error: Error) => void;
-      mockReplaceSessionRequest.mockReturnValueOnce(
-        new Promise<void>((_resolve, reject) => {
-          failReplacement = reject;
-        }),
-      );
-
-      const failedInitialization = bridge.initializeApps();
-      await vi.waitFor(() =>
-        expect(mockReplaceSessionRequest).toHaveBeenCalledOnce(),
-      );
-      const control = sendSlider();
-      failReplacement(new Error("replacement failed"));
-      await expect(failedInitialization).rejects.toThrow("replacement failed");
-
-      mockSingleApp("generated app 3", "app-2");
-      await Promise.all([bridge.initializeApps(), control]);
-
-      expect(mockBridge).toHaveBeenCalledWith(
-        expect.objectContaining({
-          appId: "app-1",
-          sessionGeneration: 2,
-        }),
-      );
-    });
   });
 
   describe("sendComponentValues", () => {
-    beforeEach(initializeSingleApp);
-
     it("should include type field and token in control request", async () => {
       const request = {
         objectIds: [uiElementId("Hbol-0")],
@@ -505,7 +232,6 @@ describe("IslandsPyodideBridge", () => {
       await bridge.sendComponentValues(request);
 
       expect(mockBridge).toHaveBeenCalledWith({
-        appId: "app-1",
         functionName: "put_control_request",
         payload: {
           type: "update-ui-element",
@@ -513,7 +239,6 @@ describe("IslandsPyodideBridge", () => {
           values: [58],
           token: "test-uuid-12345",
         },
-        sessionGeneration: 1,
       });
     });
 
@@ -526,21 +251,17 @@ describe("IslandsPyodideBridge", () => {
       await bridge.sendComponentValues(request);
 
       expect(mockBridge).toHaveBeenCalledWith({
-        appId: "app-1",
         functionName: "put_control_request",
         payload: expect.objectContaining({
           type: "update-ui-element",
           objectIds: ["slider-1", "slider-2"],
           values: [10, 20],
         }),
-        sessionGeneration: 1,
       });
     });
   });
 
   describe("sendFunctionRequest", () => {
-    beforeEach(initializeSingleApp);
-
     it("should include type field in control request", async () => {
       const request = {
         functionCallId: requestId("call-123"),
@@ -552,7 +273,6 @@ describe("IslandsPyodideBridge", () => {
       await bridge.sendFunctionRequest(request);
 
       expect(mockBridge).toHaveBeenCalledWith({
-        appId: "app-1",
         functionName: "put_control_request",
         payload: {
           type: "invoke-function",
@@ -561,14 +281,11 @@ describe("IslandsPyodideBridge", () => {
           functionName: "my_function",
           args: { x: 1, y: 2 },
         },
-        sessionGeneration: 1,
       });
     });
   });
 
   describe("sendRun", () => {
-    beforeEach(initializeSingleApp);
-
     it("should include type field in control request", async () => {
       const request = {
         cellIds: [cellId("cell-1"), cellId("cell-2")],
@@ -578,14 +295,12 @@ describe("IslandsPyodideBridge", () => {
       await bridge.sendRun(request);
 
       expect(mockBridge).toHaveBeenCalledWith({
-        appId: "app-1",
         functionName: "put_control_request",
         payload: {
           type: "execute-cells",
           cellIds: ["cell-1", "cell-2"],
           codes: ["print('hello')", "print('world')"],
         },
-        sessionGeneration: 1,
       });
     });
 
@@ -598,11 +313,7 @@ describe("IslandsPyodideBridge", () => {
       await bridge.sendRun(request);
 
       // Verify loadPackages was called with joined codes
-      expect(mockLoadPackages).toHaveBeenCalledWith({
-        appId: "app-1",
-        code: "import pandas",
-        sessionGeneration: 1,
-      });
+      expect(mockLoadPackages).toHaveBeenCalledWith("import pandas");
 
       // Verify order: loadPackages should be called before bridge
       const loadPackagesCallOrder =
@@ -613,8 +324,6 @@ describe("IslandsPyodideBridge", () => {
   });
 
   describe("sendModelValue", () => {
-    beforeEach(initializeSingleApp);
-
     it("should include type field in control request", async () => {
       const request = {
         modelId: widgetModelId("widget-1"),
@@ -629,7 +338,6 @@ describe("IslandsPyodideBridge", () => {
       await bridge.sendModelValue(request);
 
       expect(mockBridge).toHaveBeenCalledWith({
-        appId: "app-1",
         functionName: "put_control_request",
         payload: {
           type: "model",
@@ -641,14 +349,11 @@ describe("IslandsPyodideBridge", () => {
           },
           buffers: [],
         },
-        sessionGeneration: 1,
       });
     });
   });
 
   describe("control request message format", () => {
-    beforeEach(initializeSingleApp);
-
     it("should always include the type field required by msgspec", async () => {
       // Test all methods to ensure they include the type field
       await bridge.sendComponentValues({ objectIds: [], values: [] });