fix(cache): mo.cache raises KeyError: '__scratch__' in scratchpad by everettroeth · Pull Request #9664 · marimo-team/marimo

everettroeth · 2026-05-22T14:00:04Z

Why

@mo.cache, @mo.lru_cache, and with mo.persistent_cache(...) raise KeyError: '__scratch__' when used from the /api/kernel/execute scratchpad endpoint:

File "marimo/_save/save.py", line 232, in _set_context
    self.scoped_refs |= graph.cells[cell_id].defs & set(glbls.keys())
KeyError: '__scratch__'

ctx.graph returns the kernel's main graph (marimo/_runtime/context/kernel_context.py:46), but the scratchpad cell is registered only in the Runner's local graph (marimo/_runtime/runtime.py:1706-1707), so the decorator's cell lookup misses. The same root cause raises CacheException from _cache_context.trace (save.py:711) and KeyError from mo.defs() / mo.refs() (runtime.py:195,222).

What

In Kernel.run_scratchpad, also register SCRATCH_CELL_ID in self.graph for the duration of the run, unregister in a try/finally. The scratch cell's defs and refs are cleared at runtime.py:1702-1703, so this adds no edges, triggers no reactive recomputation, and fires no callbacks. State updates are flushed after cleanup so reactive cells never observe __scratch__.

Three regression tests added under TestExecution in tests/_runtime/test_runtime.py alongside the existing test_run_scratch* tests.

… KeyError: '__scratch__' Using @mo.cache (or @mo.lru_cache, or `with mo.persistent_cache(...)`) from inside the HTTP scratchpad endpoint /api/kernel/execute raised `KeyError: '__scratch__'` at decoration time, from `_cache_call._set_context` at marimo/_save/save.py:232: self.scoped_refs |= graph.cells[cell_id].defs & set(glbls.keys()) `graph` is `ctx.graph`, which always returns the kernel's main DirectedGraph (KernelRuntimeContext.graph at kernel_context.py:46). `cell_id` is SCRATCH_CELL_ID == "__scratch__". The scratchpad cell is registered only in a separate local graph constructed inside Kernel.run_scratchpad (runtime.py:1706-1707) and never in `self.graph`, so the lookup always misses. The same mismatch affected the persistent_cache context manager path at save.py:711 (CacheException), and `mo.defs()` / `mo.refs()` at runtime.py:195,222 (same KeyError). Fix: in Kernel.run_scratchpad, also register the scratch CellImpl in the kernel's main graph for the duration of the run, and unregister in a try/finally. The scratch cell has its defs and refs cleared beforehand (existing behavior at lines 1702-1703), so this registration adds no edges to the topology, triggers no reactive recomputation, and fires no callbacks (register_cell / delete_cell have neither). State updates queued during scratchpad execution are still flushed AFTER the cleanup so reactive cells never observe `__scratch__` in their graph. Verified: 1370 tests across tests/_save/ and tests/_runtime/ pass with no regressions. mypy clean across all 703 source files. ruff check and ruff format clean.

vercel · 2026-05-22T14:00:11Z

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

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

github-actions · 2026-05-22T14:00:19Z

All contributors have signed the CLA ✍️ ✅
_{Posted by the CLA Assistant Lite bot.}

Copilot

Pull request overview

Note

Copilot was unable to run its full agentic suite in this review.

This PR fixes scratchpad execution failures for cache-related features by ensuring the scratch cell is temporarily visible in the kernel’s main graph during run_scratchpad, and adds regression coverage for @mo.cache and mo.persistent_cache usage in scratchpad.

Changes:

Temporarily registers SCRATCH_CELL_ID in the kernel main graph during scratchpad execution and guarantees cleanup via try/finally.
Adds runtime regression tests covering @mo.cache, mo.persistent_cache(...), and scratchpad crash cleanup behavior.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.

File	Description
tests/_runtime/test_runtime.py	Adds regression tests ensuring scratchpad + cache features don’t error and don’t pollute globals/graph.
marimo/_runtime/runtime.py	Registers `__scratch__` in the main graph during scratchpad execution and cleans it up reliably.

