Repository: marimo-team/marimo
Pull request #4023: mep/8: introduce setup block
Author: login: dmadisetti
Created: 2025-03-07T13:56:29Z | Merged: 2025-03-13T02:02:00Z

## 📝 Summary

Towards mep/8;  This introduces the setup block

 - refactored `_save` to share code
 - added the `app.setup` context block to register cells
 - introduce `_ast/toplevel.py` to statically resolve the dependency graph for name errors (opposed to the quick check I was doing before)

setup block is just found by its name `setup` for now, but this can be changed when setup block somehow reserved (or we could just leverage the same mechanism as an easy work around, but build on top).


@akshayka OR @mscolnick



Diff:
diff --git a/marimo/__init__.py b/marimo/__init__.py
index c6ff7974a2a..650e469ce90 100644
--- a/marimo/__init__.py
+++ b/marimo/__init__.py
@@ -46,7 +46,6 @@
     "icon",
     "iframe",
     "image",
-    "import_guard",
     "latex",
     "lazy",
     "left",
@@ -124,7 +123,7 @@
     redirect_stderr,
     redirect_stdout,
 )
-from marimo._runtime.context.utils import import_guard, running_in_notebook
+from marimo._runtime.context.utils import running_in_notebook
 from marimo._runtime.control_flow import MarimoStopError, stop
 from marimo._runtime.runtime import (
     app_meta,
diff --git a/marimo/_ast/app.py b/marimo/_ast/app.py
index 2ae82bb2e7b..955b7316f81 100644
--- a/marimo/_ast/app.py
+++ b/marimo/_ast/app.py
@@ -3,6 +3,7 @@
 
 import ast
 import base64
+import builtins
 import inspect
 import sys
 import threading
@@ -57,9 +58,11 @@
     FunctionCallRequest,
     SetUIElementValueRequest,
 )
+from marimo._utils.with_skip import SkipContext
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
+    from types import TracebackType
 
     from marimo._messaging.ops import HumanReadableStatus
     from marimo._plugins.core.web_component import JSONType
@@ -149,6 +152,40 @@ def _mime_(self) -> tuple[KnownMimeType, str]:
         return tree(self._dict)._mime_()
 
 
+class _SetupContext(SkipContext):
+    """
+    A context manager that controls imports from being executed in top level code.
+    See design discussion in MEP-0008 (github:marimo-team/meps/pull/8).
+    """
+
+    def __init__(self, cell: Cell):
+        super().__init__()
+        self._cell = cell
+
+    def trace(self, _frame: Any) -> None:
+        if self._cell.refs - set(builtins.__dict__.keys()):
+            self.skip()
+
+    def __exit__(
+        self,
+        exception: Optional[type[BaseException]],
+        instance: Optional[BaseException],
+        _tracebacktype: Optional[TracebackType],
+    ) -> Literal[False]:
+        # Must be a Literal[False], for linters.
+        # Whether to suppress a given exception.
+        self.teardown()
+
+        if exception is not None:
+            LOGGER.warning(
+                "The setup cell was unable to execute, your notebook may not "
+                "work as expected."
+            )
+            LOGGER.debug("Exception: %s", exception)
+            return True  # type: ignore
+        return False
+
+
 @dataclass
 class AppEmbedResult:
     output: Html
@@ -377,6 +414,28 @@ def multiply(a: int, b: int) -> int:
             ),
         )
 
