Repository: marimo-team/marimo

Issue:
Embedded App Errors on Browser Refresh
### Describe the bug

I've been working on deployment of an app that includes several subapps written in their own marimo notebooks and integrated in the main notebook using `app.embed()`. When running the main notebook we've noticed multiple issue stemming from the embedded apps when a user refreshes the browser.

- Components based off input values seem to "stick" from the inputs made before refreshing the browser
- Input entered/selected after the browser refresh do not take
- Run buttons do not execute the cells associated with them after the browser refresh

First encountered this when trying to run the app with FastAPI then confirmed the behavior happens when running `marimo run main.py` but not `marimo edit main.py` which is why we hadn't originally caught it in development. 

I've created a simplified example to illustrate the issues. We've seen this issue on Windows 11, Ubuntu 20.04.06 and in Azure App Services "Debian GNU/Linux 11 (bullseye)".

Thanks in advance, we've really been enjoying using marimo!

### Environment

<details>

```
{
  "marimo": "0.10.19",
  "OS": "Windows",
  "OS Version": "11",
  "Processor": "Intel64 Family 6 Model 183 Stepping 1, GenuineIntel",
  "Python Version": "3.11.9",
  "Binaries": {
    "Browser": "132.0.6834.160",
    "Node": "--"
  },
  "Dependencies": {
    "click": "8.1.8",
    "docutils": "0.21.2",
    "itsdangerous": "2.2.0",
    "jedi": "0.19.2",
    "markdown": "3.7",
    "narwhals": "1.23.0",
    "packaging": "24.2",
    "psutil": "6.1.1",
    "pygments": "2.19.1",
    "pymdown-extensions": "10.14.1",
    "pyyaml": "6.0.2",
    "ruff": "0.9.2",
    "starlette": "0.41.3",
    "tomlkit": "0.13.2",
    "typing-extensions": "4.12.2",
    "uvicorn": "0.34.0",
    "websockets": "14.2"
  },
  "Optional Dependencies": {
    "pandas": "1.5.3"
  },
  "Experimental Flags": {}
}
```

</details>


### Code to reproduce

# subapp.py code
```
import marimo

__generated_with = "0.10.19"
app = marimo.App(width="medium")


@app.cell
def _():
    import marimo as mo
    return (mo,)


@app.cell
def _(mo):
    mo.md("## This is the subapp")
    return


@app.cell
def _(mo):
    category_list = ['a','b','c','d','e','f','g','h']
    category_nums = mo.ui.dictionary(
            {
                alpha: mo.ui.number(1, 1000, 1, full_width=True)
                for alpha in category_list
            }
    )
    category_select = mo.ui.multiselect(category_list,label='Categories')
    mo.vstack([
        mo.md("Select categories and enter number of each."),
        category_select
    ])
    return category_list, category_nums, category_select


@app.cell
def _(category_nums, category_select, mo):
    mo.hstack(
        [
            mo.hstack([mo.md(f"<i>{alpha} - </i>"), category_nums[alpha].style(max_width="80px")])
            for alpha in category_select.value
        ],
        justify="start",
    )
    return


@app.cell
def _(mo):
    submit_button = mo.ui.run_button()
    submit_button
    return (submit_button,)


@app.cell
def _(category_nums, category_select, mo, submit_button):
    submit_display = ""
    if submit_button.value:
        submit_display = mo.vstack(
            [
                mo.md(f"You entered {category_nums[alpha].value} of category {alpha}")
                for alpha in category_select.value
            ]
        )
    submit_display
    return (submit_display,)


@app.cell
def _():
    return


if __name__ == "__main__":
    app.run()

```
# main.py code
```
import marimo

__generated_with = "0.10.19"
app = marimo.App(width="medium")


@app.cell
def _():
    import marimo as mo
    from subapp import app
    return app, mo


@app.cell
def _(mo):
    mo.md("""# This is the main app""")
    return


@app.cell
def _(mo):
    mo.md("""Below is where the sub app is displayed""")
    return


@app.cell
async def _(app):
    sub = await app.embed()
    sub.output
    return (sub,)


@app.cell
def _():
    return


if __name__ == "__main__":
    app.run()

```

Discussion:

Solution patch:
diff --git a/marimo/_ast/app.py b/marimo/_ast/app.py
index f5245978..d28f4e32 100644
--- a/marimo/_ast/app.py
+++ b/marimo/_ast/app.py
@@ -4,6 +4,7 @@ from __future__ import annotations
 import ast
 import base64
 import inspect
+import threading
 from dataclasses import asdict, dataclass, field
 from pathlib import Path
 from textwrap import dedent
