@marimo-team/islands 0.25.1-dev2 → 0.25.1-dev4

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

@@ -1,10 +1,119 @@
 /* Copyright 2026 Marimo. All rights reserved. */
-import { describe, expect, it } from "vitest";
+import { afterEach, describe, expect, it, vi } from "vitest";
+import type { AppMode } from "@/core/mode";
 import { generateSessionId, isSessionId } from "../session";
 
+const URL_SESSION_ID = "s_abc123";
+
+/**
+ * Load a fresh copy of the session module against the given URL and mode.
+ * The module caches its session id, so each case needs its own instance.
+ */
+async function loadSession(opts: { mode: AppMode | undefined; url: string }) {
+  vi.resetModules();
+  window.history.replaceState({}, "", opts.url);
+  const { store } = await import("@/core/state/jotai");
+  const { initialModeAtom } = await import("@/core/mode");
+  store.set(initialModeAtom, opts.mode);
+  return await import("../session");
+}
+
+function currentSearchParams(): URLSearchParams {
+  return new URL(window.location.href).searchParams;
+}
+
 describe("Session", () => {
+  afterEach(() => {
+    window.history.replaceState({}, "", "/");
+  });
+
   it("should create a session", () => {
     const id = generateSessionId();
     expect(isSessionId(id)).toBe(true);
   });
+
+  it("resumes the session named in the URL in edit mode", async () => {
+    const { getSessionId } = await loadSession({
+      mode: "edit",
+      url: `/?session_id=${URL_SESSION_ID}`,
+    });
+    expect(getSessionId()).toBe(URL_SESSION_ID);
+    expect(currentSearchParams().has("session_id")).toBe(false);
+  });
+
+  it("keeps the URL session id in kiosk mode so a refresh resumes it", async () => {
+    const { getSessionId } = await loadSession({
+      mode: "edit",
+      url: `/?kiosk=true&session_id=${URL_SESSION_ID}`,
+    });
+    expect(getSessionId()).toBe(URL_SESSION_ID);
+    expect(currentSearchParams().get("session_id")).toBe(URL_SESSION_ID);
+    expect(currentSearchParams().get("kiosk")).toBe("true");
+  });
+
+  it("keeps the kiosk flag in edit mode when the URL names no session", async () => {
+    const { getSessionId } = await loadSession({
+      mode: "edit",
+      url: "/?kiosk=true",
+    });
+    expect(isSessionId(getSessionId())).toBe(true);
+    expect(currentSearchParams().get("kiosk")).toBe("true");
+  });
+
+  it("starts a new session when the URL session id is malformed", async () => {
+    const { getSessionId } = await loadSession({
+      mode: "edit",
+      url: "/?session_id=not-a-session",
+    });
+    const id = getSessionId();
+    expect(isSessionId(id)).toBe(true);
+    expect(id).not.toBe("not-a-session");
+  });
+
+  it("ignores the URL session id in read mode", async () => {
+    const { getSessionId } = await loadSession({
+      mode: "read",
+      url: `/?session_id=${URL_SESSION_ID}`,
+    });
+    const id = getSessionId();
+    expect(isSessionId(id)).toBe(true);
+    expect(id).not.toBe(URL_SESSION_ID);
+    // The param is still scrubbed so a refresh cannot reintroduce it.
+    expect(currentSearchParams().has("session_id")).toBe(false);
+  });
+
+  it("drops both the session id and the kiosk flag in read mode", async () => {
+    const { getSessionId } = await loadSession({
+      mode: "read",
+      url: `/?kiosk=true&session_id=${URL_SESSION_ID}`,
+    });
+    expect(getSessionId()).not.toBe(URL_SESSION_ID);
+    expect(currentSearchParams().has("session_id")).toBe(false);
+    // The kiosk flag is forwarded to the websocket URL, and the server
+    // refuses kiosk connections to an app, so it must not survive here.
+    expect(currentSearchParams().has("kiosk")).toBe(false);
+  });
+
+  it("drops a lone kiosk flag in read mode", async () => {
+    const { getSessionId } = await loadSession({
+      mode: "read",
+      url: "/?kiosk=true",
+    });
+    expect(isSessionId(getSessionId())).toBe(true);
+    expect(currentSearchParams().has("kiosk")).toBe(false);
+  });
+
+  it("ignores the URL session id when the mode is not yet known", async () => {
+    const { getSessionId } = await loadSession({
+      mode: undefined,
+      url: `/?session_id=${URL_SESSION_ID}`,
+    });
+    expect(getSessionId()).not.toBe(URL_SESSION_ID);
+    expect(currentSearchParams().has("session_id")).toBe(false);
+  });
+
+  it("returns the same id on every call", async () => {
+    const { getSessionId } = await loadSession({ mode: "read", url: "/" });
+    expect(getSessionId()).toBe(getSessionId());
+  });
 });