…ratchpad Three new tests under TestExecution in tests/_runtime/test_runtime.py, all using the existing `mocked_kernel` fixture and `run_scratchpad`: - test_run_scratch_with_mo_cache_decorator: asserts that decorating a function with @mo.cache inside the scratchpad does not leak KeyError: '__scratch__' to stderr, that the scratch cell is removed from the kernel graph after the run, and that the scratchpad does not pollute globals. - test_run_scratch_with_persistent_cache_context: same regression guard for the parallel `_cache_context.trace` path used by `with mo.persistent_cache(...)`. - test_run_scratch_with_mo_cache_cleans_up_after_crash: verifies the try/finally correctness — if scratchpad code raises after the cache decoration, __scratch__ is still unregistered from the main graph. These sit beside the three pre-existing test_run_scratch* tests in the same class and follow the same style.

dmadisetti · 2026-05-22T17:05:19Z

+        # cell's defs and refs are cleared above, so this adds no edges
+        # and has no reactive side effects. We unregister in a `finally`
+        # so a scratchpad crash cannot leave the main graph polluted.
+        did_register_in_main_graph = False


I think it would be better to put this directly in cache opposed to runtime?

There's also an executor refactor that could be a better place to put this if we did want to keep it in runtime.

But thanks for opening this!

Thanks for the speedy review @dmadisetti! Sounds good yep, cleaner in cache. Updated.

dmadisetti · 2026-05-22T17:22:48Z

Thanks @everettroeth ! This also just came on our radar. Let me know what you think about my comment of handling this in cache itself. Patching scratch into the graph may do weird things

Route SCRATCH_CELL_ID through the existing _external=True path in @mo.cache._set_context, mirroring how embedded app cells (UUID-prefixed) are already handled. In _cache_context.trace, fall back to linecache when the cell isn't in the kernel's main graph. Add defensive guards in mo.defs() / mo.refs() for the same root cause. Reverts the run_scratchpad main-graph mutation from earlier in this PR.

everettroeth · 2026-05-22T19:23:51Z

I have read the CLA Document and I hereby sign the CLA

dmadisetti · 2026-05-28T22:53:43Z

Nice! Thank you

github-actions · 2026-05-28T23:19:04Z

🚀 Development release published. You may be able to view the changes at https://marimo.app?v=0.23.9-dev18

Copilot AI review requested due to automatic review settings May 22, 2026 14:00

vercel Bot deployed to Preview May 22, 2026 14:01 View deployment

Copilot AI reviewed May 22, 2026

View reviewed changes

Comment thread marimo/_runtime/runtime.py Outdated

Comment thread tests/_runtime/test_runtime.py Outdated

Comment thread tests/_runtime/test_runtime.py

everettroeth force-pushed the fix/mo-cache-scratchpad-keyerror branch from 3b55599 to 271f80a Compare May 22, 2026 15:22

vercel Bot deployed to Preview May 22, 2026 15:23 View deployment

kirangadhave added the bug Something isn't working label May 22, 2026

dmadisetti reviewed May 22, 2026

View reviewed changes

vercel Bot deployed to Preview May 22, 2026 19:23 View deployment

everettroeth requested a review from dmadisetti May 22, 2026 19:25

dmadisetti approved these changes May 28, 2026

View reviewed changes

dmadisetti merged commit 5c1bb80 into marimo-team:main May 28, 2026
38 of 40 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

fix(cache): mo.cache raises KeyError: 'scratch' in scratchpad#9664

fix(cache): mo.cache raises KeyError: 'scratch' in scratchpad#9664
dmadisetti merged 3 commits into
marimo-team:mainfrom
everettroeth:fix/mo-cache-scratchpad-keyerror

everettroeth commented May 22, 2026

Uh oh!

vercel Bot commented May 22, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented May 22, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

dmadisetti May 22, 2026

Uh oh!

everettroeth May 22, 2026

Uh oh!

dmadisetti commented May 22, 2026

Uh oh!

everettroeth commented May 22, 2026

Uh oh!

dmadisetti commented May 28, 2026

Uh oh!

Uh oh!

github-actions Bot commented May 28, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

everettroeth commented May 22, 2026

Why

What

Uh oh!

vercel Bot commented May 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions Bot commented May 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

dmadisetti May 22, 2026

Choose a reason for hiding this comment

Uh oh!

everettroeth May 22, 2026

Choose a reason for hiding this comment

Uh oh!

dmadisetti commented May 22, 2026

Uh oh!

everettroeth commented May 22, 2026

Uh oh!

dmadisetti commented May 28, 2026

Uh oh!

Uh oh!

github-actions Bot commented May 28, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

vercel Bot commented May 22, 2026 •

edited

Loading

github-actions Bot commented May 22, 2026 •

edited

Loading