From 911caa5a14e5381432a5d02a78758d2d1a37bdca Mon Sep 17 00:00:00 2001 From: Kiran Gadhave Date: Wed, 22 Jul 2026 10:11:52 -0700 Subject: [PATCH 1/2] fix(document): apply transactions atomically Invalid cell references could surface as raw `KeyError`s, while a failure after earlier changes had been applied could leave the document partially mutated. Validate targets and anchors in batch order, then stage changes on an isolated copy so failed transactions preserve cells and versions. --- marimo/_messaging/notebook/document.py | 62 +++++-- tests/_messaging/notebook/test_document.py | 191 ++++++++++++++++++++- 2 files changed, 240 insertions(+), 13 deletions(-) diff --git a/marimo/_messaging/notebook/document.py b/marimo/_messaging/notebook/document.py index 2f0f792cd67..2dd9b280105 100644 --- a/marimo/_messaging/notebook/document.py +++ b/marimo/_messaging/notebook/document.py @@ -149,19 +149,33 @@ def __iter__(self) -> Iterator[CellId_t]: def apply(self, tx: Transaction) -> Transaction: """Validate and apply *tx*, return it with `version` assigned. - Raises `ValueError` for validation failures and `KeyError` - when a change references a non-existent cell. + Raises `ValueError` when a change has an invalid target, anchor, + or intra-transaction conflict. """ if not tx.changes: return structs_replace(tx, version=self._version) _validate(tx.changes, self._cells) + staged = NotebookDocument( + [ + structs_replace( + cell, + config=structs_replace(cell.config), + ) + for cell in self._cells + ] + ) for change in tx.changes: - self._apply_change(change) + staged._apply_change(change) - self._version += 1 - return structs_replace(tx, version=self._version) + next_version = self._version + 1 + applied = structs_replace(tx, version=next_version) + + # Commit only after staging and result construction both succeed. + self._cells = staged._cells + self._version = next_version + return applied def _apply_change(self, change: DocumentChange) -> None: # TODO: refactor to use match/case (min Python is 3.10) once @@ -322,22 +336,32 @@ def notebook_document_context( def _validate( changes: tuple[DocumentChange, ...], cells: list[NotebookCell] ) -> None: - """Check for conflicting changes. Raises `ValueError`.""" - existing_ids = {c.id for c in cells} - created: set[CellId_t] = set() + """Check targets, anchors, and conflicts. Raises `ValueError`.""" + live_ids = {cell.id for cell in cells} + # Claimed IDs are never released within a transaction, preserving the + # rule that delete + recreate cannot reuse an ID in the same batch. + claimed_ids = set(live_ids) deleted: set[CellId_t] = set() updated: set[CellId_t] = set() moved: set[CellId_t] = set() for change in changes: if isinstance(change, CreateCell): - if change.cell_id in existing_ids or change.cell_id in created: + if change.cell_id in claimed_ids: raise ValueError(f"Cell {change.cell_id!r} already exists") if change.before is not None and change.after is not None: raise ValueError( "CreateCell cannot specify both 'before' and 'after'" ) - created.add(change.cell_id) + anchor = ( + change.before if change.before is not None else change.after + ) + if anchor is not None and anchor not in live_ids: + raise ValueError( + f"CreateCell anchor {anchor!r} does not exist" + ) + live_ids.add(change.cell_id) + claimed_ids.add(change.cell_id) elif isinstance(change, DeleteCell): if change.cell_id in deleted: @@ -354,7 +378,10 @@ def _validate( f"Cannot delete cell {change.cell_id!r} that is also " f"moved in the same transaction" ) + if change.cell_id not in live_ids: + raise ValueError(f"Cell {change.cell_id!r} does not exist") deleted.add(change.cell_id) + live_ids.remove(change.cell_id) elif isinstance(change, MoveCell): if change.cell_id in deleted: @@ -362,12 +389,23 @@ def _validate( f"Cannot move cell {change.cell_id!r} that is also " f"deleted in the same transaction" ) + if change.cell_id not in live_ids: + raise ValueError(f"Cell {change.cell_id!r} does not exist") if change.before is not None and change.after is not None: raise ValueError( "MoveCell cannot specify both 'before' and 'after'" ) - if change.before is None and change.after is None: + anchor = ( + change.before if change.before is not None else change.after + ) + if anchor is None: raise ValueError("MoveCell requires 'before' or 'after'") + if anchor == change.cell_id: + raise ValueError( + f"MoveCell cannot anchor cell {change.cell_id!r} to itself" + ) + if anchor not in live_ids: + raise ValueError(f"MoveCell anchor {anchor!r} does not exist") moved.add(change.cell_id) elif isinstance(change, ReorderCells): @@ -379,6 +417,8 @@ def _validate( f"Cannot update cell {change.cell_id!r} that is also " f"deleted in the same transaction" ) + if change.cell_id not in live_ids: + raise ValueError(f"Cell {change.cell_id!r} does not exist") updated.add(change.cell_id) else: diff --git a/tests/_messaging/notebook/test_document.py b/tests/_messaging/notebook/test_document.py index f2feebd01e1..e968196bfa6 100644 --- a/tests/_messaging/notebook/test_document.py +++ b/tests/_messaging/notebook/test_document.py @@ -42,6 +42,23 @@ def _ids(doc: NotebookDocument) -> list[str]: return [str(cid) for cid in doc.cell_ids] +def _state( + doc: NotebookDocument, +) -> list[tuple[CellId_t, str, str, int | None, bool, bool, int]]: + return [ + ( + cell.id, + cell.code, + cell.name, + cell.config.column, + cell.config.disabled, + cell.config.hide_code, + cell.version, + ) + for cell in doc.cells + ] + + # ------------------------------------------------------------------ # CreateCell # ------------------------------------------------------------------ @@ -148,6 +165,42 @@ def test_after_pending(self) -> None: ) assert _ids(doc) == snapshot(["a", "x", "y"]) + def test_missing_after_anchor_raises(self) -> None: + doc = _doc("a") + with pytest.raises( + ValueError, + match="CreateCell anchor 'missing' does not exist", + ): + doc.apply( + _tx( + CreateCell( + cell_id=CellId_t("new"), + code="x", + name="__", + config=CellConfig(), + after=CellId_t("missing"), + ) + ) + ) + + def test_missing_before_anchor_raises(self) -> None: + doc = _doc("a") + with pytest.raises( + ValueError, + match="CreateCell anchor 'missing' does not exist", + ): + doc.apply( + _tx( + CreateCell( + cell_id=CellId_t("new"), + code="x", + name="__", + config=CellConfig(), + before=CellId_t("missing"), + ) + ) + ) + def test_duplicate_id_raises(self) -> None: doc = _doc("a") with pytest.raises(ValueError, match="already exists"): @@ -251,7 +304,7 @@ def test_delete_all(self) -> None: def test_delete_not_found(self) -> None: doc = _doc("a") - with pytest.raises(KeyError): + with pytest.raises(ValueError, match="Cell 'missing' does not exist"): doc.apply(_tx(DeleteCell(cell_id=CellId_t("missing")))) @@ -276,6 +329,66 @@ def test_no_anchor_raises(self) -> None: with pytest.raises(ValueError, match="before.*after"): doc.apply(_tx(MoveCell(cell_id=CellId_t("a")))) + def test_missing_cell_raises(self) -> None: + doc = _doc("a") + with pytest.raises(ValueError, match="Cell 'missing' does not exist"): + doc.apply( + _tx( + MoveCell( + cell_id=CellId_t("missing"), + after=CellId_t("a"), + ) + ) + ) + + def test_missing_anchor_raises(self) -> None: + doc = _doc("a", "b") + with pytest.raises( + ValueError, + match="MoveCell anchor 'missing' does not exist", + ): + doc.apply( + _tx( + MoveCell( + cell_id=CellId_t("a"), + after=CellId_t("missing"), + ) + ) + ) + + def test_self_anchor_raises(self) -> None: + doc = _doc("a", "b") + with pytest.raises( + ValueError, + match="MoveCell cannot anchor cell 'a' to itself", + ): + doc.apply( + _tx( + MoveCell( + cell_id=CellId_t("a"), + after=CellId_t("a"), + ) + ) + ) + + def test_anchor_created_earlier_in_batch(self) -> None: + doc = _doc("a", "b") + doc.apply( + _tx( + CreateCell( + cell_id=CellId_t("new"), + code="x", + name="__", + config=CellConfig(), + ), + MoveCell( + cell_id=CellId_t("a"), + after=CellId_t("new"), + ), + ) + ) + assert _ids(doc) == snapshot(["b", "new", "a"]) + # ------------------------------------------------------------------ # SetCode @@ -291,7 +404,7 @@ def test_update_code(self) -> None: def test_not_found(self) -> None: doc = _doc("a") - with pytest.raises(KeyError): + with pytest.raises(ValueError, match="Cell 'missing' does not exist"): doc.apply(_tx(SetCode(cell_id=CellId_t("missing"), code="x"))) @@ -522,6 +635,80 @@ def test_valid_mixed_ops(self) -> None: assert _ids(doc) == snapshot(["a", "new"]) +# ------------------------------------------------------------------ +# Atomicity +# ------------------------------------------------------------------ + + +class TestAtomicity: + def test_invalid_late_anchor_does_not_mutate_document(self) -> None: + doc = _doc("a", "b") + doc.apply(_tx(SetCode(cell_id=CellId_t("a"), code="original"))) + before_state = _state(doc) + before_version = doc.version + + with pytest.raises( + ValueError, + match="MoveCell anchor 'missing' does not exist", + ): + doc.apply( + _tx( + SetName(cell_id=CellId_t("a"), name="renamed"), + CreateCell( + cell_id=CellId_t("new"), + code="new code", + name="__", + config=CellConfig(disabled=True), + ), + MoveCell( + cell_id=CellId_t("b"), + after=CellId_t("missing"), + ), + ) + ) + + assert _state(doc) == before_state + assert doc.version == before_version + + def test_unexpected_apply_failure_does_not_mutate_document( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + doc = _doc("a", "b") + before_state = _state(doc) + before_version = doc.version + original_apply_change = NotebookDocument._apply_change + applied_changes = 0 + + def fail_after_second_change( + staged: NotebookDocument, change: DocumentChange + ) -> None: + nonlocal applied_changes + original_apply_change(staged, change) + applied_changes += 1 + if applied_changes == 2: + raise RuntimeError("injected apply failure") + + monkeypatch.setattr( + NotebookDocument, "_apply_change", fail_after_second_change + ) + + with pytest.raises(RuntimeError, match="injected apply failure"): + doc.apply( + _tx( + SetName(cell_id=CellId_t("a"), name="renamed"), + CreateCell( + cell_id=CellId_t("new"), + code="new code", + name="__", + config=CellConfig(disabled=True), + ), + ) + ) + + assert _state(doc) == before_state + assert doc.version == before_version + + # ------------------------------------------------------------------ # Versioning # ------------------------------------------------------------------ From 67544acc9a105cea2817576c21cf1e556475302c Mon Sep 17 00:00:00 2001 From: Kiran Gadhave Date: Wed, 22 Jul 2026 12:43:52 -0700 Subject: [PATCH 2/2] fix(document): stage transactions with copy-on-write Atomic staging deep-cloned every cell and config, making common single-cell edits allocate in proportion to notebook size. Shallow-copy the ordered list and clone only cells that are updated so failed transactions remain isolated without full-document allocation overhead. --- marimo/_messaging/notebook/document.py | 27 +++++++++----- tests/_messaging/notebook/test_document.py | 43 ++++++++++++++++++++++ 2 files changed, 61 insertions(+), 9 deletions(-) diff --git a/marimo/_messaging/notebook/document.py b/marimo/_messaging/notebook/document.py index 2dd9b280105..1e5558f0367 100644 --- a/marimo/_messaging/notebook/document.py +++ b/marimo/_messaging/notebook/document.py @@ -157,18 +157,19 @@ def apply(self, tx: Transaction) -> Transaction: _validate(tx.changes, self._cells) - staged = NotebookDocument( - [ - structs_replace( - cell, - config=structs_replace(cell.config), - ) - for cell in self._cells - ] - ) + staged = NotebookDocument(self._cells) + owned_cell_ids: set[CellId_t] = set() for change in tx.changes: + if isinstance(change, (SetCode, SetName, SetConfig)): + if change.cell_id not in owned_cell_ids: + staged._clone_cell(change.cell_id) + owned_cell_ids.add(change.cell_id) + staged._apply_change(change) + if isinstance(change, CreateCell): + owned_cell_ids.add(change.cell_id) + next_version = self._version + 1 applied = structs_replace(tx, version=next_version) @@ -284,6 +285,14 @@ def _find_index(self, cell_id: CellId_t) -> int: return i raise KeyError(f"Cell {cell_id!r} not found in document") + def _clone_cell(self, cell_id: CellId_t) -> None: + idx = self._find_index(cell_id) + cell = self._cells[idx] + self._cells[idx] = structs_replace( + cell, + config=structs_replace(cell.config), + ) + def _find_cell(self, cell_id: CellId_t) -> NotebookCell: for cell in self._cells: if cell.id == cell_id: diff --git a/tests/_messaging/notebook/test_document.py b/tests/_messaging/notebook/test_document.py index e968196bfa6..d0bccc04249 100644 --- a/tests/_messaging/notebook/test_document.py +++ b/tests/_messaging/notebook/test_document.py @@ -641,6 +641,49 @@ def test_valid_mixed_ops(self) -> None: class TestAtomicity: + def test_only_updated_cells_are_cloned(self) -> None: + doc = _doc("a", "b") + original_a = doc.get_cell(CellId_t("a")) + original_b = doc.get_cell(CellId_t("b")) + + doc.apply(_tx(SetCode(cell_id=CellId_t("a"), code="updated"))) + + updated_a = doc.get_cell(CellId_t("a")) + untouched_b = doc.get_cell(CellId_t("b")) + assert updated_a is not original_a + assert updated_a.config is not original_a.config + assert untouched_b is original_b + assert untouched_b.config is original_b.config + + def test_each_existing_cell_is_cloned_at_most_once( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + doc = _doc("a") + cloned_ids: list[CellId_t] = [] + original_clone_cell = NotebookDocument._clone_cell + + def record_clone(staged: NotebookDocument, cell_id: CellId_t) -> None: + cloned_ids.append(cell_id) + original_clone_cell(staged, cell_id) + + monkeypatch.setattr(NotebookDocument, "_clone_cell", record_clone) + + doc.apply( + _tx( + SetCode(cell_id=CellId_t("a"), code="updated"), + SetName(cell_id=CellId_t("a"), name="renamed"), + CreateCell( + cell_id=CellId_t("new"), + code="new code", + name="__", + config=CellConfig(), + ), + SetName(cell_id=CellId_t("new"), name="created"), + ) + ) + + assert cloned_ids == [CellId_t("a")] + def test_invalid_late_anchor_does_not_mutate_document(self) -> None: doc = _doc("a", "b") doc.apply(_tx(SetCode(cell_id=CellId_t("a"), code="original")))