Snapshot document and outputs in MCP execute_code#9654
Merged
Conversation
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
mscolnick
approved these changes
May 21, 2026
![cubic-dev-ai[bot]](https://avatars.githubusercontent.com/in/1082092?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
1 issue found across 1 file
Architecture diagram
sequenceDiagram
participant MCP as MCP Tool
participant Server as Code MCP Server
participant Session as Session
participant Doc as Document
participant SV as SessionView
participant CO as CellOutputs
participant Runtime as Runtime
participant CM as CodeMode (cm.get_context)
Note over MCP,CM: execute_code MCP tool call flow
MCP->>Server: execute_code(session_id, code)
Server->>Server: Create ScratchCellListener, acquire scratchpad_lock
Server->>Doc: snapshot_document()
Doc-->>Server: cell_ids, cells
Server->>SV: get_cell_outputs(cell_ids)
SV-->>Server: output dictionary
Server->>SV: get_cell_console_outputs(cell_ids)
SV-->>Server: console_outputs dictionary
Server->>CO: CellOutputs(outputs, console_outputs)
CO-->>Server: cell_outputs object
Server->>Runtime: put_control_request(ExecuteScratchpadCommand with notebook_cells, cell_outputs, code, run_id)
Runtime->>Runtime: Execute code
Runtime->>CM: cm.get_context()
CM->>CM: Read notebook_cells and cell_outputs from ExecuteScratchpadCommand
CM-->>Runtime: Context available
Runtime-->>Server: Execution complete
Server-->>MCP: CodeExecutionResult
Reply with feedback, questions, or to request a fix.
Re-trigger cubic
mscolnick
previously approved these changes
May 21, 2026
Mirrors what `/api/execute` already does. `cm.get_context()` previously raised "NotebookDocument not available" when invoked from code run via the MCP tool, because `execute_code` didn't populate `notebook_cells` or `cell_outputs` on `ExecuteScratchpadCommand`.
Pre-commit's no-double-backticks-in-docstrings hook (and AGENTS.md) require Markdown inline code, not reStructuredText.
Contributor
There was a problem hiding this comment.
Pull request overview
This PR aligns MCP execute_code behavior with the existing HTTP /execute flow by ensuring scratchpad executions carry a snapshot of the notebook document and per-cell outputs, enabling cm.get_context() to expose cells and their last outputs.
Changes:
- Added
snapshot_for_scratchpad(session)helper to capture(notebook_cells, cell_outputs)from a session. - Refactored the HTTP
/executeSSE path to use the shared snapshot helper. - Updated MCP
execute_codeto populateExecuteScratchpadCommand.notebook_cellsandcell_outputs, plus added a unit test for the helper.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
| tests/_server/test_scratchpad.py | Adds coverage for the new scratchpad snapshot helper. |
| marimo/_server/scratchpad.py | Introduces snapshot_for_scratchpad() to package document + outputs for scratchpad runs. |
| marimo/_server/api/endpoints/execution.py | Reuses the shared snapshot helper when building ExecuteScratchpadCommand for /execute. |
| marimo/_mcp/code_server/main.py | Updates MCP execute_code to include notebook cells + outputs in the scratchpad command. |
Comment on lines
+72
to
+88
| output = CellOutput( | ||
| channel=CellChannel.OUTPUT, mimetype="text/plain", data="1" | ||
| ) | ||
| console = [CellOutput.stdout("hi\n")] | ||
|
|
||
| session = MagicMock() | ||
| session.document.cells = (cell,) | ||
| session.document.cell_ids = (cell.id,) | ||
| session.session_view.get_cell_outputs.return_value = {cell.id: output} | ||
| session.session_view.get_cell_console_outputs.return_value = { | ||
| cell.id: console | ||
| } | ||
|
|
||
| notebook_cells, cell_outputs = snapshot_for_scratchpad(session) | ||
| assert notebook_cells == (cell,) | ||
| assert cell_outputs.output == snapshot({"a": output}) | ||
| assert cell_outputs.console_outputs == snapshot({"a": console}) |
|
🚀 Development release published. You may be able to view the changes at https://marimo.app?v=0.23.7-dev87 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Mirrors what
/api/executealready does.cm.get_context()previously raised "NotebookDocument not available" when invoked from code run via the MCP tool, becauseexecute_codedidn't populatenotebook_cellsorcell_outputsonExecuteScratchpadCommand.