@@ -37,6 +38,7 @@ from marimo._runtime import dataflow
 from marimo._runtime.app.kernel_runner import AppKernelRunner
 from marimo._runtime.app.script_runner import AppScriptRunner
 from marimo._runtime.context.types import (
+    ContextNotInitializedError,
     get_context,
     runtime_context_installed,
 )
@@ -132,6 +134,31 @@ class AppEmbedResult:
     defs: Mapping[str, object]
 
 
+class AppKernelRunnerRegistry:
+    def __init__(self) -> None:
+        # Mapping from thread to its kernel runners, so that app.embed() calls are
+        # isolated across run sessions.
+        self._runners: dict[int, dict[App, AppKernelRunner]] = {}
+
+    def get_runner(self, app: App) -> AppKernelRunner:
+        app_runners = self._runners.setdefault(threading.get_ident(), {})
+        runner = app_runners.get(app, None)
+        if runner is None:
+            runner = AppKernelRunner(InternalApp(app))
+            app_runners[app] = runner
+        return runner
+
+    def remove_runner(self, app: App) -> None:
+        app_runners = self._runners.get(tid := threading.get_ident(), {})
+        if app in app_runners:
+            del app_runners[app]
+        if tid in self._runners and not (self._runners[tid]):
+            del self._runners[tid]
+
+    def shutdown(self) -> None:
+        self._runners.clear()
+
+
 @mddoc
 class App:
     """A marimo notebook.
@@ -175,7 +202,12 @@ class App:
             self._filename = inspect.getfile(inspect.stack()[1].frame)
         except Exception:
             ...
-        self._app_kernel_runner: AppKernelRunner | None = None
+
+    def __del__(self) -> None:
+        try:
+            get_context().app_kernel_runner_registry.remove_runner(self)
+        except ContextNotInitializedError:
+            ...
 
     def clone(self) -> App:
         """Clone an app to embed multiple copies of it.
@@ -328,9 +360,7 @@ class App:
             self._initialized = True
 
     def _get_kernel_runner(self) -> AppKernelRunner:
-        if self._app_kernel_runner is None:
-            self._app_kernel_runner = AppKernelRunner(InternalApp(self))
-        return self._app_kernel_runner
+        return get_context().app_kernel_runner_registry.get_runner(self)
 
     def _flatten_outputs(self, outputs: dict[CellId_t, Any]) -> Sequence[Any]:
         return tuple(
diff --git a/marimo/_runtime/app/kernel_runner.py b/marimo/_runtime/app/kernel_runner.py
index 0e05d89d..e840ccff 100644
--- a/marimo/_runtime/app/kernel_runner.py
+++ b/marimo/_runtime/app/kernel_runner.py
@@ -92,10 +92,9 @@ class AppKernelRunner:
             parent=ctx,
         )
         ctx.add_child(self._runtime_context)
-        finalizer = weakref.finalize(
+        weakref.finalize(
             self, ctx.remove_child, self._runtime_context
-        )
-        finalizer.atexit = False
+        ).atexit = False
 
         # Register cells through the kernel runner, so that compilation only
         # occurs once.
diff --git a/marimo/_runtime/context/kernel_context.py b/marimo/_runtime/context/kernel_context.py
index e291c144..428428b5 100644
--- a/marimo/_runtime/context/kernel_context.py
+++ b/marimo/_runtime/context/kernel_context.py
@@ -5,6 +5,7 @@ from contextlib import contextmanager
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Iterator, Optional
 
+from marimo._ast.app import AppKernelRunnerRegistry
 from marimo._config.config import MarimoConfig
 from marimo._messaging.types import Stderr, Stdout
 from marimo._plugins.ui._core.ids import IDProvider, NoIDProviderException
