fix(document): apply transactions atomically

[kirangadhave]

📝 Summary

Make NotebookDocument.apply() atomic when validation or change application fails. Previously, invalid cell references could surface as raw KeyErrors, and an unexpected failure after earlier changes had been applied could leave cells or versions partially mutated.

Validate targets and anchors in batch order, then apply changes to an isolated copy and replace live state only after the complete transaction succeeds. This preserves create-then-reference behavior within a batch while ensuring failed transactions leave the document unchanged.

Add focused regression coverage for missing targets and anchors, self-anchored moves, earlier-created anchors, validation rollback, and unexpected execution failures.

Closes MO-6621

[vercel]

The latest updates on your projects. Learn more about Vercel for GitHub.

Project	Deployment	Actions	Updated (UTC)
marimo-docs	Ready	Preview, Comment	Jul 22, 2026 8:04pm

[cubic-dev-ai]

No issues found across 2 files

Architecture diagram

sequenceDiagram
    participant Client as Notebook Client
    participant API as Document.apply(tx)
    participant Validate as _validate()
    participant Staged as Staged NotebookDocument
    participant Live as Live State
    
    Note over Client,Live: Transaction Application Flow
    
    Client->>API: apply(tx)
    
    Note over API,Validate: Validation Phase (NEW: comprehensive checks)
    API->>Validate: _validate(tx.changes, self._cells)
    Validate->>Validate: Build live_ids from current cells
    Validate->>Validate: Track claimed_ids, deleted, moved, updated
    
    loop For each change
        alt CreateCell
            Validate->>Validate: Check cell_id not in claimed_ids
            Validate->>Validate: Check anchor (before/after) exists in live_ids
            Validate->>Validate: Add cell_id to live_ids and claimed_ids
        else DeleteCell
            Validate->>Validate: Check cell_id not already deleted or moved
            Validate->>Validate: Check cell_id exists in live_ids
            Validate->>Validate: Remove cell_id from live_ids
        else MoveCell
            Validate->>Validate: Check cell_id not deleted
            Validate->>Validate: Check cell_id exists in live_ids
            Validate->>Validate: Check anchor exists in live_ids and not self
            Validate->>Validate: Check anchor not referencing itself
        else SetCode/SetName/SetCellConfig
            Validate->>Validate: Check cell_id not deleted
            Validate->>Validate: Check cell_id exists in live_ids
        end
    end
    
    alt Validation fails
        Validate-->>API: Raise ValueError
        API-->>Client: Raise ValueError (document unchanged)
    else All valid
        Validate-->>API: OK
        
        Note over API,Staged: Staging Phase (NEW: isolated copy)
        API->>Staged: Create deep copy of current cells
        Staged-->>API: Staged copy
        
        Note over API,Staged: Apply Phase (NEW: staged only)
        loop For each change
            API->>Staged: _apply_change(change)
        end
        
        alt Application fails (e.g., RuntimeError)
            Staged-->>API: Exception
            API-->>Client: Raise exception (document unchanged)
        else All changes applied
            API->>API: next_version = self._version + 1
            API->>API: Create return Transaction with next_version
            
            Note over API,Live: Commit Phase (NEW: atomic swap)
            API->>Live: Replace self._cells with staged._cells
            API->>Live: Set self._version = next_version
            Live-->>API: Complete
            
            API-->>Client: Return applied Transaction
        end
    end

Loading

_{Re-trigger cubic}

[Copilot]

Pull request overview

This PR makes NotebookDocument.apply() transactional/atomic: it validates targets/anchors up front, stages changes, and only commits the new state if the full transaction succeeds—preventing partial mutation on invalid references or unexpected apply-time failures.

Changes:

• Added batch-order validation for missing targets/anchors and intra-transaction conflicts, raising ValueError (instead of surfacing raw KeyError).
• Implemented staging + commit in NotebookDocument.apply() to ensure atomic document updates.
• Added regression tests covering missing targets/anchors, self-anchored moves, create-then-reference within a batch, and rollback on failures.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
marimo/_messaging/notebook/document.py	Adds stricter validation and stages transactions so document state updates commit atomically.
tests/_messaging/notebook/test_document.py	Expands validation and atomicity test coverage to prevent regressions.