@marimo-team/islands 0.25.3-dev17 → 0.25.3-dev170

package/src/core/islands/worker/worker.tsx

@@ -1,21 +1,17 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import type { PyodideInterface } from "pyodide";
-import type { PyProxy } from "pyodide/ffi";
 import {
   createRPC,
   createRPCRequestHandler,
   createWorkerParentTransport,
   type RPCSchema,
 } from "rpc-anywhere";
-import type { NotificationPayload } from "@/core/kernel/messages";
 import type { ParentSchema } from "@/core/wasm/rpc";
 import { shouldLoadDuckDBPackages } from "@/core/wasm/utils";
 import { TRANSPORT_ID } from "@/core/wasm/worker/constants";
 import { getPyodideVersion } from "@/core/wasm/worker/getPyodideVersion";
-import { MessageBuffer } from "@/core/wasm/worker/message-buffer";
 import type { RawBridge, SerializedBridge } from "@/core/wasm/worker/types";
-import type { JsonString } from "@/utils/json/base64";
 import { Logger } from "@/utils/Logger";
 import { prettyError } from "../../../utils/errors";
 import { invariant } from "../../../utils/invariant";
@@ -45,14 +41,11 @@ async function loadPyodideAndPackages() {
 }
 
 const pyodideReadyPromise = loadPyodideAndPackages();
