@marimo-team/islands 0.25.3-dev129 → 0.25.3-dev134

package/src/components/tracing/tracing-keyboard.test.tsx

@@ -0,0 +1,85 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { act, fireEvent, render, screen } from "@testing-library/react";
+import { createStore, Provider } from "jotai";
+import { describe, expect, it, vi } from "vitest";
+import { TooltipProvider } from "@/components/ui/tooltip";
+import { type Run, type RunId, runsAtom } from "@/core/cells/runs";
+import { Tracing } from "./tracing";
+
+vi.mock("react-vega", () => ({ useVegaEmbed: () => undefined }));
+vi.mock("vega-lite", () => ({ compile: () => ({ spec: {} }) }));
+vi.mock("@/components/editor/chrome/panels/panel-context", () => ({
+  usePanelSection: () => "sidebar",
+}));
+
+const run: Run = {
+  runId: "run-1" as RunId,
+  runStartTime: 1_704_067_200,
+  cellRuns: new Map(),
+};
+
+describe("trace disclosure keyboard focus", () => {
+  it.each([false, true])(
+    "keeps the header focused across toggles (inline=%s)",
+    (inline) => {
+      const store = createStore();
+      store.set(runsAtom, {
+        runIds: [run.runId],
+        runMap: new Map([[run.runId, run]]),
+      });
+      render(
+        <Provider store={store}>
+          <TooltipProvider>
+            <Tracing />
+          </TooltipProvider>
+        </Provider>,
+      );
+      if (inline) {
+        fireEvent.click(screen.getByRole("checkbox", { name: "Inline chart" }));
+      }
+      const header = screen.getByRole("button", { name: /^Run -/ });
+      act(() => header.focus());
+      for (const expanded of [false, true, false]) {
+        // Native buttons emit a click with detail=0 for keyboard activation.
+        fireEvent.click(header, { detail: 0 });
+        expect(screen.getByRole("button", { name: /^Run -/ })).toBe(header);
+        expect(header).toHaveFocus();
+        expect(header).toHaveAttribute("aria-expanded", String(expanded));
+      }
+    },
+  );
+
+  it("does not steal focus when a newer run expands automatically", () => {
+    const store = createStore();
+    store.set(runsAtom, {
+      runIds: [run.runId],
+      runMap: new Map([[run.runId, run]]),
+    });
+    render(
+      <Provider store={store}>
+        <TooltipProvider>
+          <Tracing />
+        </TooltipProvider>
+      </Provider>,
+    );
+    const header = screen.getByRole("button", { name: /^Run -/ });
+    act(() => header.focus());
+    const newer: Run = {
+      ...run,
+      runId: "run-2" as RunId,
+      runStartTime: run.runStartTime + 1,
+    };
+    act(() => {
+      store.set(runsAtom, {
+        runIds: [newer.runId, run.runId],
+        runMap: new Map([
+          [run.runId, run],
+          [newer.runId, newer],
+        ]),
+      });
+    });
+    expect(header).toHaveFocus();
+    expect(header).toHaveAttribute("aria-expanded", "false");
+  });
+});

package/src/components/tracing/tracing.tsx

