feat(code-mode): read-before-write protection for cell edits#9585
Conversation
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)
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
There was a problem hiding this comment.
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
Reply with feedback, questions, or to request a fix.
Re-trigger cubic
There was a problem hiding this comment.
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.versionand version reconciliation for document changes/reloads. - Introduces
AgentReadTrackerstored on eachKernel. - 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. |
test_ctx_run_cell_cascade_error edited cell_a without first materializing it, which now trips the read-before-write staleness check. Add a ctx.cells["cell_a"] read before the edit.
edit_cell(target, name="setup") silently fills `code` from self.graph.cells when the agent passes code=None, so the check was being skipped on a path that still writes to the doc. Move the check below the migration block so it sees the post-migration code.
| f"Read it first (e.g. `ctx.cells[{cell_id!r}].code`) before " | ||
| f"editing.{other_hint}\n" | ||
| f"To override and overwrite without re-reading, pass " | ||
| f"skip_staleness_check=True to cm.get_context()." |
There was a problem hiding this comment.
nice, this was going to be my main comment.
| current = self._document.get_cell_version(written_id) | ||
| if current is not None: | ||
| self._note_read(written_id, current) |
There was a problem hiding this comment.
nit, but we probably don't need a special method for this:
| current = self._document.get_cell_version(written_id) | |
| if current is not None: | |
| self._note_read(written_id, current) | |
| if cell := self._document.get(written_id): | |
| self._note_read(written_id, cell.version) |
|
|
||
|
|
||
| @dataclass | ||
| class Agent: |
There was a problem hiding this comment.
Nice. I really like that we have some carved out space for this kind of state now.
| def get_context( | ||
| *, | ||
| skip_validation: bool = False, | ||
| skip_staleness_check: bool = False, | ||
| ) -> AsyncCodeModeContext: |
There was a problem hiding this comment.
Maybe for a followup, but could/should we emit (over stderror) the cells that changed externally as a kind of indicator to the model (hey these have changed)? Maybe context bloat.
|
🚀 Development release published. You may be able to view the changes at https://marimo.app?v=0.23.7-dev74 |
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.
post-apply for the agent's own writes (so create→edit across contexts
works without re-reading)
preserves the prior version, changed code bumps it