@@ -152,6 +153,7 @@ def create_kernel_context(
         state_registry=StateRegistry(),
         function_registry=FunctionRegistry(),
         cell_lifecycle_registry=CellLifecycleRegistry(),
+        app_kernel_runner_registry=AppKernelRunnerRegistry(),
         virtual_file_registry=VirtualFileRegistry(),
         virtual_files_supported=virtual_files_supported,
         stream=stream,
diff --git a/marimo/_runtime/context/script_context.py b/marimo/_runtime/context/script_context.py
index 521936ef..57946523 100644
--- a/marimo/_runtime/context/script_context.py
+++ b/marimo/_runtime/context/script_context.py
@@ -6,6 +6,7 @@ from contextlib import contextmanager
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Iterator, Optional
 
+from marimo._ast.app import AppKernelRunnerRegistry
 from marimo._cli.parse_args import args_from_argv
 from marimo._config.config import MarimoConfig
 from marimo._config.manager import get_default_config_manager
@@ -140,6 +141,7 @@ def initialize_script_context(
         state_registry=StateRegistry(),
         function_registry=FunctionRegistry(),
         cell_lifecycle_registry=CellLifecycleRegistry(),
+        app_kernel_runner_registry=AppKernelRunnerRegistry(),
         virtual_file_registry=VirtualFileRegistry(),
         virtual_files_supported=False,
         stream=stream,
diff --git a/marimo/_runtime/context/types.py b/marimo/_runtime/context/types.py
index 217da534..8fd17277 100644
--- a/marimo/_runtime/context/types.py
+++ b/marimo/_runtime/context/types.py
@@ -21,7 +21,7 @@ from marimo._runtime.functions import FunctionRegistry
 from marimo._runtime.requests import HTTPRequest
 
 if TYPE_CHECKING:
-    from marimo._ast.app import InternalApp
+    from marimo._ast.app import AppKernelRunnerRegistry, InternalApp
     from marimo._ast.cell import CellId_t
     from marimo._messaging.types import Stream
     from marimo._output.hypertext import Html
@@ -71,6 +71,7 @@ class RuntimeContext(abc.ABC):
     cell_lifecycle_registry: CellLifecycleRegistry
     virtual_file_registry: VirtualFileRegistry
     virtual_files_supported: bool
+    app_kernel_runner_registry: AppKernelRunnerRegistry
     # stream, stdout, stderr are _not_ owned by the context
     stream: Stream
     stdout: Stdout | None
@@ -154,6 +155,7 @@ class RuntimeContext(abc.ABC):
 
     def remove_child(self, runtime_context: RuntimeContext) -> None:
         self.children.remove(runtime_context)
+        runtime_context.virtual_file_registry.shutdown()
         assert runtime_context not in self.children
 
     @contextmanager
diff --git a/marimo/_runtime/runtime.py b/marimo/_runtime/runtime.py
index d0e677f1..6b1668e9 100644
--- a/marimo/_runtime/runtime.py
+++ b/marimo/_runtime/runtime.py
@@ -2310,7 +2310,14 @@ def launch_kernel(
         profiler.disable()
         profiler.dump_stats(profile_path)
 
+    # Defensively clear context data structures, in case a leak prevents
+    # the context from being destroyed.
+    #
+    # TODO(akshayka): define ownership semantics for contexts, so the
+    # context knows how to shut itself down. The virtual file registry
+    # is shared between the main thread and mo.Thread's right now ...
     get_context().virtual_file_registry.shutdown()
+    get_context().app_kernel_runner_registry.shutdown()
     teardown_context()
     kernel.teardown()
     if isinstance(pipe, connection.Connection):


Test patch:
diff --git a/tests/_ast/test_app.py b/tests/_ast/test_app.py
index e59de6ef..bf5add3c 100644
--- a/tests/_ast/test_app.py
+++ b/tests/_ast/test_app.py
@@ -1,14 +1,14 @@
 # Copyright 2024 Marimo. All rights reserved.
 
 from __future__ import annotations
-
 import pathlib
 import subprocess
 import textwrap
-from typing import Any
+from typing import Any, TYPE_CHECKING
 
 import pytest
-from marimo._ast.app import App, _AppConfig, InternalApp
+
+from marimo._ast.app import App, AppKernelRunnerRegistry, _AppConfig, InternalApp
 from marimo._ast.errors import (
     CycleError,
     DeleteNonlocalError,
@@ -18,9 +18,11 @@ from marimo._ast.errors import (
 from marimo._dependencies.dependencies import DependencyManager
 from marimo._plugins.stateless.flex import vstack
 from marimo._runtime.requests import SetUIElementValueRequest
-from marimo._runtime.runtime import Kernel
 from tests.conftest import ExecReqProvider
 
+if TYPE_CHECKING:
+    from marimo._runtime.runtime import Kernel
+
 
 # don't complain for useless expressions (cell outputs)
 # ruff: noqa: B018
@@ -734,3 +736,20 @@ class TestAppComposition:
 
         assert x.value == 2
         assert y.value == 3
+
+
+class TestAppKernelRunnerRegistry:
+    def test_get_runner(self, k: Kernel) -> None:
+        # `k` fixture installs a context, needed for AppKernelRunner
+        del k
+        app = App()
+        registry = AppKernelRunnerRegistry()
+        # Calling with the same app yields the same runner
+        assert registry.get_runner(app) == registry.get_runner(app)
+
+        # Calling with different app objects yields different runners
+        assert registry.get_runner(app) != registry.get_runner(other := App())
+
+        registry.remove_runner(app)
+        registry.remove_runner(other)
+        assert not registry._runners