@@ -52,6 +52,7 @@ export const Tracing: React.FC = () => {
 
   const { theme } = useTheme();
   const [chartPosition, setChartPosition] = useState<ChartPosition>("above");
+  const chartPositionId = React.useId();
   const panelSection = usePanelSection();
 
   const toggleChartPosition = () => {
@@ -76,11 +77,12 @@ export const Tracing: React.FC = () => {
     <div className="py-1 px-2 overflow-y-scroll h-full">
       <div className="flex flex-row justify-start gap-3">
         <div className="flex flex-row gap-1 items-center">
-          <label htmlFor="chartPosition" className="text-xs">
+          <label htmlFor={chartPositionId} className="text-xs">
             Inline chart
           </label>
           <input
             type="checkbox"
+            id={chartPositionId}
             name="chartPosition"
             data-testid="chartPosition"
             onClick={toggleChartPosition}
@@ -163,28 +165,31 @@ const TraceBlock: React.FC<{
   const chevron = <Icon height={16} className="inline" />;
 
   const traceTitle = (
-    <span className="text-sm cursor-pointer" onClick={onToggleExpanded}>
+    <button
+      type="button"
+      aria-expanded={isExpanded}
+      className="text-sm cursor-pointer"
+      onClick={onToggleExpanded}
+    >
       Run - {formatLogTimestamp(run.runStartTime)}
       {chevron}
-    </span>
+    </button>
   );
 
-  if (!isExpanded) {
-    return (
-      <div key={run.runId} className="flex flex-col">
-        <pre className="font-mono font-semibold">{traceTitle}</pre>
-      </div>
-    );
-  }
-
   return (
-    <TraceBlockBody
-      key={run.runId}
-      run={run}
-      chartPosition={chartPosition}
-      theme={theme}
-      title={traceTitle}
-    />
+    <div
+      className={cn(
+        "grid grid-cols-1",
+        isExpanded &&
+          chartPosition === "sideBySide" &&
+          "grid-cols-[auto_minmax(0,1fr)] grid-rows-[auto_1fr]",
+      )}
+    >
+      <pre className="font-mono font-semibold">{traceTitle}</pre>
+      {isExpanded && (
+        <TraceBlockBody run={run} chartPosition={chartPosition} theme={theme} />
+      )}
+    </div>
   );
 };
 
@@ -192,8 +197,7 @@ const TraceBlockBody: React.FC<{
   run: Run;
   chartPosition: ChartPosition;
   theme: ResolvedTheme;
-  title: React.ReactNode;
-}> = ({ run, chartPosition, theme, title }) => {
+}> = ({ run, chartPosition, theme }) => {
   const [hoveredCellId, setHoveredCellId] = useState<CellId | null>();
   const vegaRef = useRef<HTMLDivElement>(null);
   const { ref, width = 300 } = useResizeObserver<HTMLDivElement>();
@@ -264,16 +268,22 @@ const TraceBlockBody: React.FC<{
   }, [embed]);
 
   const traceRows = (
-    <TraceRows
-      run={run}
-      hoveredCellId={hoveredCellId}
-      hiddenInputElementId={hiddenInputElementId}
-    />
+    <div className="font-mono font-semibold">
+      <TraceRows
+        run={run}
+        hoveredCellId={hoveredCellId}
+        hiddenInputElementId={hiddenInputElementId}
+      />
+    </div>
   );
 
   const chartElement = (
     <div
-      className={chartPosition === "sideBySide" ? "-mt-0.5 flex-1" : ""}
+      className={
+        chartPosition === "sideBySide"
+          ? "-mt-0.5 col-start-2 row-start-1 row-span-2 min-w-0"
+          : "min-w-0"
+      }
       ref={ref}
     >
       <Suspense>
@@ -284,24 +294,18 @@ const TraceBlockBody: React.FC<{
 
   if (chartPosition === "above") {
     return (
-      <div key={run.runId} className="flex flex-col">
-        <pre className="font-mono font-semibold">
-          {title}
-          {chartElement}
-          {traceRows}
-        </pre>
-      </div>
+      <>
+        {chartElement}
+        {traceRows}
+      </>
     );
   }
 
   return (
-    <div key={run.runId} className="flex flex-row">
-      <pre className="font-mono font-semibold">
-        {title}
-        {traceRows}
-      </pre>
+    <>
+      {traceRows}
       {chartElement}
-    </div>
+    </>
   );
 };
 

package/src/components/ui/__tests__/draggable-popover.test.tsx

@@ -0,0 +1,24 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { fireEvent, render, screen } from "@testing-library/react";
+import { expect, it } from "vitest";
+import { DraggablePopover } from "../draggable-popover";
+
+it("moves from the keyboard handle without intercepting child input arrows", () => {
+  render(
+    <DraggablePopover open={true}>
+      <input aria-label="Filter" />
+    </DraggablePopover>,
+  );
+  const handle = screen.getByRole("button", {
+    name: "Move popup with arrow keys",
+  });
+  const popup = screen.getByRole("dialog");
+  fireEvent.keyDown(handle, { key: "ArrowRight" });
+  fireEvent.keyDown(handle, { key: "ArrowDown" });
+  expect(popup).toHaveStyle({ left: "10px", top: "10px" });
+  expect(
+    fireEvent.keyDown(screen.getByRole("textbox"), { key: "ArrowLeft" }),
+  ).toBe(true);
+  expect(popup).toHaveStyle({ left: "10px", top: "10px" });
+});

package/src/components/ui/draggable-popover.tsx

@@ -53,14 +53,36 @@ export const DraggablePopover = ({
           top: position.y,
         }}
       >
-        <div
+        <button
+          type="button"
+          aria-label="Move popup with arrow keys"
           onMouseDown={handleMouseDown}
+          onKeyDown={(e) => {
+            const directions: Record<
+              string,
+              { x: number; y: number } | undefined
+            > = {
+              ArrowLeft: { x: -10, y: 0 },
+              ArrowRight: { x: 10, y: 0 },
+              ArrowUp: { x: 0, y: -10 },
+              ArrowDown: { x: 0, y: 10 },
+            };
+            const offset = directions[e.key];
+            if (offset) {
+              e.preventDefault();
+              e.stopPropagation();
+              setPosition((previous) => ({
+                x: previous.x + offset.x,
+                y: previous.y + offset.y,
+              }));
+            }
+          }}
           className={`flex items-center justify-center absolute top-0 left-1/2 -translate-x-1/2 ${
             isDragging ? "cursor-grabbing" : "cursor-grab"
           }`}
         >
           <GripHorizontalIcon className="h-5 w-5 mt-1 text-muted-foreground/40" />
-        </div>
+        </button>
         {children}
       </PopoverContent>
     </Popover>

package/src/core/attachments/state.ts

@@ -2,12 +2,12 @@
 
 import { atom } from "jotai";
 import type { Attachment } from "@marimo-team/client/protocol";
-import { notebookClientAtom } from "@/core/websocket/state";
+import { notebookReplicaAtom } from "@/core/websocket/state";
 
 export type { Attachment } from "@marimo-team/client/protocol";
 
 export const agentAttachmentsAtom = atom((get) => {
-  const notebook = get(notebookClientAtom);
+  const notebook = get(notebookReplicaAtom);
   return (notebook ? get(notebook.attachmentsAtom) : []).filter(
     (attachment) => attachment.kind === "agent",
   );

package/src/core/document/__tests__/codemirror.test.ts

@@ -24,6 +24,7 @@ import { SCHEMA_CACHE } from "../../codemirror/language/languages/sql/completion
 import type { Document } from "../../host/resources/document";
 import { cellCodeBinding } from "../codemirror";
 import { createDocumentModel } from "../model";
+import { fakeBackend } from "../../host/__tests__/fixtures";
 import { cell, document, withCells } from "./fixtures";
 
 const views: EditorView[] = [];
@@ -87,7 +88,8 @@ function setup({
   lifetimes.push(controller);
   const model = createDocumentModel({
     store,
-    notebook: { documentAtom: source, edit: () => new Promise(() => {}) },
+    documentAtom: source,
+    backendAtom: atom(fakeBackend({ edit: () => new Promise(() => {}) })),
     signal: controller.signal,
   });
   const editor = (text = doc) => {

package/src/core/document/__tests__/hooks.test.tsx

@@ -15,11 +15,10 @@ import {
   useCellOrder,
   useDocumentEdit,
   useDocumentSync,
-  type DocumentHost,
 } from "..";
 import { cell, CONFIG, document, withCells } from "./fixtures";
+import { type FakeEdit, fakeBackend } from "../../host/__tests__/fixtures";
 
-type HostEdit = DocumentHost["edit"];
 const lifetimes: AbortController[] = [];
 
 afterEach(() => {
@@ -28,7 +27,7 @@ afterEach(() => {
 
 function setup(
   initial: Document | null = document(),
-  edit: HostEdit = () => new Promise(() => {}),
+  edit: FakeEdit = () => new Promise(() => {}),
 ) {
   const store = createStore();
   const source = atom(initial);
@@ -38,7 +37,8 @@ function setup(
   lifetimes.push(controller);
   const model = createDocumentModel({
     store,
-    notebook: { documentAtom: source, edit },
+    documentAtom: source,
+    backendAtom: atom(fakeBackend({ edit })),
     signal: controller.signal,
   });
   const wrapper = ({ children }: PropsWithChildren) => (
@@ -53,14 +53,14 @@ describe("document hooks", () => {
   it("flushes edits through the model and aborts work when the session ends", async () => {
     const accepted = withCells(document(), { c1: cell("c1", "local") });
     const edit = vi
-      .fn<HostEdit>()
+      .fn<FakeEdit>()
       .mockResolvedValueOnce({
         ok: true,
         status: 200,
         data: { document: accepted.revision, cells: {} },
       })
       .mockImplementation(
-        () => new Deferred<Awaited<ReturnType<HostEdit>>>().promise,
+        () => new Deferred<Awaited<ReturnType<FakeEdit>>>().promise,
       );
     const { store, source, wrapper, controller } = setup(document(), edit);
     const { result, unmount } = renderHook(
@@ -98,7 +98,7 @@ describe("document hooks", () => {
       await rejected;
     });
     expect(edit).toHaveBeenCalledTimes(2);
-    expect(edit.mock.calls[1][2].signal.aborted).toBe(true);
+    expect(edit.mock.calls[1][2]?.signal?.aborted).toBe(true);
   });
 
   it("publishes compound edits atomically without changing host records", () => {

package/src/core/document/__tests__/model.test.ts

@@ -5,8 +5,10 @@ import type * as Protocol from "@marimo-team/client/protocol";
 import { atom, createStore } from "jotai/vanilla";
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 import { Deferred } from "@/utils/Deferred";
+import type { NotebookBackend } from "../../host/backend";
 import type { Document } from "../../host/resources/document";
-import { createDocumentModel, type DocumentHost } from "../model";
+import { createDocumentModel } from "../model";
+import { type FakeEdit, fakeBackend } from "../../host/__tests__/fixtures";
 import { cell, document, withCells } from "./fixtures";
 
 function response(confirmed: Document): Result<Protocol.EditResponseOutput> {
@@ -30,11 +32,11 @@ async function expectAbort(promise: Promise<unknown>, reason: Error) {
 interface SentEdit {
   notebookId: Protocol.NotebookId;
   changes: readonly Protocol.Change[];
-  idempotencyKey: string;
-  signal: AbortSignal;
+  idempotencyKey?: string;
+  signal?: AbortSignal;
 }
 
-function setup() {
+function setup({ connected = true } = {}) {
   const store = createStore();
   const source = atom<Document | null>(document());
   const unmounted = vi.fn();
@@ -45,21 +47,24 @@ function setup() {
     replies.push(reply);
     return reply.promise;
   });
-  const host: DocumentHost = {
-    documentAtom: source,
+  const backend = fakeBackend({
     edit: (notebookId, changes, options) =>
       sent({ notebookId, changes, ...options }),
-  };
+  });
+  const backendAtom = atom<NotebookBackend | null>(connected ? backend : null);
   const controller = new AbortController();
   lifetimes.push(controller);
   const model = createDocumentModel({
     store,
-    notebook: host,
+    documentAtom: source,
+    backendAtom,
     signal: controller.signal,
   });
   return {
     store,
     source,
+    backend,
+    backendAtom,
     model,
     sent,
     replies,
@@ -88,11 +93,12 @@ describe("document synchronization", () => {
     const controller = new AbortController();
     const reason = new Error("Session ended");
     controller.abort(reason);
-    const edit = vi.fn<DocumentHost["edit"]>();
+    const edit = vi.fn<FakeEdit>();
     expect(() =>
       createDocumentModel({
         store,
-        notebook: { documentAtom: source, edit },
+        documentAtom: source,
+        backendAtom: atom(fakeBackend({ edit })),
         signal: controller.signal,
       }),
     ).toThrow(reason);
@@ -152,7 +158,34 @@ describe("document synchronization", () => {
     await vi.advanceTimersByTimeAsync(0);
     store.set(source, document("n2"));
     await rejected;
-    expect(sent.mock.calls[0][0].signal.aborted).toBe(true);
+    expect(sent.mock.calls[0][0].signal?.aborted).toBe(true);
     expect(status()).toMatchObject({ status: "error", retryable: false });
   });
+
+  it("aborts an in-flight edit when the backend changes and never sends it to the new one", async () => {
+    const { store, backendAtom, sent, edit, model } = setup();
+    edit("local");
+    const rejected = expectFailure(model.flush(), "notebook changed");
+    await vi.advanceTimersByTimeAsync(0);
+    const other = vi.fn<FakeEdit>();
+    store.set(backendAtom, fakeBackend({ edit: other }));
+    await rejected;
+    expect(sent).toHaveBeenCalledOnce();
+    expect(sent.mock.calls[0][0].signal?.aborted).toBe(true);
+    expect(other).not.toHaveBeenCalled();
+  });
+
+  it("keeps edits as drafts until a backend is connected", async () => {
+    const { store, backend, backendAtom, sent, edit, status, value } = setup({
+      connected: false,
+    });
+    edit("local");
+    await vi.advanceTimersByTimeAsync(1000);
+    expect(status()).toEqual({ status: "unavailable" });
+    expect(value()?.code).toBe("local");
+
+    store.set(backendAtom, backend);
+    await vi.advanceTimersByTimeAsync(400);
+    expect(sent).toHaveBeenCalledOnce();
+  });
 });

package/src/core/document/__tests__/scenario-fixture.ts

@@ -13,7 +13,8 @@ function createSession(initial: Document, controller: AbortController) {
   const { notebook, control } = createTestNotebook(store, initial);
   const model = createDocumentModel({
     store,
-    notebook,
+    documentAtom: notebook.documentAtom,
+    backendAtom: notebook.backendAtom,
     signal: controller.signal,
   });
   return {

package/src/core/document/__tests__/test-notebook.ts

@@ -6,8 +6,8 @@ import { atom, type createStore } from "jotai/vanilla";
 import { vi } from "vitest";
 import { Deferred } from "@/utils/Deferred";
 import { invariant } from "@/utils/invariant";
+import { type FakeEdit, fakeBackend } from "../../host/__tests__/fixtures";
 import type { Document } from "../../host/resources/document";
-import type { DocumentHost } from "../model";
 import { createTestHost, type EditRequest } from "./test-host";
 
 export interface PendingRequest extends EditRequest {
@@ -33,19 +33,25 @@ export function createTestNotebook(
       documents.push(update);
     }
   }
-  const edit = vi.fn<DocumentHost["edit"]>((notebookId, changes, options) => {
-    options.signal.throwIfAborted();
+  const edit = vi.fn<FakeEdit>((notebookId, changes, options) => {
+    const signal = options?.signal;
+    const idempotencyKey = options?.idempotencyKey;
+    invariant(
+      signal !== undefined && idempotencyKey !== undefined,
+      "The test host answers edits that carry a key and a signal",
+    );
+    signal.throwIfAborted();
     const reply = new Deferred<Result<Protocol.EditResponseOutput>>();
     let outcome: Result<Protocol.EditResponseOutput> | undefined;
-    const abort = () => reply.reject(options.signal.reason);
-    const finish = () => options.signal.removeEventListener("abort", abort);
-    options.signal.addEventListener("abort", abort, { once: true });
+    const abort = () => reply.reject(signal.reason);
+    const finish = () => signal.removeEventListener("abort", abort);
+    signal.addEventListener("abort", abort, { once: true });
     void reply.promise.then(finish, finish);
     const request: PendingRequest = {
       notebookId,
       changes: structuredClone(changes),
-      idempotencyKey: options.idempotencyKey,
-      signal: options.signal,
+      idempotencyKey,
+      signal,
       process() {
         invariant(outcome === undefined, "Request was already processed");
         // Cancellation of the client's wait does not roll back a server commit.
@@ -74,7 +80,11 @@ export function createTestNotebook(
     }
     return reply.promise;
   });
-  const notebook = { documentAtom: source, edit } satisfies DocumentHost;
+  const notebook = {
+    documentAtom: source,
+    backendAtom: atom(fakeBackend({ edit })),
+    edit,
+  };
   return {
     notebook,
     control: {

package/src/core/document/draft.ts

@@ -1,6 +1,6 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import type { Document } from "../host/resources/document";
+import type { Document } from "../host";
 import * as Protocol from "@marimo-team/client/protocol";
 
 /** The fields a client edits. Null code means the host withheld source. */

package/src/core/document/index.ts

@@ -1,5 +1,7 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
+/** The editable document, its drafts, and its synchronization with the host. */
+
 export {
   DocumentProvider,
   useCellCode,
@@ -10,13 +12,4 @@ export {
   useDocumentEdit,
   useDocumentSync,
 } from "./hooks";
-export { createDocumentModel } from "./model";
-export type {
-  DocumentEdit,
-  DocumentEditor,
-  DocumentHost,
-  DocumentModel,
-  DocumentSyncStatus,
-  EditableCell,
-  EditableDocument,
-} from "./model";
+export { createDocumentModel, type DocumentModel } from "./model";

package/src/core/document/model.ts

@@ -2,16 +2,13 @@
 
 import type * as Protocol from "@marimo-team/client/protocol";
 import type { Atom, createStore } from "jotai/vanilla";
+import type { Document, NotebookBackend } from "../host";
 import { createDocumentState, type DocumentState } from "./state";
-import {
-  createDocumentSync,
-  type DocumentHost,
-  type DocumentSyncStatus,
-} from "./sync";
+import { createDocumentSync, type DocumentSyncStatus } from "./sync";
 
 export type { DocumentEdit, DocumentEditor } from "./state";
 export type { EditableCell, EditableDocument } from "./draft";
-export type { DocumentHost, DocumentSyncStatus } from "./sync";
+export type { DocumentSyncStatus } from "./sync";
 
 /** The editable document and its synchronization with the host. */
 export interface DocumentModel extends Pick<
@@ -28,20 +25,26 @@ export interface DocumentModel extends Pick<
 /** Create once per editing session; abort to stop editing and synchronization together. */
 export function createDocumentModel({
   store,
-  notebook,
+  documentAtom,
+  backendAtom,
   signal,
 }: {
   store: ReturnType<typeof createStore>;
-  notebook: DocumentHost;
+  /** The document the host last confirmed. */
+  documentAtom: Atom<Document | null>;
+  /** The backend that confirmed it, whose commands reach its host. */
+  backendAtom: Atom<NotebookBackend | null>;
   signal: AbortSignal;
 }): DocumentModel {
   signal.throwIfAborted();
-  const state = createDocumentState({
+  const state = createDocumentState({ store, source: documentAtom, signal });
+  const sync = createDocumentSync({
     store,
-    source: notebook.documentAtom,
+    documentAtom,
+    backendAtom,
+    state,
     signal,
   });
-  const sync = createDocumentSync({ store, notebook, state, signal });
   return {
     documentAtom: state.documentAtom,
     editAtom: state.editAtom,

package/src/core/document/state.ts

@@ -6,7 +6,7 @@ import {
   type createStore,
   type WritableAtom,
 } from "jotai/vanilla";
-import type { Document } from "../host/resources/document";
+import type { Document } from "../host";
 import * as Protocol from "@marimo-team/client/protocol";
 import {
   acknowledge,