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

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,

package/src/core/document/sync.ts

@@ -2,13 +2,10 @@
 
 import { idempotencyKey, type Result } from "@marimo-team/client";
 import type * as Protocol from "@marimo-team/client/protocol";
-import { atom, type createStore } from "jotai/vanilla";
-import type { NotebookClient } from "../host/notebook-client";
+import { type Atom, atom, type createStore } from "jotai/vanilla";
+import type { Document, NotebookBackend } from "../host";
 import type { DocumentState, DocumentSubmission } from "./state";
 
-/** The notebook state and operation needed by document editing. */
-export type DocumentHost = Pick<NotebookClient, "documentAtom" | "edit">;
-
 export type DocumentSyncStatus =
   | {
       readonly status:
@@ -27,6 +24,8 @@ export type DocumentSyncStatus =
 
 interface Submission {
   readonly submission: DocumentSubmission;
+  /** The backend that confirmed the submission's base. */
+  readonly backend: NotebookBackend;
   readonly idempotencyKey: string;
   readonly startedAt: number;
   readonly controller: AbortController;
@@ -74,20 +73,22 @@ function delay(ms: number, signal: AbortSignal): Promise<void> {
 /** Own submission scheduling; follow the existing source without opening a stream. */
 export function createDocumentSync({
   store,
-  notebook,
+  documentAtom: source,
+  backendAtom,
   state,
   signal,
 }: {
   store: ReturnType<typeof createStore>;
-  notebook: DocumentHost;
+  documentAtom: Atom<Document | null>;
+  backendAtom: Atom<NotebookBackend | null>;
   state: DocumentState;
   signal: AbortSignal;
 }) {
   signal.throwIfAborted();
-  const source = notebook.documentAtom;
   const statusAtom = atom<DocumentSyncStatus>({ status: "idle" });
   const observed = atom((get) => ({
     document: get(source),
+    backend: get(backendAtom),
     pending: get(state.pendingAtom),
     changes: get(state.changesAtom),
   }));
@@ -117,8 +118,9 @@ export function createDocumentSync({
       const current = store.get(source);
       if (
         active &&
-        current &&
-        current.revision.notebook_id !== active.submission.notebookId
+        (store.get(backendAtom) !== active.backend ||
+          (current &&
+            current.revision.notebook_id !== active.submission.notebookId))
       ) {
         active.controller.abort(
           new Error("The notebook changed during synchronization"),
@@ -129,7 +131,7 @@ export function createDocumentSync({
       }
       clearTimer();
       const changes = store.get(state.changesAtom);
-      if (changes === null) {
+      if (changes === null || store.get(backendAtom) === null) {
         store.set(statusAtom, { status: "unavailable" });
       } else if (changes.length === 0) {
         // Release unsent reversals instead of keeping an invisible override.
@@ -170,14 +172,12 @@ export function createDocumentSync({
         signal.throwIfAborted();
         request.controller.signal.throwIfAborted();
         result = await abortable(
-          notebook.edit(
-            request.submission.notebookId,
-            request.submission.changes,
-            {
+          request.backend
+            .notebook(request.submission.notebookId)
+            .edit(request.submission.changes, {
               idempotencyKey: request.idempotencyKey,
               signal: controller.signal,
-            },
-          ),
+            }),
           controller.signal,
         );
       } catch (error) {
@@ -248,7 +248,12 @@ export function createDocumentSync({
           throw new Error("The notebook changed during synchronization");
         }
         if (active === null) {
-          if (current === null || store.get(state.changesAtom) === null) {
+          const backend = store.get(backendAtom);
+          if (
+            current === null ||
+            backend === null ||
+            store.get(state.changesAtom) === null
+          ) {
             throw unavailable();
           }
           const submission = store.set(state.submitAtom);
@@ -257,6 +262,7 @@ export function createDocumentSync({
           }
           active = {
             submission,
+            backend,
             idempotencyKey: idempotencyKey(),
             startedAt: Date.now(),
             controller: new AbortController(),