@marimo-team/islands 0.25.3-dev20 → 0.25.3-dev23

package/src/core/websocket/__tests__/useNotebookSession.hook.test.tsx

@@ -41,11 +41,6 @@ import { Logger } from "@/utils/Logger";
 import { Deferred } from "@/utils/Deferred";
 import { cellId } from "@/__tests__/branded";
 import { alertAtom, getPackageAlert } from "@/core/alerts/state";
-import {
-  type Attachment,
-  attachmentsAtom,
-  agentAttachmentsAtom,
-} from "@/core/attachments/state";
 import { notebookAtom } from "@/core/cells/cells";
 import { AppConfigSchema } from "@/core/config/config-schema";
 import { ConnectionNotice } from "@/components/editor/alerts/connection-notice";
@@ -239,28 +234,6 @@ describe("useNotebookSession messages", () => {
     });
   }
 
-  it("replaces attachment snapshots, including the last detach", () => {
-    const agent: Attachment = {
-      id: "a1",
-      kind: "agent",
-      name: "Pi",
-      since: "2026-10-14T00:00:00Z",
-      notebook_id: "nb-test",
-      identity: "pi-thread",
-    };
-    const client: Attachment = {
-      id: "c1",
-      kind: "client",
-      name: null,
-      since: "2026-10-14T00:00:01Z",
-      notebook_id: "nb-test",
-    };
-    for (const attachments of [[agent, client], [client], []]) {
-      receive({ data: { op: "attachments", attachments } });
-      expect(store.get(attachmentsAtom)).toEqual(attachments);
-    }
-  });
-
   it.each(["{broken", "null"])(
     "reports a malformed legacy frame and continues applying notifications: %s",
     (data) => {
@@ -408,150 +381,6 @@ it.each(["kernel-ready", "reconnected"])(
   },
 );
 
-describe("attachment snapshots across browser connections", () => {
-  const agent: Attachment = {
-    id: "a1",
-    kind: "agent",
-    name: "Pi",
-    since: "2026-10-14T00:00:00Z",
-    notebook_id: "nb-test",
-    identity: "pi-thread",
-  };
-
-  function connectedSession() {
-    const store = createStore();
-    store.set(connectionAtom, { state: ConnectionState.OPEN });
-    store.set(attachmentsAtom, [agent]);
-    const transport = makeTransport(WebSocket.OPEN);
-    vi.mocked(createConnectionTransport).mockReturnValue(transport);
-    vi.mocked(useRuntimeManager).mockReturnValue(makeRuntimeManager());
-    renderConnectionHook(store);
-    const options = getTransportHandlers();
-    const receive = (data: NotificationPayload["data"]) =>
-      act(() => {
-        options.onMessage(
-          new MessageEvent("message", {
-            data: JSON.stringify({ op: data.op, data }),
-          }),
-        );
-      });
-    return { store, transport, options, receive };
-  }
-
-  it.each([
-    "MARIMO_NO_FILE_KEY",
-    "MARIMO_NO_SESSION_ID",
-    "MARIMO_NO_SESSION",
-    "MARIMO_SHUTDOWN",
-    "MARIMO_KERNEL_STARTUP_ERROR",
-    "MARIMO_UNAUTHORIZED",
-    "MARIMO_KIOSK_NOT_ALLOWED",
-    "MARIMO_TRANSPORT_EXHAUSTED",
-  ])("clears cached attachments on %s", (reason) => {
-    const { store, transport, options } = connectedSession();
-    act(() => options.onClose(new CloseEvent("close", { reason })));
-    expect(store.get(attachmentsAtom)).toEqual([]);
-    expect(store.get(agentAttachmentsAtom)).toEqual([]);
-    expect(store.get(connectionAtom).state).toBe(ConnectionState.CLOSED);
-    expect(transport.reconnect).not.toHaveBeenCalled();
-  });
-
-  it("clears the snapshot on transport error", () => {
-    const { store, options } = connectedSession();
-    act(() => options.onError(new Event("error")));
-    expect(store.get(attachmentsAtom)).toEqual([]);
-    expect(store.get(agentAttachmentsAtom)).toEqual([]);
-  });
-
-  it("waits for a fresh snapshot after a transient reconnect", async () => {
-    const { store, transport, options, receive } = connectedSession();
-    act(() => options.onClose(new CloseEvent("close")));
-    expect(store.get(attachmentsAtom)).toEqual([]);
-    expect(store.get(connectionAtom).state).toBe(ConnectionState.CONNECTING);
-    expect(transport.reconnect).toHaveBeenCalledOnce();
-
-    await act(async () => options.onOpen(new Event("open")));
-    receive({ op: "reconnected" });
-    expect(store.get(connectionAtom).state).toBe(ConnectionState.OPEN);
-    expect(store.get(agentAttachmentsAtom)).toEqual([]);
-
-    const refreshed = { ...agent, id: "a2", name: "Codex" };
-    receive({ op: "attachments", attachments: [refreshed] });
-    expect(store.get(agentAttachmentsAtom)).toEqual([refreshed]);
-    receive({ op: "attachments", attachments: [] });
-    expect(store.get(agentAttachmentsAtom)).toEqual([]);
-  });
-
-  it("keeps a fresh snapshot hidden until the session is ready", async () => {
-    const { store, options, receive } = connectedSession();
-    act(() => options.onClose(new CloseEvent("close")));
-    await act(async () => options.onOpen(new Event("open")));
-    receive({ op: "attachments", attachments: [agent] });
-    expect(store.get(attachmentsAtom)).toEqual([agent]);
-    expect(store.get(agentAttachmentsAtom)).toEqual([]);
-    receive({ op: "reconnected" });
-    expect(store.get(agentAttachmentsAtom)).toEqual([agent]);
-  });
-
-  it("invalidates an old snapshot when a new transport opens", async () => {
-    const { store, options } = connectedSession();
-    await act(async () => options.onOpen(new Event("open")));
-    expect(store.get(attachmentsAtom)).toEqual([]);
-    expect(store.get(agentAttachmentsAtom)).toEqual([]);
-  });
-});
-
-it.each([true, false])(
-  "keeps attachments only when kernel-ready resumes the session: %s",
-  (resumed) => {
-    const store = createStore();
-    const attachments: Attachment[] = [
-      {
-        id: "a1",
-        kind: "agent",
-        name: "Pi",
-        since: "2026-10-14T00:00:00Z",
-        notebook_id: "nb-test",
-        identity: "pi-thread",
-      },
-    ];
-    store.set(attachmentsAtom, attachments);
-    vi.mocked(createConnectionTransport).mockReturnValue(
-      makeTransport(WebSocket.OPEN),
-    );
-    vi.mocked(useRuntimeManager).mockReturnValue(makeRuntimeManager());
-    renderConnectionHook(store);
-    const options = getTransportHandlers();
-    act(() =>
-      options.onMessage(
-        new MessageEvent("message", {
-          data: JSON.stringify({
-            op: "kernel-ready",
-            data: {
-              op: "kernel-ready",
-              cell_ids: [],
-              codes: [],
-              names: [],
-              configs: [],
-              layout: null,
-              resumed,
-              ui_values: {},
-              last_executed_code: {},
-              last_execution_time: {},
-              app_config: { width: "normal" },
-              kiosk: false,
-              capabilities: { terminal: false },
-              auto_instantiated: true,
-              consumer_capabilities: { edit: true, interact: true },
-            },
-          }),
-        }),
-      ),
-    );
-    expect(store.get(attachmentsAtom)).toEqual(resumed ? attachments : []);
-  },
-);
-
 describe("connection notice", () => {
   beforeEach(() => vi.useFakeTimers());
   afterEach(() => vi.useRealTimers());

package/src/core/websocket/state.ts

@@ -0,0 +1,7 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { atom } from "jotai";
+import type { NotebookClient } from "../host/notebook-client";
+
+/** The client owned by the mounted notebook session. */
+export const notebookClientAtom = atom<NotebookClient | null>(null);

package/src/core/websocket/useNotebookClient.ts

@@ -14,6 +14,7 @@ import {
   type NotebookConnection,
 } from "../host/notebook-client";
 import { classifyCloseEvent } from "./connection-state";
+import { notebookClientAtom } from "./state";
 import {
   createConnectionTransport,
   type ConnectionTransportOptions,
@@ -41,6 +42,7 @@ export function useNotebookClient(options: Options): NotebookClient {
 
   useEffect(() => {
     const lifetime = new AbortController();
+    store.set(notebookClientAtom, controller.client);
     controller.connect({
       createTransport: (url) =>
         createConnectionTransport({
@@ -61,9 +63,15 @@ export function useNotebookClient(options: Options): NotebookClient {
       decodeJson: jsonParseWithSpecialChar,
       classifyClose: classifyCloseEvent,
     });
-    return () => lifetime.abort();
+    return () => {
+      lifetime.abort();
+      if (store.get(notebookClientAtom) === controller.client) {
+        store.set(notebookClientAtom, null);
+      }
+    };
   }, [
     controller,
+    store,
     sessionId,
     isStatic,
     transportType,

package/src/core/websocket/useNotebookSession.tsx

@@ -5,7 +5,6 @@ import { useAtom, useAtomValue, useSetAtom, useStore } from "jotai";
 import { useLayoutEffect } from "react";
 import { useErrorBoundary } from "react-error-boundary";
 import { toast } from "@/components/ui/use-toast";
-import { attachmentsAtom } from "@/core/attachments/state";
 import { getNotebook, useCellActions } from "@/core/cells/cells";
 import { applyTransactionChanges } from "@/core/cells/document-changes";
 import { AUTOCOMPLETER } from "@/core/codemirror/completion/Autocompleter";
@@ -140,7 +139,6 @@ export function useNotebookSession(opts: {
   const transportType = useAtomValue(connectionTransportTypeAtom);
   const setCacheInfo = useSetAtom(cacheInfoAtom);
   const setKernelStartupError = useSetAtom(kernelStartupErrorAtom);
-  const setAttachments = useSetAtom(attachmentsAtom);
   const {
     setNamespaces: setStorageNamespaces,
     filterFromVariables: filterStorageFromVariables,
@@ -148,9 +146,6 @@ export function useNotebookSession(opts: {
 
   const handleNotification = (msg: NotificationPayload) => {
     switch (msg.data.op) {
-      case "attachments":
-        setAttachments(msg.data.attachments);
-        return;
       case "reload":
         reloadSafe();
         return;
@@ -185,7 +180,6 @@ export function useNotebookSession(opts: {
           existingCells,
         });
         if (!msg.data.resumed) {
-          setAttachments([]);
           // A freshly started kernel may expose a different environment
           // (e.g. new env vars); re-run discovery instead of reusing stale
           // suggestions from the previous kernel instance.
@@ -452,7 +446,6 @@ export function useNotebookSession(opts: {
       if (status.state === ConnectionState.NOT_STARTED) {
         return;
       }
-      setAttachments([]);
       if (status.state === ConnectionState.OPEN) {
         updateStartupProgress(null);
         return; // kernel-ready, not the socket, establishes session readiness.
@@ -487,7 +480,6 @@ export function useNotebookSession(opts: {
     notebook,
     store,
     setConnection,
-    setAttachments,
     setKernelStartupError,
     updateStartupProgress,
   ]);

package/src/utils/once.ts

@@ -2,36 +2,36 @@
 
 import { arrayShallowEquals } from "./arrays";
 
-// oxlint-disable-next-line typescript/no-explicit-any
-export function once<T extends (...args: any[]) => any>(fn: T): T {
-  let result: ReturnType<T>;
+export function once<F extends (...args: never[]) => unknown>(fn: F): F;
+export function once<Args extends unknown[], Result, This>(
+  fn: (this: This, ...args: Args) => Result,
+): (this: This, ...args: Args) => Result {
+  let result: Result;
   let called = false;
-  return function (
-    this: ThisParameterType<T>,
-    ...args: Parameters<T>
-  ): ReturnType<T> {
+  return function (this: This, ...args: Args): Result {
     if (!called) {
       called = true;
-      result = fn.apply(this, args) as ReturnType<T>;
+      result = fn.apply(this, args);
     }
     return result;
-  } as T;
+  };
 }
 
-// oxlint-disable-next-line typescript/no-explicit-any
-export function memoizeLastValue<T extends (...args: any[]) => any>(fn: T): T {
-  let result: ReturnType<T>;
-  let lastArgs: Parameters<T> | undefined;
+export function memoizeLastValue<F extends (...args: never[]) => unknown>(
+  fn: F,
+): F;
+export function memoizeLastValue<Args extends unknown[], Result, This>(
+  fn: (this: This, ...args: Args) => Result,
+): (this: This, ...args: Args) => Result {
+  let result: Result;
+  let lastArgs: Args | undefined;
   let lastError: unknown;
   let hasError = false;
 
-  return function (
-    this: ThisParameterType<T>,
-    ...args: Parameters<T>
-  ): ReturnType<T> {
+  return function (this: This, ...args: Args): Result {
     if (lastArgs === undefined || !arrayShallowEquals(args, lastArgs)) {
       try {
-        result = fn.apply(this, args) as ReturnType<T>;
+        result = fn.apply(this, args);
         hasError = false;
         lastError = undefined;
       } catch (error) {
@@ -45,5 +45,5 @@ export function memoizeLastValue<T extends (...args: any[]) => any>(fn: T): T {
       throw lastError;
     }
     return result;
-  } as T;
+  };
 }