+    @property
+    def setup(self) -> _SetupContext:
+        """Provides a context manager to initialize the setup cell.
+
+        This block should only be utilized at the start of a marimo notebook.
+        It's used as following:
+
+        ```
+        with app.setup:
+            import my_libraries
+            from typing import Any
+
+            CONSTANT = "my constant"
+        ```
+        """
+        # Get the calling context to extract the location of the cell
+        frame = inspect.stack()[1].frame
+        cell = self._cell_manager.cell_context(
+            app=InternalApp(self), frame=frame
+        )
+        return _SetupContext(cell)
+
     def _unparsable_cell(
         self,
         code: str,
diff --git a/marimo/_ast/cell.py b/marimo/_ast/cell.py
index 9ddfc035e06..95f6e27d1f5 100644
--- a/marimo/_ast/cell.py
+++ b/marimo/_ast/cell.py
@@ -290,48 +290,34 @@ def namespace_to_variable(self, namespace: str) -> Name | None:
     def is_coroutine(self) -> bool:
         return _is_coroutine(self.body) or _is_coroutine(self.last_expr)
 
-    def is_toplevel_acceptable(
-        self, allowed_refs: Optional[set[Name]] = None
-    ) -> bool:
-        # Check no defs aside from the single function
-        if len(self.defs) != 1:
-            return False
-
-        # Must be parsable
+    @property
+    def toplevel_variable(self) -> Optional[VariableData]:
+        """Return the single, scoped, toplevel variable defined if found."""
         try:
             tree = ast.parse(self.code)
         except SyntaxError:
-            return False
+            return None
+
+        if len(self.defs) != 1:
+            return None
 
-        # Must be a bare function.
         if not (
             len(tree.body) == 1
             and isinstance(
                 tree.body[0], (ast.FunctionDef, ast.AsyncFunctionDef)
             )
         ):
-            return False
+            return None
 
         # Check that def matches the single definition
         name = tree.body[0].name
         if not (name == list(self.defs)[0] and name in self.variable_data):
-            return False
+            return None
 
-        # No required_refs are allowed for now
-        refs = set().union(
-            *[v.required_refs for v in self.variable_data[name]]
-        )
-        # NOTE: Builtins are allowed, but should be passed in under
-        # allowed_refs. Defers to allowed_refs because shadowed builtins
-        # are accounted for.
-        if allowed_refs is None:
-            allowed_refs = set(globals()["__builtins__"].keys())
-        # Allow recursion
-        refs -= {name} | allowed_refs
-        if refs:
-            return False
-
-        return True
+        if len(variable_data := self.variable_data[name]) != 1:
+            return None
+
+        return list(variable_data)[0]
 
     def set_runtime_state(
         self, status: RuntimeStateType, stream: Stream | None = None
diff --git a/marimo/_ast/cell_manager.py b/marimo/_ast/cell_manager.py
index db96b2e5b9f..2ee858b5dd9 100644
--- a/marimo/_ast/cell_manager.py
+++ b/marimo/_ast/cell_manager.py
@@ -18,7 +18,11 @@
     from typing import ParamSpec, TypeAlias
 
 from marimo._ast.cell import Cell, CellConfig
-from marimo._ast.compiler import cell_factory, toplevel_cell_factory
+from marimo._ast.compiler import (
+    cell_factory,
+    context_cell_factory,
+    toplevel_cell_factory,
+)
 from marimo._ast.models import CellData
 from marimo._ast.names import DEFAULT_CELL_NAME
 from marimo._ast.pytest import process_for_pytest
@@ -26,6 +30,7 @@
 
 if TYPE_CHECKING:
     from collections.abc import Iterable
+    from types import FrameType
 
     from marimo._ast.app import InternalApp
 
@@ -134,6 +139,21 @@ def decorator(func: Fn[P, R]) -> Cell | Fn[P, R]:
         else:
             return _register(func)
 
+    def cell_context(
+        self,
+        frame: FrameType,
+        app: InternalApp | None = None,
+    ) -> Cell:
+        """Registers cells when called through a context block."""
+        cell = context_cell_factory(
+            cell_id=self.create_cell_id(),
+            # NB. carry along the frame from the initial call.
+            frame=frame,
+        )
+        cell._cell.configure(CellConfig())
+        self._register_cell(cell, app=app)
+        return cell
+
     def _register_cell(
         self, cell: Cell, app: InternalApp | None = None
     ) -> None:
diff --git a/marimo/_ast/codegen.py b/marimo/_ast/codegen.py
index 3b28f082965..4d79c8b960d 100644
--- a/marimo/_ast/codegen.py
+++ b/marimo/_ast/codegen.py
@@ -14,8 +14,8 @@
 from marimo._ast.app import App, _AppConfig
 from marimo._ast.cell import CellConfig, CellImpl
 from marimo._ast.compiler import compile_cell
-from marimo._ast.names import DEFAULT_CELL_NAME
-from marimo._ast.transformers import RemoveImportTransformer
+from marimo._ast.names import DEFAULT_CELL_NAME, SETUP_CELL_NAME
+from marimo._ast.toplevel import TopLevelExtraction, TopLevelStatus
 from marimo._ast.visitor import Name
 from marimo._types.ids import CellId_t
 
@@ -30,6 +30,24 @@
 )
 
 
+def get_setup_cell(
+    code: list[str],
+    names: list[str],
+    configs: list[CellConfig],
+    toplevel_fn: bool,
+) -> Optional[CellImpl]:
+    # Find the cell named setup, compile, and remove the index from all lists.
+    if SETUP_CELL_NAME not in names or not toplevel_fn:
+        return None
+    setup_index = names.index(SETUP_CELL_NAME)
+    setup_code = code.pop(setup_index)
+    names.pop(setup_index)
+    setup_config = configs.pop(setup_index)
+    return compile_cell(
+        setup_code, cell_id=CellId_t(SETUP_CELL_NAME)
+    ).configure(setup_config)
+
+
 def indent_text(text: str) -> str:
     return textwrap.indent(text, INDENT)
 
@@ -107,38 +125,19 @@ def to_decorator(
         )
 
 
-def build_import_section(import_blocks: list[str]) -> str:
-    from marimo._utils.formatter import Formatter, ruff
-
-    stripped_block = RemoveImportTransformer("marimo").strip_imports(
-        "\n".join(import_blocks)
-    )
-    if not stripped_block:
+def build_setup_section(setup_cell: Optional[CellImpl]) -> str:
+    if setup_cell is None:
         return ""
-
-    code = "\n".join(
+    block = setup_cell.code
+    if not block.strip():
+        return ""
+    return "\n".join(
         [
-            "with marimo.import_guard():",
-            indent_text(NOTICE),
-            indent_text(stripped_block),
+            "with app.setup:",
+            indent_text(block),
         ]
     )
 
-    stub_cell_id = CellId_t("sub")
-
-    formatted = Formatter(MAX_LINE_LENGTH).format({stub_cell_id: code})
-    if not formatted:
-        return code
-    try:
-        # Note F401 strips unused imports, which is this whole block.
-        tidied = ruff(formatted, "check", "--fix-only", "--ignore=F401")
-        if tidied:
-            return tidied[stub_cell_id] + "\n\n"
-    # Thrown in WASM
-    except OSError:
-        pass
-    return formatted[stub_cell_id]
-
 
 def to_functiondef(
     cell: CellImpl,
@@ -193,9 +192,7 @@ def to_top_functiondef(
     # the cell, it must pass basic checks in the cell impl.
     if allowed_refs is None:
         allowed_refs = set(builtins.__dict__.keys())
-    assert cell.is_toplevel_acceptable(allowed_refs), (
-        "Cell is not a top-level function"
-    )
+    assert cell.toplevel_variable, "Cell is not a top-level function"
     if cell.code:
         decorator = to_decorator(cell.config, fn="function")
         return "\n".join([decorator, cell.code.strip()])
@@ -232,6 +229,34 @@ def generate_unparsable_cell(
     return "\n".join(text)
 
 
+def serialize(
+    extraction: TopLevelExtraction, status: TopLevelStatus, toplevel_fn: bool
+) -> str:
+    if status.is_unparsable:
+        return generate_unparsable_cell(
+            code=status.code, config=status.cell_config, name=status.name
+        )
+
+    cell = status._cell
+    assert cell is not None
+    if not toplevel_fn:
+        return to_functiondef(
+            cell, status.name, extraction.unshadowed, None, fn="cell"
+        )
+    elif status.is_cell:
+        return to_functiondef(
+            cell,
+            status.name,
+            extraction.allowed_refs,
+            extraction.used_refs,
+            fn="cell",
+        )
+    elif status.is_toplevel:
+        return to_top_functiondef(cell, extraction.allowed_refs)
+    else:
+        raise ValueError("Unknown cell status, please report this issue.")
+
+
 def generate_app_constructor(config: Optional[_AppConfig]) -> str:
     default_config = _AppConfig().asdict()
     updates = {}
@@ -292,71 +317,16 @@ def generate_filecontents(
         if name == "__":
             names[idx] = DEFAULT_CELL_NAME
 
-    # We require 3 sweeps.
-    #  - One for compilation and import collection
-    #  - One for some basic static determination of top-level functions.
-    #  - And a final sweep for cleaner argument requirements
-    # (since we now know what's top-level).
-    defs: set[Name] = set()
-    toplevel_imports: set[Name] = set()
-    used_refs: Optional[set[Name]] = set()
-    import_blocks: list[str] = []
-
-    definition_stubs: list[Optional[CellImpl]] = [None] * len(codes)
-    definitions: list[Optional[str]] = [None] * len(codes)
-
-    cell: Optional[CellImpl]
-    for idx, (code, cell_config) in enumerate(zip(codes, cell_configs)):
-        try:
-            cell = compile_cell(code, cell_id=CellId_t(str(idx))).configure(
-                cell_config
-            )
-            defs |= cell.defs
-            assert isinstance(used_refs, set)
-            used_refs |= cell.refs
-            if cell.import_workspace.is_import_block:
-                # maybe a bug, but import_workspace.imported_defs does not
-                # contain the information we need.
-                toplevel_imports |= cell.defs
-                definitions[idx] = to_functiondef(cell, names[idx])
-                if toplevel_fn:
-                    import_blocks.append(code.strip())
-            else:
-                definition_stubs[idx] = cell
-        except SyntaxError:
-            definitions[idx] = generate_unparsable_cell(
-                code=code, config=cell_config, name=names[idx]
-            )
-
-    unshadowed_builtins = set(builtins.__dict__.keys()) - defs
-    allowed_refs = unshadowed_builtins | toplevel_imports
-
-    for idx, cell in enumerate(definition_stubs):
-        # We actually don't care about the graph, we just want to see if we can
-        # render the top-level functions without a name error.
-        # Let graph issues be delegated to the runtime.
-        if cell and toplevel_fn and cell.is_toplevel_acceptable(allowed_refs):
-            definitions[idx] = to_top_functiondef(cell, allowed_refs)
-            definition_stubs[idx] = None
-            # Order does matter since feasibly, an app.function could be a
-            # decorator for another.
-            allowed_refs.add(names[idx])
-
-    # Let's hide the new behavior for now.
-    # Removing the toplevel_fn check may produce a bit of churn,
-    # so let's release the serialization changes all together.
-    if not toplevel_fn:
-        allowed_refs = unshadowed_builtins
-        used_refs = None
-
-    for idx, cell in enumerate(definition_stubs):
-        if cell:
-            definitions[idx] = to_functiondef(
-                cell, names[idx], allowed_refs, used_refs, fn="cell"
-            )
+    # TODO: replace with dedicated cell
+    setup_cell = get_setup_cell(codes, names, cell_configs, toplevel_fn)
+    toplevel_defs: set[Name] = set()
+    if setup_cell:
+        toplevel_defs = set(setup_cell.defs)
+    extraction = TopLevelExtraction(codes, names, cell_configs, toplevel_defs)
+    cell_blocks = [
+        serialize(extraction, status, toplevel_fn) for status in extraction
+    ]
 
-    assert all(isinstance(d, str) for d in definitions)
-    cell_blocks: list[str] = cast(list[str], definitions)
     if not toplevel_fn:
         return _classic_export(cell_blocks, header_comments, config)
 
@@ -368,10 +338,10 @@ def generate_filecontents(
         [
             "import marimo",
             "",
-            build_import_section(import_blocks),
-            "",
             f'__generated_with = "{__version__}"',
             generate_app_constructor(config),
+            "",
+            build_setup_section(setup_cell),
             "\n",
             "\n\n\n".join(cell_blocks),
             "\n",
diff --git a/marimo/_ast/compiler.py b/marimo/_ast/compiler.py
index 05bd626b39c..89dd3c1ecab 100644
--- a/marimo/_ast/compiler.py
+++ b/marimo/_ast/compiler.py
@@ -19,6 +19,8 @@
     ImportWorkspace,
     SourcePosition,
 )
+from marimo._ast.names import SETUP_CELL_NAME
+from marimo._ast.transformers import ContainedExtractWithBlock
 from marimo._ast.variables import is_local
 from marimo._ast.visitor import ImportData, Name, ScopedVisitor
 from marimo._types.ids import CellId_t
@@ -28,6 +30,7 @@
 
 if TYPE_CHECKING:
     from collections.abc import Iterator
+    from types import FrameType
 
 
 def code_key(code: str) -> int:
@@ -296,6 +299,51 @@ def get_source_position(
     )
 
 
+def context_cell_factory(
+    cell_id: CellId_t,
+    frame: FrameType,
+    anonymous_file: bool = False,
+) -> Cell:
+    # Frame is from the initiating context block.
+    _, lnum = inspect.getsourcelines(frame)
+    source = inspect.getsource(frame)
+    lines = source.split("\n")
+
+    entry_line = frame.f_lineno
+    # Offset needed when called from within a function (e.g. for tests)
+    if lnum > 0:
+        entry_line += 1 - lnum
+
+    _, with_block = ContainedExtractWithBlock(entry_line).visit(
+        ast.parse(textwrap.dedent(source)).body  # type: ignore[arg-type]
+    )
+
+    start_node = with_block.body[0]
+    end_node = with_block.body[-1]
+    col_offset = start_node.col_offset
+    cell_code = textwrap.dedent(
+        "\n".join(lines[entry_line : end_node.end_lineno + 1])
+    ).rstrip()
+
+    source_position = None
+    if not anonymous_file:
+        source_position = SourcePosition(
+            filename=frame.f_code.co_filename,
+            lineno=start_node.lineno - 1,
+            col_offset=col_offset,
+        )
+
+    return Cell(
+        _name=SETUP_CELL_NAME,
+        _cell=compile_cell(
+            cell_code,
+            cell_id=cell_id,
+            source_position=source_position,
+            test_rewrite=False,
+        ),
+    )
+
+
 def toplevel_cell_factory(
     f: Callable[..., Any],
     cell_id: CellId_t,
diff --git a/marimo/_ast/names.py b/marimo/_ast/names.py
index 7464dfabd92..f05f3f54293 100644
--- a/marimo/_ast/names.py
+++ b/marimo/_ast/names.py
@@ -1,5 +1,6 @@
 # Copyright 2024 Marimo. All rights reserved.
 DEFAULT_CELL_NAME = "_"
+SETUP_CELL_NAME = "setup"
 
 
 def is_internal_cell_name(name: str) -> bool:
diff --git a/marimo/_ast/toplevel.py b/marimo/_ast/toplevel.py
new file mode 100644
index 00000000000..7ee9bc1a74e
--- /dev/null
+++ b/marimo/_ast/toplevel.py
@@ -0,0 +1,304 @@
+# Copyright 2024 Marimo. All rights reserved.
+from __future__ import annotations
+
+import builtins
+from enum import Enum
+from typing import TYPE_CHECKING, Optional
+
+from marimo._ast.app import InternalApp
+from marimo._ast.cell import CellConfig, CellImpl
+from marimo._ast.compiler import compile_cell
+from marimo._ast.visitor import Name
+from marimo._types.ids import CellId_t
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
+# Constant for easy reuse in tests
+HINT_UNPARSABLE = "Cannot parse cell."
+HINT_BAD_NAME = (
+    "Top level definitions cannot be named 'app', '__name__' or "
+    "'__generated_with'"
+)
+HINT_NOT_SINGLE = "Cell must contain exactly one function definition"
+HINT_ORDER_DEPENDENT = (
+    "Signature and decorators depend on {} defined out of correct cell order"
+)
+HINT_HAS_REFS = (
+    "Function contains references to variables {} which are not top level."
+)
+HINT_HAS_CLOSE_REFS = (
+    "Function contains references to variables {} which failed to be toplevel."
+)
+
+
+class TopLevelType(Enum):
+    CELL = "cell"
+    TOPLEVEL = "toplevel"
+    UNPARSABLE = "unparsable"
+    UNRESOLVED = "unresolved"
+
+
+class TopLevelStatus:
+    """Status of a cell during serialization"""
+
+    def __init__(
+        self,
+        cell_id: CellId_t,
+        code: str,
+        name: Name,
+        cell_config: CellConfig,
+        allowed_refs: set[Name],
+    ):
+        self.cell_id = cell_id
+        self.name = name
+        self._name = name
+        self._type: TopLevelType = TopLevelType.UNPARSABLE
+        self.dependencies: set[Name] = set()
+        self._cell: Optional[CellImpl] = None
+        self.hint: Optional[str] = None
+
+        self.code = code
+        self.cell_config = cell_config
+
+        self.update(allowed_refs)
+
+    @classmethod
+    def from_cell(
+        cls, cell: CellImpl, allowed_refs: set[Name]
+    ) -> TopLevelStatus:
+        return cls(cell.cell_id, cell.code, "_", CellConfig(), allowed_refs)
+
+    def update(
+        self,
+        allowed_refs: set[Name],
+        unresolved: set[Name] | None = None,
+        potential_refs: set[Name] | None = None,
+    ) -> None:
+        if potential_refs is None:
+            potential_refs = set()
+        if unresolved is None:
+            unresolved = set()
+
+        if self._cell is None:
+            try:
+                self._cell = compile_cell(
+                    self.code, cell_id=self.cell_id
+                ).configure(self.cell_config)
+            except SyntaxError:
+                # Keep default
+                self.type = TopLevelType.UNPARSABLE
+                self.hint = HINT_UNPARSABLE
+                return
+
+        # Check that def matches the single definition
+        var = self._cell.toplevel_variable
+        if var is None:
+            self.demote(HINT_NOT_SINGLE)
+            return
+
+        (self.name,) = self._cell.defs
+        if self.name in ("app", "__generated_with", "__name__"):
+            self.demote(HINT_BAD_NAME)
+            return
+
+        allowed_refs = set(allowed_refs) | {self.name}
+        order_dependent_refs = var.unbounded_refs - allowed_refs
+        if order_dependent_refs and order_dependent_refs - unresolved:
+            self.dependencies = order_dependent_refs
+            if not order_dependent_refs - (unresolved | potential_refs):
+                self.demote(HINT_ORDER_DEPENDENT.format(self.dependencies))
+            else:
+                self.demote(HINT_HAS_REFS.format(self.dependencies))
+            return
+
+        dependent_refs = var.required_refs - allowed_refs
+        if not dependent_refs:
+            self.type = TopLevelType.TOPLEVEL
+            return
+
+        defined_refs = dependent_refs - potential_refs
+        if not (defined_refs):
+            self.type = TopLevelType.UNRESOLVED
+            self.dependencies = dependent_refs
+            return
+
+        self.demote(HINT_HAS_REFS.format(defined_refs))
+
+    def demote(self, hint: str) -> None:
+        self.type = TopLevelType.CELL
+        self.hint = hint
+
+    @property
+    def defs(self) -> set[Name]:
+        if self._cell is None:
+            return set()
+        return self._cell.defs
+
+    @property
+    def refs(self) -> set[Name]:
+        if self._cell is None:
+            return set()
+        return self._cell.refs
+
+    @property
+    def type(self) -> TopLevelType:
+        return self._type
+
+    @type.setter
+    def type(self, value: TopLevelType) -> None:
+        self._type = value
+        self.hint = None
+
+    @property
+    def is_toplevel(self) -> bool:
+        return self.type == TopLevelType.TOPLEVEL
+
+    @property
+    def is_cell(self) -> bool:
+        return self.type == TopLevelType.CELL
+
+    @property
+    def is_unresolved(self) -> bool:
+        return self.type == TopLevelType.UNRESOLVED
+
+    @property
+    def is_unparsable(self) -> bool:
+        return self.type == TopLevelType.UNPARSABLE
+
+
+class TopLevelExtraction:
+    """Graph representation of cells for serialization"""
+
+    def __init__(
+        self,
+        codes: list[str],
+        names: list[str],
+        cell_configs: list[CellConfig],
+        toplevel_defs: set[Name],
+    ):
+        self.statuses: list[TopLevelStatus] = []
+
+        # Track resolution status
+        self.unparsable: dict[Name, TopLevelStatus] = {}
+        self.toplevel: dict[Name, TopLevelStatus] = {}
+        self.unresolved: dict[Name, TopLevelStatus] = {}
+        self.cells: dict[Name, TopLevelStatus] = {}
+        self._statuses: dict[TopLevelType, dict[Name, TopLevelStatus]] = {
+            TopLevelType.UNPARSABLE: self.unparsable,
+            TopLevelType.TOPLEVEL: self.toplevel,
+            TopLevelType.UNRESOLVED: self.unresolved,
+            TopLevelType.CELL: self.cells,
+        }
+
+        # Track definitions and references
+        defs: set[Name] = set()
+        refs: set[Name] = set()
+        self.allowed_refs: set[Name] = set(toplevel_defs)
+        # Run through and get deff + refs, and a naive attempt at resolving cell
+        # status.
+        for idx, (code, name, config) in enumerate(
+            zip(codes, names, cell_configs)
+        ):
+            status = TopLevelStatus(
+                CellId_t(str(idx)), code, name, config, self.allowed_refs
+            )
+            self._statuses[status.type][status.name] = status
+            self.statuses.append(status)
+            if not status.is_unparsable:
+                defs.update(status.defs)
+                refs.update(status.refs)
+            if status.is_toplevel:
+                self.allowed_refs.add(status.name)
+
+        # Refresh names
+        names = [status.name for status in self.statuses]
+
+        self.unshadowed = set(builtins.__dict__.keys()) - defs
+        self.allowed_refs.update(self.unshadowed)
+        self.used_refs = refs
+
+        # Now toplevel, "allowed" defs have been determined, we can resolve
+        # references, and potentially promote cells to toplevel.
+        unresolved: set[Name] = set()
+        for status in self.statuses:
+            if status.is_unresolved or status.is_cell:
+                self._statuses[status.type].pop(status.name)
+                status.update(self.allowed_refs, unresolved, set(names))
+                if status.is_toplevel:
+                    self.allowed_refs.add(status.name)
+                elif status.is_unresolved:
+                    unresolved.add(status.name)
+                self._statuses[status.type][status.name] = status
+
+        self._resolve_dependencies()
+        # Don't change names of objects that are not toplevel.
+        for status in self.cells.values():
+            status.name = status._name
+
+    def _resolve_dependencies(self) -> None:
+        """Resolve the dependencies of unresolved cells"""
+        visited: dict[Name, bool] = {}
+
+        def resolve(name: Name) -> bool:
+            """Recursively resolve dependencies"""
+            # Already processed
+            if name not in self.unresolved:
+                if name in visited:
+                    return visited[name]
+                elif name in self.toplevel:
+                    return True
+                elif name in self.cells or name in self.unparsable:
+                    return False
+                # We are in a cycle. It's safe to serialize, but marimo will
+                # throw an error if running directly. Note, that in library
+                # form, multi-recursion is still supported (since it's just raw
+                # python).
+                return True
+
+            status = self.unresolved.pop(name)
+            refs = status.refs
+
+            # Check if all dependencies are resolved
+            resolved = True
+            invalid = set()
+            for ref in refs:
+                if not resolve(ref):
+                    resolved = False
+                    invalid.add(ref)
+
+            if resolved:
+                # All dependencies resolved, mark as top-level
+                status.type = TopLevelType.TOPLEVEL
+                self.allowed_refs.add(name)
+            else:
+                # Could not resolve all dependencies, mark as cell
+                status.demote(HINT_HAS_CLOSE_REFS.format(invalid))
+            self._statuses[status.type][name] = status
+
+            visited[name] = resolved
+            return resolved
+
+        # Try to resolve all unresolved cells
+        names_to_resolve = list(self.unresolved.keys())
+        for name in names_to_resolve:
+            if name in self.unresolved:
+                _ = resolve(name)
+        assert not self.unresolved
+
+    @classmethod
+    def from_app(cls, app: InternalApp) -> TopLevelExtraction:
+        codes = list(app.cell_manager.codes())
+        names = list(app.cell_manager.names())
+        cell_configs = list(app.cell_manager.configs())
+
+        from marimo._ast.codegen import get_setup_cell
+
+        setup = get_setup_cell(codes, names, cell_configs, True)
+        if setup:
+            return cls(codes, names, cell_configs, setup.defs)
+        return cls(codes, names, cell_configs, set())
+
+    def __iter__(self) -> Iterator[TopLevelStatus]:
+        """Iterate over cell statuses"""
+        return iter(self.statuses)
diff --git a/marimo/_ast/transformers.py b/marimo/_ast/transformers.py
index 6173c4cbabe..986d4d9ce4e 100644
--- a/marimo/_ast/transformers.py
+++ b/marimo/_ast/transformers.py
@@ -2,7 +2,95 @@
 from __future__ import annotations
 
 import ast
-from typing import Any, Optional
+import inspect
+import textwrap
+from typing import TYPE_CHECKING, Any, Callable, Optional, cast
+
+from marimo._ast.variables import unmangle_local
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+ARG_PREFIX: str = "*"
+
+
+class BlockException(Exception):
+    pass
+
+
+def compiled_ast(block: Sequence[ast.AST | ast.stmt]) -> ast.Module:
+    return cast(
+        ast.Module,
+        compile(
+            ast.Module(cast(list[ast.stmt], block), type_ignores=[]),
+            # <ast> is non-standard as a filename, but easier to debug than
+            # <module> everywhere.
+            "<ast>",
+            mode="exec",
+            flags=ast.PyCF_ONLY_AST | ast.PyCF_ALLOW_TOP_LEVEL_AWAIT,
+            optimize=0,
+            dont_inherit=True,
+        ),
+    )
+
+
+def clean_to_modules(
+    pre_block: list[ast.AST], block: ast.With
+) -> tuple[ast.Module, ast.Module]:
+    """Standardizes a `with` block to modules.
+
+    Consider
+
+    >>> <pre_block>
+    >>> ...
+    >>> </pre_block>
+    >>> with fn(val:=other) as x:
+    >>>   <block>
+    >>>   ...
+    >>>   </block>
+
+    We want to compile <pre_block>, and <block> into two separate modules,
+    however, line 4 removes the context of "x" and "val", so this is adjusted
+    to:
+
+    >>> <pre_block>
+    >>> ...
+    >>> x = fn(val := other)
+    >>> </pre_block>
+    >>> <block>
+    >>> ...
+    >>> </block>
+    """
+    assert len(block.items) == 1, "Unexpected with block structure."
+    (with_block,) = block.items
+    initializer: ast.AST = with_block.context_expr
+    if with_block.optional_vars:
+        initializer = ast.Assign(
+            targets=[with_block.optional_vars],
+            value=cast(ast.expr, initializer),
+        )
+    else:
+        # Edgecase with no "as" clause.
+        initializer = ast.Expr(value=cast(ast.expr, initializer))
+    initializer.lineno = len(pre_block) + 1
+    initializer.col_offset = 0
+    pre_block.append(initializer)
+    return (compiled_ast(pre_block), compiled_ast(block.body))
+
+
+def strip_function(fn: Callable[..., Any]) -> ast.Module:
+    code, _ = inspect.getsourcelines(fn)
+    args = set(fn.__code__.co_varnames)
+    function_ast = ast.parse(textwrap.dedent("".join(code)))
+    body = function_ast.body.pop()
+    assert isinstance(body, (ast.FunctionDef, ast.AsyncFunctionDef)), (
+        "Expected a function definition"
+    )
+    extracted = ast.Module(body.body, type_ignores=[])
+    module = RemoveReturns().visit(extracted)
+    module = MangleArguments(args).visit(module)
+    assert isinstance(module, ast.Module), "Expected a module"
+    return module
 
 
 class NameTransformer(ast.NodeTransformer):
@@ -117,3 +205,186 @@ def visit_ImportFrom(
             or (not alias.asname and alias.name != name)
         ]
         return node if node.names else None
+
+
+class ExtractWithBlock(ast.NodeTransformer):
+    def __init__(
+        self,
+        line: int,
+        allowed_types: tuple[type, ...],
+        *arg: Any,
+        name: str = "context",
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(*arg, **kwargs)
+        self.target_line = line
+        if allowed_types is None:
+            allowed_types = (ast.With,)
+        self.allowed_types = allowed_types
+        self.name = name
+
+    def generic_visit(self, node: ast.AST) -> tuple[ast.Module, ast.Module]:  # type: ignore[override]
+        pre_block = []
+
+        # There are a few strange edges cases like:
+        # >>> with open("file.txt"): pass
+        #
+        # It's difficult to properly delineate the block, so if multiple
+        # things map to the same line number, it's best to throw an error.
+        on_line = []
+        previous = None
+
+        assert isinstance(node, list), "Unexpected block structure."
+        for n in node:
+            # There's a chance that the block is first evaluated somewhere in a
+            # multiline line expression, for instance:
+            # 1 >>> with cache as c:
+            # 2 >>>    a = [
+            # 3 >>>        f(x) # <<< first frame call is here, and not line 2
+            # 4 >>>    ]
+            # so check that the "target" line is in the interval
+            # (i.e. 1 <= 3 <= 4)
+            parent = None
+            if n.lineno < self.target_line:
+                pre_block.append(n)
+                previous = n
+            # the line is contained with this
+            elif n.lineno <= self.target_line <= n.end_lineno:
+                on_line.append(n)
+                parent = n
+            # The target line can easily be skipped if there are comments or
+            # white space or if the block is contained within another block.
+            else:
+                break
+
+        # Capture the edge case when the block is contained within another.
+        # These cases are explicitly restricted to If and other With blocks,
+        # excluding by omission try, for, classes and functions.
+        if len(on_line) == 0:
+            if isinstance(
+                previous,
+                self.allowed_types,
+            ):
+                on_line.append(previous)
+                # Captures both branches of If, and the With block.
+                bodies = (
+                    [previous.body, previous.orelse]
+                    if isinstance(previous, ast.If)
+                    else [previous.body]  # type: ignore[attr-defined]
+                )
+                for body in bodies:
+                    try:
+                        # Recursion by referring to the containing block also
+                        # captures the case where the target line number was not
+                        # exactly hit.
+                        # for instance:
+                        # 1 >>> if True:          # Captured ast node
+                        # 2 >>>     with fn() as x:
+                        # 3 >>>         with cache as c:
+                        # 4 >>>             a = 1 # <<< frame line
+                        # will recurse through here thrice to get to the frame
+                        # line.
+                        # NB. the "extracted" With block is the one that
+                        # invoked this call
+                        return ExtractWithBlock(
+                            self.target_line,
+                            self.allowed_types,
+                            name=self.name,
+                        ).generic_visit(
+                            body  # type: ignore[arg-type]
+                        )
+                    except BlockException:
+                        pass
+
+            else:
+                raise BlockException(
+                    f"{self.name} cannot be invoked within a {type(previous)} "
+                    f"block (try moving the block within the {self.name} "
+                    "scope)."
+                )
+        # Intentionally not elif (on_line can be added in previous block)
+        if len(on_line) == 1:
+            if parent and not isinstance(on_line[0], ast.With):
+                raise BlockException("Detected line is not a With statement.")
+            if not isinstance(on_line[0], ast.With):
+                raise BlockException(
+                    "Unconventional formatting may lead to unexpected behavior. "
+                    "Please format your code, and/or reduce nesting.\n"
+                    "For instance, the following is not supported:\n"
+                    f">>>> with {self.name}: a = 1 # all one line"
+                )
+            return clean_to_modules(pre_block, on_line[0])
+        # It should be possible to relate the lines with the AST,
+        # but reduce potential bugs by just throwing an error.
+        raise BlockException(
+            "Unable to determine structure your call. Please"
+            " report this to github:marimo-team/marimo/issues"
+        )
+
+
+class CacheExtractWithBlock(ExtractWithBlock):
+    def __init__(self, line: int, *arg: Any, **kwargs: Any) -> None:
+        name = kwargs.pop("name", "cache")
+        super().__init__(line, (ast.With, ast.If), *arg, name=name, **kwargs)
+
+
+class ContainedExtractWithBlock(ExtractWithBlock):
+    def __init__(self, line: int, *arg: Any, **kwargs: Any) -> None:
+        name = kwargs.pop("name", "app.setup")
+        super().__init__(
+            line,
+            (
+                ast.With,
+                ast.If,
+                ast.FunctionDef,
+                ast.ClassDef,
+                ast.AsyncFunctionDef,
+            ),
+            *arg,
+            name=name,
+            **kwargs,
+        )
+
+
+class MangleArguments(ast.NodeTransformer):
+    """Mangles arguments names to prevent shadowing issues in analysis."""
+
+    def __init__(
+        self,
+        args: set[str],
+        *arg: Any,
+        prefix: str = ARG_PREFIX,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(*arg, **kwargs)
+        self.prefix = prefix
+        self.args = args
+
+    def visit_Name(self, node: ast.Name) -> ast.Name:
+        if node.id in self.args:
+            node.id = f"{self.prefix}{node.id}"
+        return node
+
+
+class DeprivateVisitor(ast.NodeTransformer):
+    """Removes the mangling of private variables from a module."""
+
+    def visit_Name(self, node: ast.Name) -> ast.Name:
+        node.id = unmangle_local(node.id).name
+        return node
+
+    def generic_visit(self, node: ast.AST) -> ast.AST:
+        if hasattr(node, "name") and node.name:
+            node.name = unmangle_local(node.name).name
+        return super().generic_visit(node)
+
+
+class RemoveReturns(ast.NodeTransformer):
+    # NB: Won't work for generators since not replacing Yield.
+    # Note that functools caches the generator, which is then dequeue'd,
+    # so in that sense, it doesn't work either.
+    def visit_Return(self, node: ast.Return) -> ast.Expr:
+        expr = ast.Expr(value=cast(ast.expr, node.value))
+        expr.lineno = node.lineno
+        expr.col_offset = node.col_offset
+        return expr
diff --git a/marimo/_ast/visitor.py b/marimo/_ast/visitor.py
index 00b6110626f..c8ac47bee0e 100644
--- a/marimo/_ast/visitor.py
+++ b/marimo/_ast/visitor.py
@@ -5,6 +5,7 @@
 import itertools
 import sys
 from collections import defaultdict
+from copy import deepcopy
 from dataclasses import dataclass, field
 from typing import Callable, Literal, Optional, Union
 from uuid import uuid4
@@ -70,6 +71,9 @@ class VariableData:
     # while traversing the tree.
     required_refs: set[Name] = field(default_factory=set)
 
+    # A reference needed for function/ class definition
+    unbounded_refs: set[Name] = field(default_factory=set)
+
     # For kind == import
     import_data: Optional[ImportData] = None
 
@@ -328,6 +332,7 @@ def generic_visit(self, node: ast.AST) -> ast.AST:
                 # `U` should not be a ref
                 for child in node.type_params:
                     self.visit(child)
+
             # This will revisit the type_params, but that's okay because
             # visiting is idempotent
             super().generic_visit(node)
@@ -379,11 +384,27 @@ def generic_visit(self, node: ast.AST) -> ast.AST:
             super().generic_visit(node)
         return node
 
-    def _visit_and_get_refs(self, node: ast.AST) -> set[Name]:
+    def _visit_and_get_refs(
+        self, node: Union[ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef]
+    ) -> tuple[set[Name], set[Name]]:
         """Create a ref scope for the variable to be declared (e.g. function,
         class), visit the children the node, propagate the refs to the higher
-        scope and then return the refs."""
+        scope and then return the body refs and unbounded refs."""
+
         self.ref_stack.append(set())
+        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
+            # Handle function refs that are evaluated in the outer scope
+            # Remove the body, which keeps signature and non-scoped parts.
+            mock = deepcopy(node)
+            mock.body.clear()
+            self.generic_visit(mock)
+            # Collect the unbounded refs
+            unbounded_refs = set(self.ref_stack[-1])
+        else:
+            # TODO: Update for class_definition toplevel decorator.
+            unbounded_refs = set()
+
+        # Process the function body
         self.generic_visit(node)
         refs = self.ref_stack.pop()
         # The scope a level up from the one just investigated also is dependent
@@ -395,12 +416,13 @@ def _visit_and_get_refs(self, node: ast.AST) -> set[Name]:
         # the variable `foo` needs to be aware that it may require the ref `x`
         # during execution.
         self.ref_stack[-1].update(refs)
-        return refs
+        # Return both sets of refs
+        return refs, unbounded_refs
 
     # ClassDef and FunctionDef nodes don't have ast.Name nodes as children
     def visit_ClassDef(self, node: ast.ClassDef) -> ast.ClassDef:
         node.name = self._if_local_then_mangle(node.name)
-        refs = self._visit_and_get_refs(node)
+        refs, _ = self._visit_and_get_refs(node)
         self._define(
             node,
             node.name,
@@ -412,21 +434,29 @@ def visit_AsyncFunctionDef(
         self, node: ast.AsyncFunctionDef
     ) -> ast.AsyncFunctionDef:
         node.name = self._if_local_then_mangle(node.name)
-        refs = self._visit_and_get_refs(node)
+        refs, unbounded_refs = self._visit_and_get_refs(node)
         self._define(
             node,
             node.name,
-            VariableData(kind="function", required_refs=refs),
+            VariableData(
+                kind="function",
+                required_refs=refs,
+                unbounded_refs=unbounded_refs,
+            ),
         )
         return node
 
     def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.FunctionDef:
         node.name = self._if_local_then_mangle(node.name)
-        refs = self._visit_and_get_refs(node)
+        refs, unbounded_refs = self._visit_and_get_refs(node)
         self._define(
             node,
             node.name,
-            VariableData(kind="function", required_refs=refs),
+            VariableData(
+                kind="function",
+                required_refs=refs,
+                unbounded_refs=unbounded_refs,
+            ),
         )
         return node
 
diff --git a/marimo/_runtime/context/utils.py b/marimo/_runtime/context/utils.py
index 2cf8750bf09..10982b145cd 100644
--- a/marimo/_runtime/context/utils.py
+++ b/marimo/_runtime/context/utils.py
@@ -2,16 +2,13 @@
 from __future__ import annotations
 
 import os
-from typing import TYPE_CHECKING, Literal, Optional
+from typing import Literal, Optional
 
 from marimo._output.rich_help import mddoc
 from marimo._runtime.context import ContextNotInitializedError, get_context
 from marimo._server.model import SessionMode
 from marimo._utils.assert_never import assert_never
 
-if TYPE_CHECKING:
-    from types import TracebackType
-
 RunMode = Literal["run", "edit", "script", "test"]
 
 
@@ -60,23 +57,3 @@ def get_mode() -> Optional[RunMode]:
         return "test"
 
     return None
-
-
-class import_guard:
-    """
-    A context manager that controls imports from being executed in top level code.
-    NB. May be replaced for `import_guard -> bool`, experimental stub pending
-    MEP-0008 (github:marimo-team/meps/pull/8)
-    """
-
-    def __enter__(self) -> None:
-        pass
-
-    def __exit__(
-        self,
-        exception: Optional[type[BaseException]],
-        instance: Optional[BaseException],
-        _tracebacktype: Optional[TracebackType],
-    ) -> Literal[False]:
-        # Whether to suppress a given exception.
-        return False
diff --git a/marimo/_save/ast.py b/marimo/_save/ast.py
deleted file mode 100644
index 3b008aa8edc..00000000000
--- a/marimo/_save/ast.py
+++ /dev/null
@@ -1,235 +0,0 @@
-# Copyright 2024 Marimo. All rights reserved.
-from __future__ import annotations
-
-import ast
-import inspect
-import textwrap
-from typing import TYPE_CHECKING, Any, Callable, cast
-
-from marimo._ast.variables import unmangle_local
-
-if TYPE_CHECKING:
-    from collections.abc import Sequence
-
-ARG_PREFIX: str = "*"
-
-
-class BlockException(Exception):
-    pass
-
-
-def compiled_ast(block: Sequence[ast.AST | ast.stmt]) -> ast.Module:
-    return cast(
-        ast.Module,
-        compile(
-            ast.Module(cast(list[ast.stmt], block), type_ignores=[]),
-            # <ast> is non-standard as a filename, but easier to debug than
-            # <module> everywhere.
-            "<ast>",
-            mode="exec",
-            flags=ast.PyCF_ONLY_AST | ast.PyCF_ALLOW_TOP_LEVEL_AWAIT,
-            optimize=0,
-            dont_inherit=True,
-        ),
-    )
-
-
-def clean_to_modules(
-    pre_block: list[ast.AST], block: ast.With
-) -> tuple[ast.Module, ast.Module]:
-    """Standardizes a `with` block to modules.
-
-    Consider
-
-    >>> <pre_block>
-    >>> ...
-    >>> </pre_block>
-    >>> with fn(val:=other) as x:
-    >>>   <block>
-    >>>   ...
-    >>>   </block>
-
-    We want to compile <pre_block>, and <block> into two separate modules,
-    however, line 4 removes the context of "x" and "val", so this is adjusted
-    to:
-
-    >>> <pre_block>
-    >>> ...
-    >>> x = fn(val := other)
-    >>> </pre_block>
-    >>> <block>
-    >>> ...
-    >>> </block>
-    """
-    assert len(block.items) == 1, "Unexpected with block structure."
-    (with_block,) = block.items
-    initializer: ast.AST = with_block.context_expr
-    if with_block.optional_vars:
-        initializer = ast.Assign(
-            targets=[with_block.optional_vars],
-            value=cast(ast.expr, initializer),
-        )
-    else:
-        # Edgecase with no "as" clause.
-        initializer = ast.Expr(value=cast(ast.expr, initializer))
-    initializer.lineno = len(pre_block) + 1
-    initializer.col_offset = 0
-    pre_block.append(initializer)
-    return (compiled_ast(pre_block), compiled_ast(block.body))
-
-
-class ExtractWithBlock(ast.NodeTransformer):
-    def __init__(self, line: int, *arg: Any, **kwargs: Any) -> None:
-        super().__init__(*arg, **kwargs)
-        self.target_line = line
-
-    def generic_visit(self, node: ast.AST) -> tuple[ast.Module, ast.Module]:  # type: ignore[override]
-        pre_block = []
-
-        # There are a few strange edges cases like:
-        # >>> with open("file.txt"): pass
-        #
-        # It's difficult to properly delineate the block, so if multiple
-        # things map to the same line number, it's best to throw an error.
-        on_line = []
-        previous = None
-
-        assert isinstance(node, list), "Unexpected block structure."
-        for n in node:
-            # There's a chance that the block is first evaluated somewhere in a
-            # multiline line expression, for instance:
-            # 1 >>> with cache as c:
-            # 2 >>>    a = [
-            # 3 >>>        f(x) # <<< first frame call is here, and not line 2
-            # 4 >>>    ]
-            # so check that the "target" line is in the interval
-            # (i.e. 1 <= 3 <= 4)
-            parent = None
-            if n.lineno < self.target_line:
-                pre_block.append(n)
-                previous = n
-            # the line is contained with this
-            elif n.lineno <= self.target_line <= n.end_lineno:
-                on_line.append(n)
-                parent = n
-            # The target line can easily be skipped if there are comments or
-            # white space or if the block is contained within another block.
-            else:
-                break
-
-        # Capture the edge case when the block is contained within another.
-        # These cases are explicitly restricted to If and other With blocks,
-        # excluding by omission try, for, classes and functions.
-        if len(on_line) == 0:
-            if isinstance(previous, (ast.With, ast.If)):
-                on_line.append(previous)
-                # Captures both branches of If, and the With block.
-                bodies = (
-                    [previous.body]
-                    if isinstance(previous, ast.With)
-                    else [previous.body, previous.orelse]
-                )
-                for body in bodies:
-                    try:
-                        # Recursion by referring to the containing block also
-                        # captures the case where the target line number was not
-                        # exactly hit.
-                        # for instance:
-                        # 1 >>> if True:          # Captured ast node
-                        # 2 >>>     with fn() as x:
-                        # 3 >>>         with cache as c:
-                        # 4 >>>             a = 1 # <<< frame line
-                        # will recurse through here thrice to get to the frame
-                        # line.
-                        # NB. the "extracted" With block is the one that
-                        # invoked this call
-                        return ExtractWithBlock(
-                            self.target_line
-                        ).generic_visit(
-                            body  # type: ignore[arg-type]
-                        )
-                    except BlockException:
-                        pass
-
-            else:
-                raise BlockException(
-                    "persistent_cache cannot be invoked within a block "
-                    "(try moving the block within the persistent_cache scope)."
-                )
-        # Intentionally not elif (on_line can be added in previous block)
-        if len(on_line) == 1:
-            if parent and not isinstance(on_line[0], ast.With):
-                raise BlockException("Detected line is not a With statement.")
-            if not isinstance(on_line[0], ast.With):
-                raise BlockException(
-                    "Unconventional formatting may lead to unexpected behavior. "
-                    "Please format your code, and/or reduce nesting.\n"
-                    "For instance, the following is not supported:\n"
-                    ">>>> with cache() as c: a = 1 # all one line"
-                )
-            return clean_to_modules(pre_block, on_line[0])
-        # It should be possible to relate the lines with the AST,
-        # but reduce potential bugs by just throwing an error.
-        raise BlockException(
-            "Unable to determine structure your call. Please"
-            " report this to github:marimo-team/marimo/issues"
-        )
-
-
-class MangleArguments(ast.NodeTransformer):
-    """Mangles arguments names to prevent shadowing issues in analysis."""
-
-    def __init__(
-        self,
-        args: set[str],
-        *arg: Any,
-        prefix: str = ARG_PREFIX,
-        **kwargs: Any,
-    ) -> None:
-        super().__init__(*arg, **kwargs)
-        self.prefix = prefix
-        self.args = args
-
-    def visit_Name(self, node: ast.Name) -> ast.Name:
-        if node.id in self.args:
-            node.id = f"{self.prefix}{node.id}"
-        return node
-
-
-class DeprivateVisitor(ast.NodeTransformer):
-    """Removes the mangling of private variables from a module."""
-
-    def visit_Name(self, node: ast.Name) -> ast.Name:
-        node.id = unmangle_local(node.id).name
-        return node
-
-    def generic_visit(self, node: ast.AST) -> ast.AST:
-        if hasattr(node, "name") and node.name:
-            node.name = unmangle_local(node.name).name
-        return super().generic_visit(node)
-
-
-class RemoveReturns(ast.NodeTransformer):
-    # NB: Won't work for generators since not replacing Yield.
-    # Note that functools caches the generator, which is then dequeue'd,
-    # so in that sense, it doesn't work either.
-    def visit_Return(self, node: ast.Return) -> ast.Expr:
-        expr = ast.Expr(value=cast(ast.expr, node.value))
-        expr.lineno = node.lineno
-        expr.col_offset = node.col_offset
-        return expr
-
-
-def strip_function(fn: Callable[..., Any]) -> ast.Module:
-    code, _ = inspect.getsourcelines(fn)
-    args = set(fn.__code__.co_varnames)
-    function_ast = ast.parse(textwrap.dedent("".join(code)))
-    body = function_ast.body.pop()
-    assert isinstance(body, (ast.FunctionDef, ast.AsyncFunctionDef)), (
-        "Expected a function definition"
-    )
-    extracted = ast.Module(body.body, type_ignores=[])
-    module = RemoveReturns().visit(extracted)
-    module = MangleArguments(args).visit(module)
-    assert isinstance(module, ast.Module), "Expected a module"
-    return module
diff --git a/marimo/_save/hash.py b/marimo/_save/hash.py
index 7dbde1ebc24..ae46aede019 100644
--- a/marimo/_save/hash.py
+++ b/marimo/_save/hash.py
@@ -10,6 +10,7 @@
 import types
 from typing import TYPE_CHECKING, Any, Callable, NamedTuple, Optional
 
+from marimo._ast.transformers import DeprivateVisitor, strip_function
 from marimo._ast.variables import (
     get_cell_from_local,
     if_local_then_mangle,
@@ -27,7 +28,6 @@
     is_pure_function,
 )
 from marimo._runtime.state import SetFunctor, State
-from marimo._save.ast import DeprivateVisitor, strip_function
 from marimo._save.cache import Cache, CacheType
 from marimo._types.ids import CellId_t
 
diff --git a/marimo/_save/save.py b/marimo/_save/save.py
index 016109e22fe..b31544cbab2 100644
--- a/marimo/_save/save.py
+++ b/marimo/_save/save.py
@@ -22,11 +22,15 @@
     overload,
 )
 
+from marimo._ast.transformers import (
+    ARG_PREFIX,
+    CacheExtractWithBlock,
+    strip_function,
+)
 from marimo._ast.variables import is_mangled_local, unmangle_local
 from marimo._messaging.tracebacks import write_traceback
 from marimo._runtime.context import get_context
 from marimo._runtime.state import State
-from marimo._save.ast import ARG_PREFIX, ExtractWithBlock, strip_function
 from marimo._save.cache import Cache, CacheException
 from marimo._save.hash import (
     DEFAULT_HASH,
@@ -43,6 +47,7 @@
     MemoryLoader,
 )
 from marimo._types.ids import CellId_t
+from marimo._utils.with_skip import SkipContext
 
 # Many assertions are for typing and should always pass. This message is a
 # catch all to motive users to report if something does fail.
@@ -56,16 +61,9 @@
 if TYPE_CHECKING:
     from types import FrameType, TracebackType
 
-    from _typeshed import TraceFunction
-    from typing_extensions import Self
-
     from marimo._runtime.dataflow import DirectedGraph
 
 
-class SkipWithBlock(Exception):
-    """Special exception to get around executing the with block body."""
-
-
 class _cache_call:
     """Like functools.cache but notebook-aware. See `cache` docstring`"""
 
@@ -215,7 +213,7 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any:
         return response
 
 
-class _cache_context:
+class _cache_context(SkipContext):
     def __init__(
         self,
         name: str,
@@ -224,55 +222,33 @@ def __init__(
         pin_modules: bool = False,
         hash_type: str = DEFAULT_HASH,
     ) -> None:
-        # For an implementation sibling regarding the block skipping, see
-        # `withhacks` in pypi.
+        super().__init__()
         self.name = name
 
-        self._skipped = True
         self._cache: Optional[Cache] = None
-        self._entered_trace = False
-        self._old_trace: Optional[TraceFunction] = None
-        self._frame: Optional[FrameType] = None
         self._body_start: int = MAXINT
         # TODO: Consider having a user level setting.
         self.pin_modules = pin_modules
         self.hash_type = hash_type
         self._loader = loader
 
-    def __enter__(self) -> Self:
-        sys.settrace(lambda *_args, **_keys: None)
-        frame = sys._getframe(1)
-        # Hold on to the previous trace.
-        self._old_trace = frame.f_trace
-        # Setting the frametrace, will cause the function to be run on _every_
-        # single context call until the trace is cleared.
-        frame.f_trace = self._trace
-        return self
-
-    def _trace(
-        self, with_frame: FrameType, _event: str, _arg: Any
-    ) -> Union[TraceFunction, None]:
+    def trace(self, with_frame: FrameType) -> None:
         # General flow is as follows:
         #   1) Follow the stack trace backwards to the first instance of a
         # "<module>" function call, which corresponds to a cell level block.
         #   2) Run static analysis to determine whether the call meets our
         # criteria. The procedure is a little brittle as such, certain contexts
-        # are not allow (e.g. called within a function or a loop).
+        # are not allowed (e.g. called within a function or a loop).
         #  3) Hash the execution and lookup the cache, and return!
         #  otherwise) Set _skipped such that the block continues to execute.
 
-        self._entered_trace = True
-
-        if not self._skipped:
-            return self._old_trace
-
         # This is possible if `With` spans multiple lines.
         # This behavior arguably a python bug.
         # Note the behavior does subtly change in 3.14, but will still be
         # captured by this check.
         if self._cache and self._cache.hit:
             if with_frame.f_lineno >= self._body_start:
-                raise SkipWithBlock()
+                self.skip()
 
         stack = traceback.extract_stack()
 
@@ -294,7 +270,9 @@ def _trace(
                     )
                 graph = ctx.graph
                 cell_id = ctx.cell_id or ctx.execution_context.cell_id
-                pre_module, save_module = ExtractWithBlock(lineno - 1).visit(
+                pre_module, save_module = CacheExtractWithBlock(
+                    lineno - 1
+                ).visit(
                     ast.parse(graph.cells[cell_id].code).body  # type: ignore[arg-type]
                 )
 
@@ -315,13 +293,11 @@ def _trace(
                 self._body_start = save_module.body[0].lineno
                 if self._cache and self._cache.hit:
                     if lineno >= self._body_start:
-                        raise SkipWithBlock()
-                    return self._old_trace
-                self._skipped = False
-                return self._old_trace
-            elif i > 1:
+                        self.skip()
+                return
+            elif i > 2:
                 raise CacheException(
-                    "`persistent_cache` must be invoked from cell level "
+                    "`cache` must be invoked from cell level "
                     "(cannot be in a function or class)"
                 )
         raise CacheException(
@@ -335,8 +311,8 @@ def __exit__(
         instance: Optional[BaseException],
         _tracebacktype: Optional[TracebackType],
     ) -> bool:
-        sys.settrace(self._old_trace)  # Clear to previous set trace.
-        if not self._entered_trace:
+        self.teardown()
+        if not self.entered_trace:
             raise CacheException(
                 f"Unexpected block format {UNEXPECTED_FAILURE_BOILERPLATE}"
             )
@@ -350,9 +326,6 @@ def __exit__(
 
         # NB: exception is a type.
         if exception:
-            assert not isinstance(instance, SkipWithBlock), (
-                f"Cache was not correctly set {UNEXPECTED_FAILURE_BOILERPLATE}"
-            )
             if isinstance(instance, BaseException):
                 raise instance from CacheException("Failure during save.")
             raise exception
diff --git a/marimo/_utils/with_skip.py b/marimo/_utils/with_skip.py
new file mode 100644
index 00000000000..12942c3ddb2
--- /dev/null
+++ b/marimo/_utils/with_skip.py
@@ -0,0 +1,69 @@
+# Copyright 2025 Marimo. All rights reserved.
+from __future__ import annotations
+
+import sys
+from abc import ABC, abstractmethod
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    NoReturn,
+    Optional,
+    Union,
+)
+
+from typing_extensions import Self
+
+if TYPE_CHECKING:
+    from types import FrameType
+
+    from _typeshed import TraceFunction
+
+
+class SkipContext(ABC):
+    def __init__(self) -> None:
+        # For an implementation sibling regarding the block skipping, see
+        # `withhacks` in pypi.
+        self._entered_trace = False
+        self._old_trace: Optional[TraceFunction] = None
+        self._frame: Optional[FrameType] = None
+        self._skipped = True
+
+    def __enter__(self) -> Self:
+        sys.settrace(lambda *_args, **_keys: None)
+        frame = sys._getframe(1)
+        # Hold on to the previous trace.
+        self._old_trace = frame.f_trace
+        # Setting the frametrace, will cause the function to be run on _every_
+        # single context call until the trace is cleared.
+        frame.f_trace = self._trace
+        return self
+
+    def skip(self) -> NoReturn:
+        raise SkipWithBlock()
+
+    def _trace(
+        self, with_frame: FrameType, _event: str, _arg: Any
+    ) -> Union[TraceFunction, None]:
+        self._entered_trace = True
+
+        if not self._skipped:
+            return self._old_trace
+
+        self.trace(with_frame)
+        self._skipped = False
+        return self._old_trace
+
+    @abstractmethod
+    def trace(self, with_frame: FrameType) -> None:
+        pass
+
+    @property
+    def entered_trace(self) -> bool:
+        return self._entered_trace
+
+    def teardown(self) -> None:
+        sys.settrace(self._old_trace)  # Clear to previous set trace.
+
+
+class SkipWithBlock(Exception):
+    """Special exception to get around executing the with block body."""
diff --git a/tests/_ast/codegen_data/test_generate_filecontents_toplevel.py b/tests/_ast/codegen_data/test_generate_filecontents_toplevel.py
index 0e4d958d203..acbb70a9798 100644
--- a/tests/_ast/codegen_data/test_generate_filecontents_toplevel.py
+++ b/tests/_ast/codegen_data/test_generate_filecontents_toplevel.py
@@ -6,9 +6,10 @@
 
 import marimo
 
-with marimo.import_guard():
-    # Imports used by your notebook.
-    # This section is automatically generated by marimo.
+__generated_with = "0.0.0"
+app = marimo.App(_toplevel_fn=True)
+
+with app.setup:
     import io
     import textwrap
     import typing
@@ -17,12 +18,8 @@
     import marimo as mo
 
 
-__generated_with = "0.0.0"
-app = marimo.App(_toplevel_fn=True)
-
-
 @app.cell
-def _(fun_that_uses_another_but_out_of_order):
+def _():
     shadow = 1
     globe = 1
     (
@@ -49,6 +46,7 @@ def shadow_case(shadow):
 def _(shadow):
     def reference_case():
         return shadow
+
     return
 
 
@@ -57,6 +55,7 @@ def _(globe):
     def global_case():
         global globe
         return globe
+
     return
 
 
@@ -65,11 +64,9 @@ def fun_that_uses_mo():
     return mo.md("Hello there!")
 
 
-@app.cell
+@app.function
 def fun_that_uses_another_but_out_of_order():
-    def fun_that_uses_another_but_out_of_order():
-        return fun_that_uses_another()
-    return (fun_that_uses_another_but_out_of_order,)
+    return fun_that_uses_another()
 
 
 @app.function
@@ -77,22 +74,5 @@ def fun_that_uses_another():
     return fun_that_uses_mo()
 
 
-@app.cell
-def _():
-    import io
-    import textwrap
-
-    # Comments are stripped out.
-    import marimo as mo
-    return io, mo, textwrap
-
-
-@app.cell
-def _():
-    import typing
-    from pathlib import Path
-    return Path, typing
-
-
 if __name__ == "__main__":
     app.run()
diff --git a/tests/_ast/test_codegen.py b/tests/_ast/test_codegen.py
index 79075d9b534..be5290eb535 100644
--- a/tests/_ast/test_codegen.py
+++ b/tests/_ast/test_codegen.py
@@ -321,7 +321,7 @@ def test_generate_filecontents_toplevel() -> None:
         assert "import marimo" in source
         split = source.split("import marimo")
         # The default one, the as mo in top level, in as mo in cell
-        assert len(split) == 4
+        assert len(split) == 3
 
     @staticmethod
     def test_generate_filecontents_toplevel_pytest() -> None:
diff --git a/tests/_ast/test_pytest_toplevel.py b/tests/_ast/test_pytest_toplevel.py
index dc22895be94..f3570d96c47 100644
--- a/tests/_ast/test_pytest_toplevel.py
+++ b/tests/_ast/test_pytest_toplevel.py
@@ -2,15 +2,13 @@
 
 import marimo
 
-with marimo.import_guard():
-    # Imports used by your notebook.
-    # This section is automatically generated by marimo.
-    import pytest
-
-
 __generated_with = "0.0.0"
 app = marimo.App(_toplevel_fn=True)
 
+with app.setup:
+    # Special setup cell
+    import pytest
+
 
 @app.function
 # Sanity check that base case works.
@@ -35,17 +33,5 @@ def test_add_bad(a, b, c):
     assert add(a, b) == c
 
 
-@app.cell
-def _():
-    # Invert the order intentionally.
-    import pytest
-
-    # Comments are stripped.
-    import marimo  # noqa: I001
-
-    # Final comments preserve total whitespace.
-    return marimo, pytest
-
-
 if __name__ == "__main__":
     app.run()
diff --git a/tests/_ast/test_toplevel.py b/tests/_ast/test_toplevel.py
new file mode 100644
index 00000000000..732389d5b7d
--- /dev/null
+++ b/tests/_ast/test_toplevel.py
@@ -0,0 +1,275 @@
+# Copyright 2024 Marimo. All rights reserved.
+
+from __future__ import annotations
+
+import builtins
+
+from marimo._ast import toplevel
+from marimo._ast.app import App, InternalApp
+from marimo._ast.toplevel import (
+    TopLevelExtraction,
+    TopLevelStatus,
+    TopLevelType,
+)
+
+BUILTINS = set(builtins.__dict__.keys())
+
+
+class TestTopLevelStatus:
+    @staticmethod
+    def test_expression_ignored(app) -> None:
+        @app.cell
+        def cell():
+            def add(a: int, b: int) -> int:
+                return a + b
+
+            pass
+
+        status = TopLevelStatus.from_cell(cell._cell, BUILTINS)
+        assert status.type == TopLevelType.CELL
+        assert status.hint == toplevel.HINT_NOT_SINGLE
+
+    @staticmethod
+    def test_function_converted(app) -> None:
+        @app.cell
+        def cell():
+            def add(a: int, b: int) -> int:
+                return a + b
+
+        status = TopLevelStatus.from_cell(cell._cell, BUILTINS)
+        assert status.type == TopLevelType.TOPLEVEL
+
+    @staticmethod
+    def test_function_uses_top(app) -> None:
+        with app.setup:
+            c = 1e-16
+
+        @app.cell
+        def cell():
+            def add(a: int, b: int, offset: float = c) -> float:
+                return a + b + offset
+
+        status = TopLevelStatus.from_cell(cell._cell, BUILTINS | {"c"})
+        assert status.type == TopLevelType.TOPLEVEL
+
+    @staticmethod
+    def test_function_uses_top_fn(app) -> None:
+        @app.cell
+        def _():
+            def c() -> float:
+                return 1e-16
+
+        @app.cell
+        def cell(c):
+            def add(a: int, b: int, offset: float = c()) -> float:
+                return a + b + offset
+
+        status = TopLevelStatus.from_cell(cell._cell, BUILTINS | {"c"})
+        assert status.type == TopLevelType.TOPLEVEL
+
+    @staticmethod
+    def test_function_uses_top_fn_bad_order(app) -> None:
+        @app.cell
+        def cell(c):
+            def add(a: int, b: int, offset: float = c()) -> float:
+                return a + b + offset
+
+        @app.cell
+        def _():
+            def c() -> float:
+                return 1e-16
+
+        status = TopLevelStatus.from_cell(cell._cell, BUILTINS)
+        status.update(BUILTINS, set(), {"c"})
+        assert status.type == TopLevelType.CELL
+        assert status.hint == toplevel.HINT_ORDER_DEPENDENT.format({"c"})
+
+    @staticmethod
+    def test_function_not_top(app) -> None:
+        @app.cell
+        def _():
+            c = 1e-16
+            return c
+
+        @app.cell
+        def cell(c):
+            def add(a: int, b: int, offset: float = c) -> float:
+                return a + b + offset
+
+        status = TopLevelStatus.from_cell(cell._cell, BUILTINS)
+        assert status.type == TopLevelType.CELL
+        assert status.hint == toplevel.HINT_HAS_REFS.format({"c"})
+
+    @staticmethod
+    def test_injection_name_ignored(app) -> None:
+        @app.cell
+        def app():
+            def app(a: int, b: int) -> int:
+                return a + b
+
+        status = TopLevelStatus.from_cell(app._cell, BUILTINS)
+        assert status.type == TopLevelType.CELL
+        assert status.hint == toplevel.HINT_BAD_NAME
+
+    @staticmethod
+    def test_function_recursion(app) -> None:
+        @app.cell
+        def cell():
+            def fib(n: int) -> int:
+                if n <= 1:
+                    return 1
+                return fib(n - 1) + fib(n - 2)
+
+        @app.cell
+        def _(fib):
+            assert fib(6) == 13
+
+        status = TopLevelStatus.from_cell(cell._cell, BUILTINS)
+        assert status.type == TopLevelType.TOPLEVEL
+
+
+class TestTopLevelExtraction:
+    @staticmethod
+    def test_function_uses_top(app) -> None:
+        with app.setup:
+            c = 1e-16
+
+        @app.cell
+        def cell():
+            def add(a: int, b: int, offset: float = c) -> float:
+                return a + b + offset
+
+        extraction = TopLevelExtraction.from_app(InternalApp(app))
+        assert [TopLevelType.TOPLEVEL] == [s.type for s in extraction]
+
+    @staticmethod
+    def test_function_uses_top_fn_unresolved(app) -> None:
+        @app.cell
+        def cell(c):
+            def add(a: int, b: int) -> int:
+                return a + b + c()
+
+        @app.cell
+        def _():
+            def c() -> float:
+                return 1e-16
+
+        status = TopLevelStatus.from_cell(cell._cell, BUILTINS)
+        status.update(BUILTINS, set(), {"c"})
+        assert status.type == TopLevelType.UNRESOLVED
+
+        extraction = TopLevelExtraction.from_app(InternalApp(app))
+        assert [TopLevelType.TOPLEVEL, TopLevelType.TOPLEVEL] == [
+            s.type for s in extraction
+        ], [s.hint for s in extraction]
+
+    @staticmethod
+    def test_function_uses_top_fn_unresolved_path(app) -> None:
+        @app.cell
+        def cell(A, Z):
+            def a() -> float:
+                return A() + Z()
+
+        @app.cell
+        def _(B):
+            def A() -> float:
+                return B()
+
+        @app.cell
+        def _():
+            def B() -> float:
+                return 1.0
+
+        @app.cell
+        def _():
+            def Y() -> float:
+                return 1.0
+
+        @app.cell
+        def _(Y):
+            def Z() -> float:
+                return Y()
+
+        status = TopLevelStatus.from_cell(cell._cell, BUILTINS)
+        status.update(BUILTINS, set(), {"A", "B", "Y", "Z"})
+        assert status.type == TopLevelType.UNRESOLVED
+
+        extraction = TopLevelExtraction.from_app(InternalApp(app))
+        assert [
+            TopLevelType.TOPLEVEL,
+        ] * 5 == [s.type for s in extraction], [s.hint for s in extraction]
+
+    @staticmethod
+    def test_function_uses_top_fn_unresolved_path_fails(app) -> None:
+        @app.cell
+        def cell(A, Z):
+            def a() -> float:
+                return A() + Z()
+
+        @app.cell
+        def _(B):
+            def A() -> float:
+                return B()
+
+        @app.cell
+        def _():
+            def B() -> float:
+                return 1.0
+
+        @app.cell
+        def _():
+            def Y() -> float:
+                return 1.0
+
+            pass  # to make it a cell
+
+        @app.cell
+        def _(Y):
+            def Z() -> float:
+                return Y()
+
+        status = TopLevelStatus.from_cell(cell._cell, BUILTINS)
+        status.update(BUILTINS, set(), {"A", "B", "Y", "Z"})
+        assert status.type == TopLevelType.UNRESOLVED
+
+        extraction = TopLevelExtraction.from_app(InternalApp(app))
+        assert [
+            TopLevelType.CELL,
+            TopLevelType.TOPLEVEL,
+            TopLevelType.TOPLEVEL,
+            TopLevelType.CELL,
+            TopLevelType.CELL,
+        ] == [s.type for s in extraction], [s.hint for s in extraction]
+        assert extraction.statuses[
+            0
+        ].hint == toplevel.HINT_HAS_CLOSE_REFS.format({"Z"})
+
+    @staticmethod
+    def test_indirect_recursion() -> None:
+        app = App()
+        # Needed for consistent stack trace paths.
+        app._anonymous_file = True
+
+        # Renders, but fails in runtime (which is fine).
+        @app.cell
+        def first(c):
+            def add(a: int, b: int) -> int:
+                return a + b + c()
+
+        @app.cell
+        def second(add):
+            def c() -> float:
+                return add(1, 1)
+
+        status = TopLevelStatus.from_cell(first._cell, BUILTINS)
+        status.update(BUILTINS, set(), {"c", "add"})
+        assert status.type == TopLevelType.UNRESOLVED
+
+        status = TopLevelStatus.from_cell(second._cell, BUILTINS)
+        status.update(BUILTINS, {"add"}, {"c", "add"})
+        assert status.type == TopLevelType.UNRESOLVED
+
+        extraction = TopLevelExtraction.from_app(InternalApp(app))
+        assert [TopLevelType.TOPLEVEL, TopLevelType.TOPLEVEL] == [
+            s.type for s in extraction
+        ], [s.hint for s in extraction]
diff --git a/tests/_ast/test_visitor.py b/tests/_ast/test_visitor.py
index 8333bcc5d41..6968066592f 100644
--- a/tests/_ast/test_visitor.py
+++ b/tests/_ast/test_visitor.py
@@ -425,7 +425,13 @@ def foo(x=y, y=x, z=a):
     assert v.refs == set(["x", "y", "a"])
     # TODO: Are these required refs?
     assert v.variable_data == {
-        "foo": [VariableData(kind="function", required_refs={"x", "y", "a"})],
+        "foo": [
+            VariableData(
+                kind="function",
+                required_refs={"x", "y", "a"},
+                unbounded_refs={"x", "y", "a"},
+            )
+        ],
     }
 
 
diff --git a/tests/_runtime/script_data/script_exception_function.py b/tests/_runtime/script_data/script_exception_function.py
new file mode 100644
index 00000000000..dff367ac398
--- /dev/null
+++ b/tests/_runtime/script_data/script_exception_function.py
@@ -0,0 +1,19 @@
+import marimo
+
+__generated_with = "0.6.19"
+app = marimo.App()
+
+
+@app.function
+def bad_divide(y: int) -> int:
+    return y / 0
+
+
+@app.cell
+def _():
+    x = bad_divide(100)
+    return (x,)
+
+
+if __name__ == "__main__":
+    app.run()
diff --git a/tests/_runtime/script_data/script_exception_setup_cell.py b/tests/_runtime/script_data/script_exception_setup_cell.py
new file mode 100644
index 00000000000..9994d391c3b
--- /dev/null
+++ b/tests/_runtime/script_data/script_exception_setup_cell.py
@@ -0,0 +1,14 @@
+import marimo
+
+__generated_with = "0.6.19"
+app = marimo.App()
+
+
+with app.setup:
+    y = 1
+    x = 0
+    y = y / x
+
+
+if __name__ == "__main__":
+    app.run()
diff --git a/tests/_runtime/test_trace.py b/tests/_runtime/test_trace.py
index 33d2b22c9bb..a961ac099b2 100644
--- a/tests/_runtime/test_trace.py
+++ b/tests/_runtime/test_trace.py
@@ -103,6 +103,39 @@ def test_script_trace_with_imported_file() -> None:
         )
         assert "y = y / x" in result
 
+    @staticmethod
+    def test_script_trace_function() -> None:
+        p = subprocess.run(
+            [
+                sys.executable,
+                "tests/_runtime/script_data/script_exception_function.py",
+            ],
+            capture_output=True,
+        )
+        assert p.returncode == 1
+
+        result = p.stderr.decode()
+        assert "ZeroDivisionError: division by zero" in result
+        assert ('script_exception_function.py", line 9') in result
+        assert "y / 0" in result
+
+    @staticmethod
+    def test_script_trace_setup_cell() -> None:
+        p = subprocess.run(
+            [
+                sys.executable,
+                "tests/_runtime/script_data/script_exception_setup_cell.py",
+            ],
+            capture_output=True,
+        )
+        assert p.returncode == 1
+
+        result = p.stderr.decode()
+        assert "The setup cell was unable to execute" in result
+        assert "ZeroDivisionError: division by zero" in result
+        assert ('script_exception_setup_cell.py", line 10') in result
+        assert "y / x" in result
+
 
 class TestAppTrace:
     @staticmethod
diff --git a/tests/_save/test_cache.py b/tests/_save/test_cache.py
index 33a9f697838..b382a8907ee 100644
--- a/tests/_save/test_cache.py
+++ b/tests/_save/test_cache.py
@@ -207,7 +207,7 @@ def called(v):
 
     @staticmethod
     def test_cache_same_line_fails() -> None:
-        from marimo._save.ast import BlockException
+        from marimo._ast.transformers import BlockException
 
         app = App()
         app._anonymous_file = True
@@ -230,7 +230,7 @@ def call(v):
 
     @staticmethod
     def test_cache_in_fn_fails() -> None:
-        from marimo._save.ast import BlockException
+        from marimo._ast.transformers import BlockException
 
         app = App()
         app._anonymous_file = True
diff --git a/tests/snapshots/api.txt b/tests/snapshots/api.txt
index 0b6d6cf7992..31b5951d442 100644
--- a/tests/snapshots/api.txt
+++ b/tests/snapshots/api.txt
@@ -32,7 +32,6 @@ hstack
 icon
 iframe
 image
-import_guard
 islands
   MarimoIslandGenerator
   MarimoIslandStub


Reviews:
user: login: mscolnick
body: this is AWESOME! 

want me to unblock the frontend?
state: DISMISSED
author_association: CONTRIBUTOR
submitted_at: 2025-03-07T16:57:21Z
user: login: mscolnick
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-07T17:11:29Z
user: login: dmadisetti
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-07T17:39:45Z
user: login: akshayka
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-07T18:03:37Z
user: login: akshayka
body: 🎉!!
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-07T18:27:01Z
user: login: dmadisetti
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-07T19:37:46Z
user: login: dmadisetti
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-07T19:40:52Z
user: login: akshayka
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-07T20:00:53Z
user: login: mscolnick
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-07T21:00:56Z
user: login: dmadisetti
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-07T21:05:03Z
user: login: akshayka
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-07T23:19:19Z
user: login: akshayka
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-09T15:58:01Z
user: login: dmadisetti
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-09T18:18:52Z
user: login: dmadisetti
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-10T15:21:17Z
user: login: dmadisetti
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-10T15:23:25Z
user: login: mscolnick
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-03-10T16:22:22Z
user: login: akshayka
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-03-13T02:01:38Z

Review comments:
diff_hunk: @@ -1,5 +1,6 @@
 # Copyright 2024 Marimo. All rights reserved.
 DEFAULT_CELL_NAME = "_"
+SETUP_CELL_NAME = "setup"
path: marimo/_ast/names.py
user: login: mscolnick
body: i wonder if it would be better to find by `cell_id` instead of name. lots of the data structures are keyed by ID, and would avoid conflicts of someone naming a cell `setup` (which we could still prevent)
created_at: 2025-03-07T17:11:29Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1,5 +1,6 @@
 # Copyright 2024 Marimo. All rights reserved.
 DEFAULT_CELL_NAME = "_"
+SETUP_CELL_NAME = "setup"
path: marimo/_ast/names.py
user: login: dmadisetti
body: Yes, and it's probably better to resolve now than get too deep.

But the display for the setup cell should maybe be different? Unsure whether it should just be a normally rendered cell
@akshayka did you have any more thoughts on this front?
created_at: 2025-03-07T17:39:44Z
author_association: COLLABORATOR
diff_hunk: @@ -1,5 +1,6 @@
 # Copyright 2024 Marimo. All rights reserved.
 DEFAULT_CELL_NAME = "_"
+SETUP_CELL_NAME = "setup"
path: marimo/_ast/names.py
user: login: akshayka
body: I think the display should be slightly different yes — with an explainer tooltip, maybe collapsed by default (not sure). Feel free to tap me in whenever the frontend is ready for styling/design.
created_at: 2025-03-07T18:03:37Z
author_association: CONTRIBUTOR
diff_hunk: @@ -149,6 +151,32 @@ def _mime_(self) -> tuple[KnownMimeType, str]:
         return tree(self._dict)._mime_()
 
 
+class _SetupContext(SkipContext):
+    """
+    A context manager that controls imports from being executed in top level code.
path: marimo/_ast/app.py
user: login: akshayka
body: This should skip everything in the setup cell — not just imports. Right? Similar to how when loading a marimo notebook for `marimo edit` none of the cell code is run.
created_at: 2025-03-07T18:17:22Z
author_association: CONTRIBUTOR
diff_hunk: @@ -377,6 +405,24 @@ def multiply(a: int, b: int) -> int:
             ),
         )
 
+    @property
path: marimo/_ast/app.py
user: login: akshayka
body: Should this be a function so it can be called with arguments in the future, should we need them?

I guess you could generalize by turning `setup` into a ContextManager that supports both parenthesized invocation and non-parenthesized, so maybe this is fine?
created_at: 2025-03-07T18:21:03Z
author_association: CONTRIBUTOR
diff_hunk: @@ -134,6 +138,20 @@ def decorator(func: Fn[P, R]) -> Cell | Fn[P, R]:
         else:
             return _register(func)
 
+    def cell_context(
+        self,
+        frame_offset: int = 2,
path: marimo/_ast/cell_manager.py
user: login: akshayka
body: Can you leave a comment explaining the frame offset and why it's `2` for us mere mortals? :)
created_at: 2025-03-07T18:22:12Z
author_association: CONTRIBUTOR
diff_hunk: @@ -296,6 +298,51 @@ def get_source_position(
     )
 
 
+def context_cell_factory(
+    cell_id: CellId_t,
+    frame_offset: int = 3,
path: marimo/_ast/compiler.py
user: login: akshayka
body: Same here, a comment explaining the frame offset please? :)
created_at: 2025-03-07T18:23:37Z
author_association: CONTRIBUTOR
diff_hunk: @@ -134,6 +138,20 @@ def decorator(func: Fn[P, R]) -> Cell | Fn[P, R]:
         else:
             return _register(func)
 
+    def cell_context(
+        self,
+        frame_offset: int = 2,
path: marimo/_ast/cell_manager.py
user: login: dmadisetti
body: I think passing the frame is better than incrementing the offset, so I will revise. I remembered the pattern in `mo.cache`, but think that was due to deferred and variable frame calls. 
created_at: 2025-03-07T19:37:46Z
author_association: COLLABORATOR
diff_hunk: @@ -149,6 +151,32 @@ def _mime_(self) -> tuple[KnownMimeType, str]:
         return tree(self._dict)._mime_()
 
 
+class _SetupContext(SkipContext):
+    """
+    A context manager that controls imports from being executed in top level code.
path: marimo/_ast/app.py
user: login: dmadisetti
body: This just skips the setup cell.

This could insert dummy variables such that the rest of the notebook doesn't break-
but more practically, I think for `marimo run` it will only load the definitions needed to prevent a name error, and leave everything else.
created_at: 2025-03-07T19:40:52Z
author_association: COLLABORATOR
diff_hunk: @@ -149,6 +151,32 @@ def _mime_(self) -> tuple[KnownMimeType, str]:
         return tree(self._dict)._mime_()
 
 
+class _SetupContext(SkipContext):
+    """
+    A context manager that controls imports from being executed in top level code.
path: marimo/_ast/app.py
user: login: akshayka
body: Got it, I was responding to the docstring, which mentions imports.

For `marimo run` — is there a reason to have different semantics than `marimo edit`?
created_at: 2025-03-07T20:00:53Z
author_association: CONTRIBUTOR
diff_hunk: @@ -377,6 +405,24 @@ def multiply(a: int, b: int) -> int:
             ),
         )
 
+    @property
path: marimo/_ast/app.py
user: login: mscolnick
body: i think its nice to have consistency with the @app.cell and can support both in future 
created_at: 2025-03-07T21:00:56Z
author_association: CONTRIBUTOR
diff_hunk: @@ -149,6 +151,32 @@ def _mime_(self) -> tuple[KnownMimeType, str]:
         return tree(self._dict)._mime_()
 
 
+class _SetupContext(SkipContext):
+    """
+    A context manager that controls imports from being executed in top level code.
path: marimo/_ast/app.py
user: login: dmadisetti
body: Maybe not

My thoughts were:

 -   edit should never break
 -   if something is wrong in the main definitions of my app, then on deploying my app should probably fail, but maybe fail a bit more gracefully
created_at: 2025-03-07T21:05:03Z
author_association: COLLABORATOR
diff_hunk: @@ -149,6 +151,32 @@ def _mime_(self) -> tuple[KnownMimeType, str]:
         return tree(self._dict)._mime_()
 
 
+class _SetupContext(SkipContext):
+    """
+    A context manager that controls imports from being executed in top level code.
path: marimo/_ast/app.py
user: login: akshayka
body: I see, that makes sense. I thought for maintainability it might just be simpler to have the same code path for both, but okay with either approach.
created_at: 2025-03-07T23:19:19Z
author_association: CONTRIBUTOR
diff_hunk: @@ -134,6 +138,20 @@ def decorator(func: Fn[P, R]) -> Cell | Fn[P, R]:
         else:
             return _register(func)
 
+    def cell_context(
+        self,
+        frame_offset: int = 2,
path: marimo/_ast/cell_manager.py
user: login: akshayka
body: @dmadisetti did you want to do this revision in this PR? 
created_at: 2025-03-09T15:58:00Z
author_association: CONTRIBUTOR
diff_hunk: @@ -134,6 +138,20 @@ def decorator(func: Fn[P, R]) -> Cell | Fn[P, R]:
         else:
             return _register(func)
 
+    def cell_context(
+        self,
+        frame_offset: int = 2,
path: marimo/_ast/cell_manager.py
user: login: dmadisetti
body: Yeah, a pretty small change. Just haven't been in front of a computer to address comments
created_at: 2025-03-09T18:18:52Z
author_association: COLLABORATOR
diff_hunk: @@ -1,5 +1,6 @@
 # Copyright 2024 Marimo. All rights reserved.
 DEFAULT_CELL_NAME = "_"
+SETUP_CELL_NAME = "setup"
path: marimo/_ast/names.py
user: login: dmadisetti
body: @akshayka let's continue this conversation in https://github.com/marimo-team/marimo/pull/4026 which Myles as already start putting work towards
created_at: 2025-03-10T15:21:17Z
author_association: COLLABORATOR
diff_hunk: @@ -149,6 +151,32 @@ def _mime_(self) -> tuple[KnownMimeType, str]:
         return tree(self._dict)._mime_()
 
 
+class _SetupContext(SkipContext):
+    """
+    A context manager that controls imports from being executed in top level code.
path: marimo/_ast/app.py
user: login: dmadisetti
body: I filled this in a bit more. Allowing for a broken setup block, but triggering a warning
created_at: 2025-03-10T15:23:25Z
author_association: COLLABORATOR

PR comments:
user: login: vercel[bot]
created_at: 2025-03-07T13:56:33Z
body: [vc]: #C8i34bfjvNxyGniwY4jZJ9c+L0/EDGQcZyb5wM+LYTU=:eyJpc01vbm9yZXBvIjp0cnVlLCJ0eXBlIjoiZ2l0aHViIiwicHJvamVjdHMiOlt7Im5hbWUiOiJtYXJpbW8tZG9jcyIsInJvb3REaXJlY3RvcnkiOiJkb2NzIiwibGl2ZUZlZWRiYWNrIjp7InJlc29sdmVkIjowLCJ1bnJlc29sdmVkIjowLCJ0b3RhbCI6MCwibGluayI6Im1hcmltby1kb2NzLWdpdC1kbS1zZXR1cC1jb250ZXh0LW1hcmltby52ZXJjZWwuYXBwIn0sImluc3BlY3RvclVybCI6Imh0dHBzOi8vdmVyY2VsLmNvbS9tYXJpbW8vbWFyaW1vLWRvY3MvQXB2ZW5WZHBETktYOXZmZERNTGRIUm9VN05tWCIsInByZXZpZXdVcmwiOiJtYXJpbW8tZG9jcy1naXQtZG0tc2V0dXAtY29udGV4dC1tYXJpbW8udmVyY2VsLmFwcCIsIm5leHRDb21taXRTdGF0dXMiOiJERVBMT1lFRCJ9LHsibmFtZSI6Im1hcmltby1zdG9yeWJvb2siLCJyb290RGlyZWN0b3J5IjoiZnJvbnRlbmQiLCJsaXZlRmVlZGJhY2siOnsicmVzb2x2ZWQiOjAsInVucmVzb2x2ZWQiOjAsInRvdGFsIjowLCJsaW5rIjoibWFyaW1vLXN0b3J5Ym9vay1naXQtZG0tc2V0dXAtY29udGV4dC1tYXJpbW8udmVyY2VsLmFwcCJ9LCJpbnNwZWN0b3JVcmwiOiJodHRwczovL3ZlcmNlbC5jb20vbWFyaW1vL21hcmltby1zdG9yeWJvb2svNURzcmZ6VHZMMXBEeG1KOHlMcWl0NkpZbTRFbiIsInByZXZpZXdVcmwiOiJtYXJpbW8tc3Rvcnlib29rLWdpdC1kbS1zZXR1cC1jb250ZXh0LW1hcmltby52ZXJjZWwuYXBwIiwibmV4dENvbW1pdFN0YXR1cyI6IkRFUExPWUVEIn1dfQ==
**The latest updates on your projects**. Learn more about [Vercel for Git ↗︎](https://vercel.link/github-learn-more)

| Name | Status | Preview | Comments | Updated (UTC) |
| :--- | :----- | :------ | :------- | :------ |
| **marimo-docs** | ✅ Ready ([Inspect](https://vercel.com/marimo/marimo-docs/ApvenVdpDNKX9vfdDMLdHRoU7NmX)) | [Visit Preview](https://marimo-docs-git-dm-setup-context-marimo.vercel.app) | 💬 [**Add feedback**](https://vercel.live/open-feedback/marimo-docs-git-dm-setup-context-marimo.vercel.app?via=pr-comment-feedback-link) | Mar 10, 2025 3:22pm |
| **marimo-storybook** | ✅ Ready ([Inspect](https://vercel.com/marimo/marimo-storybook/5DsrfzTvL1pDxmJ8yLqit6JYm4En)) | [Visit Preview](https://marimo-storybook-git-dm-setup-context-marimo.vercel.app) | 💬 [**Add feedback**](https://vercel.live/open-feedback/marimo-storybook-git-dm-setup-context-marimo.vercel.app?via=pr-comment-feedback-link) | Mar 10, 2025 3:22pm |



author_association: NONE
user: login: dmadisetti
created_at: 2025-03-07T17:36:07Z
body: Thanks! @mscolnick just tagged you in a messy frontend branch that might be a good starting point
author_association: COLLABORATOR
user: login: dmadisetti
created_at: 2025-03-13T01:55:42Z
body: @akshayka were you alright with this as a stop gap?
author_association: COLLABORATOR
user: login: github-actions[bot]
created_at: 2025-03-13T02:04:29Z
body: 🚀 Development release published. You may be able to view the changes at https://marimo.app?v=0.11.20-dev3
author_association: NONE