-export interface IslandsKernelMessage {
-  message: JsonString<NotificationPayload>;
+/** An event of one session's notebook stream. */
+export interface IslandsNotebookEvent {
+  event: string;
   sessionGeneration: number;
 }
-
-const messageBuffer = new MessageBuffer<IslandsKernelMessage>((payload) => {
-  rpc.send.kernelMessage(payload);
-});
 interface SessionRequest {
   code: string;
   appId: string;
@@ -79,21 +72,22 @@ async function startSession(
       code: opts.code,
       filename: `app-${opts.appId}.py`,
     });
-    const nextBridge = await self.controller.startSession({
-      ...notebook,
-      onMessage: (message) => {
-        messageBuffer.push({
-          message,
-          sessionGeneration: opts.sessionGeneration,
-        });
-      },
+    const session = await self.controller.prepareSession(notebook);
+    // Follow before the session instantiates. Each subscription begins with
+    // its own kernel-ready and the session's retained notifications.
+    try {
+      await follow(session.bridge, opts.sessionGeneration);
+    } catch (error) {
+      await session.stop();
+      throw error;
+    }
+    void session.start().catch((error: unknown) => {
+      rpc.send.initializedError({ error: prettyError(error) });
     });
-    if (activeSession) {
-      (nextBridge as unknown as PyProxy).destroy();
-    } else {
+    if (!activeSession) {
       activeSession = {
         appId: opts.appId,
-        bridge: nextBridge,
+        bridge: session.bridge,
         sessionGeneration: opts.sessionGeneration,
       };
     }
@@ -107,10 +101,23 @@ async function startSession(
 }
 
 async function stopActiveSession(): Promise<void> {
-  const session = activeSession;
+  // Stopping closes every session and its notebook subscriptions.
   await self.controller.stopSession();
   activeSession = undefined;
-  (session?.bridge as unknown as PyProxy | undefined)?.destroy();
+}
+
+async function follow(
+  bridge: SerializedBridge,
+  sessionGeneration: number,
+): Promise<void> {
+  const notebook = bridge.notebook;
+  try {
+    await notebook.subscribe(`island-${sessionGeneration}`, null, (event) => {
+      rpc.send.notebookEvent({ event, sessionGeneration });
+    });
+  } finally {
+    notebook.destroy();
+  }
 }
 
 function enqueueSession<T>(operation: () => Promise<T>): Promise<T> {
@@ -243,8 +250,8 @@ export type WorkerSchema = RPCSchema<
     messages: {
       // Emitted when the worker is ready
       ready: {};
-      // Emitted when the kernel sends a message
-      kernelMessage: IslandsKernelMessage;
+      // Emitted for each event of a session's notebook stream
+      notebookEvent: IslandsNotebookEvent;
       // Emitted when the Pyodide is initialized
       initialized: {};
       // Emitted when the Pyodide fails to initialize
@@ -263,13 +270,6 @@ const rpc = createRPC<WorkerSchema, ParentSchema>({
 
 rpc.send("ready", {});
 
-/// Listeners
-// When the consumer is ready, start the message buffer
-rpc.addMessageListener("consumerReady", async () => {
-  await pyodideReadyPromise; // Make sure loading is done
-  messageBuffer.start();
-});
-
 function getMarimoVersion() {
   return self.name; // We store the version in the worker name
 }

package/src/core/kernel/__tests__/RuntimeState.test.ts

@@ -25,7 +25,7 @@ describe("RuntimeState", () => {
 
   beforeEach(() => {
     vi.clearAllMocks();
-    mockSendComponentValues = vi.fn((args) => Promise.resolve(null));
+    mockSendComponentValues = vi.fn(() => Promise.resolve(null));
     uiElementRegistry = UIElementRegistry.INSTANCE;
     uiElementRegistry.entries.clear();
 

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

@@ -3,7 +3,7 @@
 import { afterEach, describe, expect, it } from "vitest";
 import { initialModeAtom } from "@/core/mode";
 import { store } from "@/core/state/jotai";
-import { ConnectionState } from "../../websocket/types";
+import { ConnectionState } from "../../session/types";
 import { connectionAtom, waitForConnectionOpenIfNotebook } from "../connection";
 
 const NEVER = Symbol("never");

package/src/core/network/connection-notice.ts

@@ -6,7 +6,7 @@ import {
   type ConnectionPhase,
   WebSocketClosedReason,
   ConnectionState,
-} from "@/core/websocket/types";
+} from "@/core/session/types";
 import { connectionAtom } from "./connection";
 
 export interface ConnectionNotice {

package/src/core/network/connection.ts

@@ -4,11 +4,13 @@ import { atomWithReducer } from "jotai/utils";
 import type { NotificationMessageData } from "../kernel/messages";
 import { isNotebookPage } from "../mode";
 import { waitFor } from "../state/jotai";
-import { type ConnectionStatus, ConnectionState } from "../websocket/types";
+import { type ConnectionStatus, ConnectionState } from "../session/types";
 
 /**
- * Atom for storing the connection status.
- * Initialized to NOT_STARTED.
+ * Shared application state, not the actual event-stream status.
+ *
+ * TODO(MO-8051) Expose read-only connection state and derive UI availability
+ * separately.
  */
 export const connectionAtom = atom<ConnectionStatus>({
   state: ConnectionState.NOT_STARTED,

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

@@ -7,7 +7,7 @@ import { memoizeLastValue } from "@/utils/once";
 import { waitForKernelToBeInstantiated } from "../kernel/state";
 import type { RuntimeManager } from "../runtime/runtime";
 import { store } from "../state/jotai";
-import { ConnectionState } from "../websocket/types";
+import { ConnectionState } from "../session/types";
 import { connectionAtom, waitForConnectionOpen } from "./connection";
 import type { EditRequests, RunRequests } from "./types";
 
@@ -144,7 +144,6 @@ const ACTIONS: Record<keyof AllRequests, Action> = {
   downloadStorage: "waitForConnectionOpen",
 
   // Sync operations that wait for connection
-  sendDocumentTransaction: "waitForConnectionOpen",
   sendCodeCompletionRequest: "waitForConnectionOpen",
 };
 

package/src/core/network/requests-network.ts

@@ -70,15 +70,6 @@ export function createNetworkRequests(): EditRequests & RunRequests {
         })
         .then(handleResponseReturnNull);
     },
-    sendDocumentTransaction: async (request) => {
-      await waitForConnectionOpen();
-      return getClient()
-        .POST("/api/document/transaction", {
-          body: request,
-          params: getParams(),
-        })
-        .then(handleResponseReturnNull);
-    },
     sendRename: (request) => {
       return getClient()
         .POST("/api/kernel/rename", {

package/src/core/network/requests-static.ts

@@ -40,7 +40,6 @@ export function createStaticRequests(): EditRequests & RunRequests {
       return null;
     },
     sendRestart: throwNotInEditMode,
-    sendDocumentTransaction: throwNotInEditMode,
     sendRun: throwNotInEditMode,
     sendRunScratchpad: throwNotInEditMode,
     sendRename: throwNotInEditMode,

package/src/core/network/requests-toasting.tsx

@@ -21,7 +21,6 @@ export function createErrorToastingRequests(
     sendInstantiate: "Failed to instantiate",
     sendFunctionRequest: "Failed to send function request",
     sendRestart: "Failed to restart",
-    sendDocumentTransaction: "Failed to sync document transaction",
     sendRun: "Failed to run",
     sendRunScratchpad: "Failed to run scratchpad",
     sendRename: "Failed to rename",

package/src/core/network/types.ts

@@ -117,8 +117,6 @@ export interface FileCreateInput {
 }
 export type UpdateUIElementRequest = schemas["UpdateUIElementRequest"];
 export type ModelRequest = schemas["ModelRequest"];
-export type NotebookDocumentTransactionRequest =
-  schemas["NotebookDocumentTransactionRequest"];
 export type UpdateUserConfigRequest = schemas["UpdateUserConfigRequest"];
 export type ShutdownSessionRequest = schemas["ShutdownSessionRequest"];
 export type Snippet = schemas["Snippet"];
@@ -179,9 +177,6 @@ export interface EditRequests {
   saveAppConfig: (request: SaveAppConfigurationRequest) => Promise<null>;
   saveCellConfig: (request: SetCellConfigRequest) => Promise<null>;
   sendRestart: () => Promise<null>;
-  sendDocumentTransaction: (
-    request: NotebookDocumentTransactionRequest,
-  ) => Promise<null>;
   sendInstallMissingPackages: (
     request: InstallPackagesRequest,
   ) => Promise<null>;

package/src/core/network/useConnectionNotice.ts

@@ -17,11 +17,17 @@ export function useConnectionNotice(
   const notice = useAtomValue(connectionNoticeAtom);
   const connected = useAtomValue(isConnectedAtom);
   const elapsed = useDelayElapsed(notice?.pending ? delayMs : 0);
-  const [previous, setPrevious] = useState<ConnectionNotice | null>(null);
+  const [previous, setPrevious] = useState<ConnectionNotice | null>(
+    elapsed ? notice : null,
+  );
+
+  const visibleNotice = elapsed ? notice : null;
+  if ((notice || !connected) && previous !== visibleNotice) {
+    setPrevious(visibleNotice);
+  }
 
   useEffect(() => {
     if (notice || !connected) {
-      setPrevious(elapsed ? notice : null);
       return;
     }
     const timeout = setTimeout(() => setPrevious(null), 1500);

package/src/core/notebook/attachments.ts

@@ -0,0 +1,38 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { dequal } from "dequal";
+import { atom, useAtomValue } from "jotai";
+import { selectAtom } from "jotai/vanilla/utils";
+import { useMemo } from "react";
+import type { Attachment } from "@marimo-team/client/protocol";
+import { activeNotebookReplicaAtom } from "../session/active-notebook";
+
+export type { Attachment } from "@marimo-team/client/protocol";
+
+const EMPTY_ATTACHMENTS: readonly Attachment[] = [];
+
+/** Attachments visible to this session, empty while no notebook is available. */
+export const notebookAttachmentsAtom = atom((get) => {
+  const replica = get(activeNotebookReplicaAtom);
+  return replica ? get(replica.attachmentsAtom) : EMPTY_ATTACHMENTS;
+});
+
+/** Current notebook attachments, optionally restricted to one kind. */
+export function useNotebookAttachments(kind?: Attachment["kind"]) {
+  const attachments = useMemo(
+    () =>
+      kind === undefined
+        ? notebookAttachmentsAtom
+        : selectAtom(
+            notebookAttachmentsAtom,
+            (value) => value.filter((attachment) => attachment.kind === kind),
+            dequal,
+          ),
+    [kind],
+  );
+  return useAtomValue(attachments);
+}
+
+export function attachmentLabel(attachment: Attachment): string {
+  return attachment.name || "Agent";
+}