package/src/core/kernel/session.ts

@@ -4,6 +4,8 @@ import { Logger } from "@/utils/Logger";
 import type { TypedString } from "@/utils/typed";
 import { updateQueryParams } from "@/utils/urls";
 import { KnownQueryParams } from "../constants";
+import { initialModeAtom } from "../mode";
+import { store } from "../state/jotai";
 
 export type SessionId = TypedString<"SessionId">;
 
@@ -20,31 +22,60 @@ export function isSessionId(value: string | null): value is SessionId {
   return /^s_[\da-z]{6}$/.test(value);
 }
 
-const sessionId = (() => {
+let sessionId: SessionId | null = null;
+
+function resolveSessionId(): SessionId {
+  // Only the editor lets the URL steer the connection: kiosk mode and
+  // external editors reattach to a running session through `session_id`.
+  // An app served with `marimo run` never does. There, the session id is the
+  // only thing that routes a browser to its kernel, so a link that picked it
+  // would let the link's author attach to the reader's session afterwards,
+  // and a stray `kiosk` flag would make the server refuse the connection.
+  const urlMaySteerConnection = store.get(initialModeAtom) === "edit";
+
   const url = new URL(window.location.href);
   const id = url.searchParams.get(
     KnownQueryParams.sessionId,
   ) as SessionId | null;
-  if (isSessionId(id)) {
-    // Remove the session_id from the URL
-    updateQueryParams((params) => {
-      // Keep the session_id if we are in kiosk mode
-      // this is so we can resume the same session if the user refreshes the page
-      if (params.has(KnownQueryParams.kiosk)) {
-        return;
-      }
+  const urlSessionId = isSessionId(id) ? id : null;
+
+  updateQueryParams((params) => {
+    if (!urlMaySteerConnection) {
       params.delete(KnownQueryParams.sessionId);
+      params.delete(KnownQueryParams.kiosk);
+      return;
+    }
+    // Keep the session_id in kiosk mode so that a refresh resumes the same
+    // session.
+    if (params.has(KnownQueryParams.kiosk)) {
+      return;
+    }
+    params.delete(KnownQueryParams.sessionId);
+  });
+
+  if (urlSessionId === null) {
+    Logger.debug("Starting a new session", { sessionId: id });
+    return generateSessionId();
+  }
+  if (!urlMaySteerConnection) {
+    Logger.debug("Ignoring session_id from URL outside the editor", {
+      sessionId: urlSessionId,
     });
-    Logger.debug("Connecting to existing session", { sessionId: id });
-    return id;
+    return generateSessionId();
   }
-  Logger.debug("Starting a new session", { sessionId: id });
-  return generateSessionId();
-})();
+  Logger.debug("Connecting to existing session", { sessionId: urlSessionId });
+  return urlSessionId;
+}
 
 /**
- * Resume an existing session or start a new one
+ * Resume an existing session or start a new one.
+ *
+ * Resolved on first use rather than at import time, because the URL is only
+ * trusted once the app mode is known.
  */
 export function getSessionId(): SessionId {
+  if (sessionId === null) {
+    sessionId = resolveSessionId();
+  }
   return sessionId;
 }