@marimo-team/frontend 0.23.16-dev4 → 0.23.16-dev41

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

@@ -10,6 +10,7 @@ import {
 } from "@/__tests__/branded";
 import { notebookAtom } from "@/core/cells/cells";
 import { islandsPendingInitialRunsAtom } from "@/core/islands/state";
+import type { IslandsKernelMessage } from "../worker/worker";
 
 type Base64String = components["schemas"]["Base64String"];
 interface TestIslandApp {
@@ -173,6 +174,19 @@ describe("IslandsPyodideBridge", () => {
     });
   }
 
+  function sendKernelMessage(sessionGeneration: number) {
+    const listener = mockMessageListeners.get("kernelMessage");
+    if (!listener) {
+      throw new Error("Missing kernel message listener");
+    }
+    const message = {
+      message: "" as IslandsKernelMessage["message"],
+      sessionGeneration,
+    };
+    listener(message as never);
+    return message;
+  }
+
   async function initializeSingleApp() {
     mockSingleApp();
     await bridge.initializeApps();
@@ -281,6 +295,99 @@ describe("IslandsPyodideBridge", () => {
   });
 
   describe("app lifecycle", () => {
+    it("ignores messages after a session yields ownership", async () => {
+      const consumeMessage = vi.fn();
+      bridge.consumeMessages(consumeMessage);
+      mockSingleApp();
+      await bridge.initializeApps();
+
+      const firstMessage = sendKernelMessage(1);
+      let finishReplacement!: () => void;
+      mockReplaceSessionRequest.mockReturnValueOnce(
+        new Promise<void>((resolve) => {
+          finishReplacement = resolve;
+        }),
+      );
+      setSingleApp("generated app 2");
+      const replacement = bridge.initializeApps();
+      await vi.waitFor(() =>
+        expect(mockReplaceSessionRequest).toHaveBeenCalledOnce(),
+      );
+      sendKernelMessage(1);
+      const currentMessage = sendKernelMessage(2);
+
+      expect(consumeMessage.mock.calls).toEqual([
+        [firstMessage],
+        [currentMessage],
+      ]);
+
+      finishReplacement();
+      await replacement;
+    });
+
+    it("restores message ownership when replacement fails", async () => {
+      const consumeMessage = vi.fn();
+      bridge.consumeMessages(consumeMessage);
+      mockSingleApp();
+      await bridge.initializeApps();
+      mockReplaceSessionRequest.mockRejectedValueOnce(
+        new Error("replacement failed"),
+      );
+      setSingleApp("generated app 2");
+
+      await expect(bridge.initializeApps()).rejects.toThrow(
+        "replacement failed",
+      );
+      const previousMessage = sendKernelMessage(1);
+      sendKernelMessage(2);
+
+      expect(consumeMessage.mock.calls).toEqual([[previousMessage]]);
+    });
+
+    it("ignores messages while their session is stopping", async () => {
+      const consumeMessage = vi.fn();
+      bridge.consumeMessages(consumeMessage);
+      mockSingleApp();
+      await bridge.initializeApps();
+      let finishStop!: () => void;
+      mockStopSessionRequest.mockReturnValueOnce(
+        new Promise<void>((resolve) => {
+          finishStop = resolve;
+        }),
+      );
+
+      const stop = bridge.stopSession();
+      await vi.waitFor(() =>
+        expect(mockStopSessionRequest).toHaveBeenCalledOnce(),
+      );
+      sendKernelMessage(1);
+
+      expect(consumeMessage).not.toHaveBeenCalled();
+
+      finishStop();
+      await stop;
+    });
+
+    it("forwards messages from every app in a multi-app document", async () => {
+      const consumeMessage = vi.fn();
+      bridge.consumeMessages(consumeMessage);
+      mockParseMarimoIslandApps.mockReturnValue([app(), app("app-2")]);
+      mockCreateMarimoFile
+        .mockReturnValueOnce("generated app 1")
+        .mockReturnValueOnce("generated app 2");
+      signalWorkerReady();
+
+      await bridge.initializeApps();
+      const firstMessage = sendKernelMessage(1);
+      const secondMessage = sendKernelMessage(2);
+      sendKernelMessage(3);
+
+      expect(consumeMessage.mock.calls).toEqual([
+        [firstMessage],
+        [secondMessage],
+      ]);
+    });
+
     it("holds controls until the first app session is ready", async () => {
       setSingleApp();
 

package/src/core/islands/bridge.ts

@@ -66,6 +66,7 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests {
   private session: AppSession | undefined;
   private sessionReady = new Deferred<AppSession>();
   private nextSessionGeneration = 0;
+  private readonly activeSessionGenerations = new Set<number>();
   private appTransition = Promise.resolve();
   private workerReady = new Deferred<void>();
 
@@ -110,6 +111,11 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests {
     );
 
     this.rpc.addMessageListener("kernelMessage", (message) => {
+      // A stopped session can still have messages in transit after the next
+      // app takes ownership of the shared frontend state.
+      if (!this.activeSessionGenerations.has(message.sessionGeneration)) {
+        return;
+      }
       this.messageConsumer?.(message);
     });
   }
@@ -149,6 +155,7 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests {
       return;
     }
     if (apps.length > 0) {
+      this.activeSessionGenerations.clear();
       this.store.set(
         islandsPendingInitialRunsAtom,
         new Set(apps.map((_, index) => this.nextSessionGeneration + index + 1)),
@@ -165,6 +172,7 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests {
       };
       const previousSession = this.session;
       const replacesSession = managesSingleApp && previousSession !== undefined;
+      this.activeSessionGenerations.add(request.sessionGeneration);
       if (replacesSession) {
         this.store.set(notebookAtom, initialNotebookState());
       }
@@ -186,8 +194,14 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests {
           this.sessionReady.resolve(this.session);
         }
       } catch (error) {
+        this.activeSessionGenerations.delete(request.sessionGeneration);
         if (this.session?.sessionGeneration === request.sessionGeneration) {
           this.session = previousSession;
+          if (previousSession) {
+            this.activeSessionGenerations.add(
+              previousSession.sessionGeneration,
+            );
+          }
         }
         Logger.error(`Failed to start session for app ${app.id}:`, error);
         if (managesSingleApp) {
@@ -203,6 +217,7 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests {
       if (session?.code === undefined || (appId && session.appId !== appId)) {
         return;
       }
+      this.activeSessionGenerations.delete(session.sessionGeneration);
       await this.rpc.proxy.request.stopSession({
         appId: session.appId,
         sessionGeneration: session.sessionGeneration,
@@ -339,6 +354,7 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests {
   sendUpdateFile = throwNotImplemented;
   sendFileDetails = throwNotImplemented;
   openTutorial = throwNotImplemented;
+  getExportAvailability = throwNotImplemented;
   exportAsHTML = throwNotImplemented;
   exportAsIPYNB = throwNotImplemented;
   exportAsMarkdown = throwNotImplemented;