feat(code-mode): read-before-write protection for cell edits

[mscolnick]

Closes marimo-team/marimo-pair#26

ctx.edit_cell now raises StaleCellError when targeting a cell the agent
hasn't materialized at its current version, preventing the agent from
clobbering human edits made between scratchpad calls.

• Per-cell version on NotebookCell, bumped on real SetCode changes
• AgentReadTracker on Kernel records reads via ctx.cells[...] and
  post-apply for the agent's own writes (so create→edit across contexts
  works without re-reading)
• _replace_cells reconciles versions on file-watch reload: same code
  preserves the prior version, changed code bumps it
• Empty/whitespace cells are exempt — no prior content to clobber
• Opt out with cm.get_context(skip_staleness_check=True)

[vercel]

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

Project	Deployment	Actions	Updated (UTC)
marimo-docs	Ready	Preview, Comment	May 18, 2026 7:46pm

[cubic-dev-ai]

1 issue found across 11 files

Architecture diagram

sequenceDiagram
    participant Agent as Agent Process
    participant CM as get_context()
    participant Ctx as AsyncCodeModeContext
    participant Tracker as AgentReadTracker
    participant Doc as NotebookDocument
    participant Kernel

    Note over Agent,Kernel: Read-Before-Write Protection Flow

    Agent->>CM: get_context(skip_staleness_check=False)
    CM->>Ctx: NEW: create with AgentReadTracker ref
    Ctx->>Kernel: access kernel.agent.read_tracker

    alt Agent reads a cell via cells[...]
        Agent->>Ctx: nb.cells[cell_id].code
        Ctx->>Doc: get NotebookCell for cell_id
        Doc-->>Ctx: NotebookCell with version
        Ctx->>Tracker: record_read(cell_id, version)
        Ctx-->>Agent: NotebookCell.code
    end

    Note over Agent,Doc: edit_cell() with staleness check

    Agent->>Ctx: nb.edit_cell(target, code="new_code")
    Ctx->>Ctx: resolve target cell_id
    alt code is not None AND not skip_staleness_check AND cell not pending_add
        Ctx->>Doc: get(cell_id)
        Doc-->>Ctx: NotebookCell or None
        alt cell exists AND code.strip() non-empty
            Ctx->>Tracker: has_read(cell_id, current_version)
            alt has_read returns False
                Tracker->>Ctx: False
                Ctx-->>Agent: raise StaleCellError
                Note over Ctx: Includes stale_cells frozenset
            else has_read returns True
                Tracker->>Ctx: True
                Ctx->>Ctx: proceed with edit
            end
        else empty cell or cell not found
            Ctx->>Ctx: skip staleness check
        end
    end

    Note over Agent,Doc: Post-apply read recording for agent writes

    Ctx->>Doc: apply transaction (SetCode changes)
    Doc->>Doc: bump cell.version if code changed
    loop for each applied op
        alt op is AddOp or UpdateOp with code
            Ctx->>Doc: get_cell_version(written_id)
            Doc-->>Ctx: current version
            Ctx->>Tracker: record_read(written_id, current_version)
        end
    end

    Note over Agent,Doc: File-watch reload path

    alt File-watch triggers _replace_cells
        Doc->>Doc: _replace_cells(new_cells)
        Doc->>Doc: reconcile versions with prior state
        alt cell code unchanged
            Doc->>Doc: inherit prior version
        else cell code changed
            Doc->>Doc: bump version to prior + 1
        end
    end

    Note over Agent,Kernel: Error recovery

    alt StaleCellError raised
        Agent->>Agent: catch StaleCellError
        Agent->>Agent: re-read stale cells via cells[...]
        Agent->>Ctx: edit_cell() again (now has read)
        Note over Agent,Ctx: Fresh materialization records read
    end

Loading

<sub>Reply with feedback, questions, or to request a fix.

Re-trigger cubic</sub>

[Copilot]

Pull request overview

Adds read-before-write protection for code-mode cell edits by tracking per-cell document versions and agent reads across scratchpad executions.

Changes:

• Adds NotebookCell.version and version reconciliation for document changes/reloads.
• Introduces AgentReadTracker stored on each Kernel.
• Updates code-mode context/cell views to record reads, raise StaleCellError, and expose an opt-out.

Reviewed changes

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

Show a summary per file

File	Description
marimo/_code_mode/_context.py	Adds staleness checks, read recording, error type, and post-write tracker updates.
marimo/_code_mode/__init__.py	Exports StaleCellError.
marimo/_messaging/notebook/document.py	Adds per-cell versions and version bump/reconciliation logic.
marimo/_runtime/agent.py	Adds per-kernel agent read tracking state.
marimo/_runtime/runtime.py	Initializes agent state on Kernel.
tests/_code_mode/test_cells_view.py	Updates cell-view test helper for read-recording calls.
tests/_code_mode/test_context.py	Disables staleness checks in existing context tests.
tests/_code_mode/test_context_autosave.py	Disables staleness checks in autosave context tests.
tests/_code_mode/test_staleness.py	Adds staleness behavior tests.
tests/_messaging/notebook/test_document.py	Adds document cell-version tests.
tests/_runtime/test_agent.py	Adds agent read-tracker tests.