Repository: marimo-team/marimo
Pull request: feat: mo.thread lifecycle
Closes #4564 

This change adds a `should_exit` property to `mo.Thread` which evaluates to `True`
when the thread's spawning cell lifecycle has ended. This makes it
possible for developers to clean up their threads.

This change also adds a utility method, mo.current_thread(), that
returns the current thread of control if it is a marimo.Thread object,
and raises otherwise. Developers could technically use
threading.current_thread() instead, but this method provides
type safety and validation.

Problem statement:
mo.Threads are not cleaned up
### Describe the bug

On changing a mo.Thread target definition, an noticed that the background thread was still running (I just changed print from 1 to 2 of the smoke test and saw both print statements)

Likewise the plain threading.Thread definition seemed to be running on cell deletions. It might be worth having thread specific deletion handlers if we're going to support mo.Thread

### Environment

<details>

```
Replace this line with the output of marimo env. Leave the backticks in place.
```

</details>


### Code to reproduce

```python


import marimo

__generated_with = "0.12.10"
app = marimo.App()


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


@app.cell
def _(sleep):
    def poll():
        while True:
            sleep(1)
            print(2)
    return (poll,)


@app.cell
def _():
    import threading
    return


@app.cell
def _(mo, poll):
    with mo.redirect_stdout():
        t = mo.Thread(target=poll)
        t.start()
    return (t,)


@app.cell
def _(t):
    t.is_alive()
    return


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

Hints from issue discussion:
Instead of forcibly terminating threads, which is difficult if not impossible to do reliably across operating systems, we could introduce some kind of condition variable/event-based API, with a guarantee that the runtime will notify/set the event when the thread *should* be cleaned up. But it would be up to the user to make sure they actually check the condition periodically.

Commit message: unit tests

Patch to review:
diff --git a/docs/api/control_flow.md b/docs/api/control_flow.md
index cfec9fdb649..5110639e23f 100644
--- a/docs/api/control_flow.md
+++ b/docs/api/control_flow.md
@@ -1,9 +1,19 @@
 # Control flow
 
-Use `mo.stop` to halt execution of a cell, and optionally output an object.
-This function is useful for validating user input.
+## Controlling when cells run
+
+* Use [`mo.stop`][marimo.stop] to halt execution of a cell when a condition is met.
+* Combine [`mo.stop`][marimo.stop] with [`mo.ui.run_button`][marimo.ui.run_button] to gate execution on button click.
+* Use [`mo.ui.refresh`][marimo.ui.refresh] to make cells run periodically.
+
+!!! tip "Lazy execution"
+
+    In addition to these utilities, you can [configure the runtime to be lazy](../guides/expensive_notebooks.md#configure-how-marimo-runs-cells), marking cells as stale instead of automatically running them.
 
 ::: marimo.stop
 
-Use [`mo.ui.refresh`][marimo.ui.refresh] to trigger other cells to run periodically, on a configurable
-interval (or on click).
+## Threading
+
+::: marimo.Thread
+
+::: marimo.current_thread
diff --git a/docs/api/miscellaneous.md b/docs/api/miscellaneous.md
index 862733f112c..f5687abe832 100644
--- a/docs/api/miscellaneous.md
+++ b/docs/api/miscellaneous.md
@@ -9,5 +9,3 @@
 ::: marimo.notebook_dir
 
 ::: marimo.notebook_location
-
-::: marimo.Thread
diff --git a/docs/guides/outputs.md b/docs/guides/outputs.md
index 430764b21bb..be3c3fdbbb5 100644
--- a/docs/guides/outputs.md
+++ b/docs/guides/outputs.md
@@ -196,6 +196,28 @@ To create a thread that can reliably communicate outputs to the frontend,
 use [`mo.Thread`][marimo.Thread], which has exactly the same API as
 as `threading.Thread`.
 
+### Cleaning up your thread
+
+When the cell that spawned a [`mo.Thread`][marimo.Thread] is invalidated
+(re-run, deleted, interrupted, or otherwise errored), the thread's
+`should_exit` property will evaluate to `True`, at which point it is your
+responsibility to clean up your thread. You can retrieve the current
+[`mo.Thread`][marimo.Thread] with [`mo.current_thread`][marimo.current_thread].
+
+**Example.**
+
+```python
+def target():
+    import marimo as mo
+
+    thread = mo.current_thread()
+    while not thread.should_exit:
+        ...
+```
+
+
+### Patching threads created by third-party code
+
 If you need to forward outputs from threads spawned by third-party code, try
 patching `threading.Thread`:
 
@@ -205,3 +227,6 @@ import marimo as mo
 
 threading.Thread = mo.Thread
 ```
+
+This however may leak threads, since the patched threads won't know to check the `mo.Thread`'s
+`should_exit` property.
diff --git a/marimo/__init__.py b/marimo/__init__.py
index 0fbe43900b0..b146df5baef 100644
--- a/marimo/__init__.py
+++ b/marimo/__init__.py
@@ -22,6 +22,7 @@
     "MarimoIslandGenerator",
     "MarimoStopError",
     "Thread",
+    "current_thread",
     # Other namespaces
     "ai",
     "ui",
@@ -137,7 +138,7 @@
     refs,
 )
 from marimo._runtime.state import state
-from marimo._runtime.threads import Thread
+from marimo._runtime.threads import Thread, current_thread
 from marimo._save.save import cache, lru_cache, persistent_cache
 from marimo._server.asgi import create_asgi_app
 from marimo._sql.sql import sql
diff --git a/marimo/_runtime/threads.py b/marimo/_runtime/threads.py
index f57adda8a61..2f567f4cb34 100644
--- a/marimo/_runtime/threads.py
+++ b/marimo/_runtime/threads.py
@@ -1,10 +1,12 @@
 # Copyright 2024 Marimo. All rights reserved.
 from __future__ import annotations
 
+import functools
 import threading
 from typing import Any
 
 from marimo._messaging.streams import ThreadSafeStream
+from marimo._runtime.cell_lifecycle_item import CellLifecycleItem
 from marimo._runtime.context.kernel_context import KernelRuntimeContext
 from marimo._runtime.context.script_context import ScriptRuntimeContext
 from marimo._runtime.context.types import (
@@ -21,7 +23,7 @@
 
 
 class Thread(threading.Thread):
-    """A Thread subclass that is aware of marimo internals.
+    """A Thread subclass that can communicate with the frontend.
 
     `mo.Thread` has the same API as threading.Thread,
     but `mo.Thread`s are able to communicate with the marimo
@@ -33,17 +35,56 @@ class Thread(threading.Thread):
 
     Writing directly to sys.stdout or sys.stderr, or to file descriptors 1 and
     2, is not yet supported.
+
+    **Thread lifecycle.** When the cell that spawned this thread is invalidated
+    (re-run, deleted, interrupted, or otherwise errored), this thread's
+    `should_exit` property will evaluate to `True`, at which point it
+    is the developer's responsibility to clean up their thread.
+
+    **Example.**
+
+    ```python
+    def target():
+        import time
+        import marimo as mo
+
+        thread = mo.current_thread()
+        while not thread.should_exit:
+            time.sleep(1)
+            print("hello")
+    ```
+
+    ```python
+    import marimo as mo
+
+    mo.Thread(target=target).start()
+    ```
     """
 
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         super().__init__(*args, **kwargs)
 
         self._marimo_ctx: RuntimeContext | None = None
+        exit_event = threading.Event()
+        self._exit_event = exit_event
 
         if not runtime_context_installed():
             return
 
+        class ThreadLifecycle(CellLifecycleItem):
+            def create(self, context: RuntimeContext) -> None:
+                del context
+                pass
+
+            def dispose(self, context: RuntimeContext, deletion: bool) -> bool:
+                del context
+                del deletion
+                exit_event.set()
+                return True
+
         ctx = get_context()
+        ctx.cell_lifecycle_registry.add(ThreadLifecycle())
+
         if isinstance(ctx, KernelRuntimeContext):
             self._marimo_ctx = KernelRuntimeContext(**ctx.__dict__)
             # standard IO is not yet threadsafe
@@ -77,6 +118,30 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
                     "Unsupported stream type " + str(type(ctx.stream))
                 )
 
+    @property
+    def should_exit(self) -> bool:
+        """Whether this thread should exit.
+
+        Returns `True` when the cell that spawned this thread has been invalidated;
+        for example, if the cell:
+
+        - was re-run
+        - was deleted
+        - was interrupted
+
+        then this property evaluates to True.
+
+        It is the developer's responsibility to clean up and finish their
+        thread when this flag is set. Retrieve the current `mo.Thread` with
+
+        ```python
+        import marimo as mo
+
+        mo.current_thread()
+        ```
+        """
+        return self._exit_event.is_set()
+
     def run(self) -> None:
         if self._marimo_ctx is not None:
             initialize_context(self._marimo_ctx)
@@ -90,3 +155,21 @@ def run(self) -> None:
         super().run()
         THREADS.remove(thread_id)
         teardown_context()
+
+
+def current_thread() -> Thread:
+    """Return the `marimo.Thread` object for the caller's thread of control.
+
+    Returns:
+        The current `marimo.Thread` object.
+
+    Raises:
+        RuntimeError: If the current thread of control is not a `marimo.Thread`.
+    """
+    thread = threading.current_thread()
+    if not isinstance(thread, Thread):
+        raise RuntimeError(
+            "mo.current_thread() must be called from a "
+            "thread created with mo.Thread."
+        )
+    return thread
diff --git a/tests/_runtime/test_threads.py b/tests/_runtime/test_threads.py
index 4bea4dc5ec1..7df085ef683 100644
--- a/tests/_runtime/test_threads.py
+++ b/tests/_runtime/test_threads.py
@@ -1,5 +1,6 @@
 import time
 
+from marimo._runtime.requests import DeleteCellRequest
 from marimo._runtime.runtime import Kernel
 from tests.conftest import ExecReqProvider
 
@@ -135,3 +136,101 @@ def target():
     assert thread_stream.messages[0][1]["console"]["channel"] == "stdout"
     assert "world" in thread_stream.messages[1][1]["console"]["data"]
     assert thread_stream.messages[1][1]["console"]["channel"] == "stdout"
+
+
+async def test_thread_should_exit_on_rerun(
+    k: Kernel, exec_req: ExecReqProvider
+) -> None:
+    """Test that a thread's exit event is set when cell lifecycle is disposed."""
+
+    await k.run(
+        [
+            exec_req.get("import marimo as mo"),
+            exec_req.get(
+                """
+                def target():
+                    ...
+                """
+            ),
+            er := exec_req.get(
+                """
+                thread = mo.Thread(target=target)
+                thread.start()
+                """
+            ),
+        ]
+    )
+
+    thread = k.globals["thread"]
+    assert not thread.should_exit
+
+    # rerunning the cell should trigger the cell lifecycle disposal and set its exit event
+    await k.run([er])
+    assert thread.should_exit
+
+
+async def test_thread_should_not_exit_on_other_cell_run(
+    k: Kernel, exec_req: ExecReqProvider
+) -> None:
+    """Test that a thread's exit event is set when cell lifecycle is disposed."""
+
+    await k.run(
+        [
+            exec_req.get("import marimo as mo"),
+            exec_req.get(
+                """
+                def target():
+                    ...
+                """
+            ),
+            exec_req.get(
+                """
+                thread = mo.Thread(target=target)
+                thread.start()
+                """
+            ),
+            er := exec_req.get(
+                """
+                ...
+                """
+            ),
+        ]
+    )
+
+    thread = k.globals["thread"]
+    assert not thread.should_exit
+
+    # rerunning a cell that is not related to the spawning cell should _not_
+    # signal the thread to exit
+    await k.run([er])
+    assert not thread.should_exit
+
+
+async def test_thread_should_exit_on_deletion(
+    k: Kernel, exec_req: ExecReqProvider
+) -> None:
+    """Test that a thread's exit event is set when cell lifecycle is disposed."""
+
+    await k.run(
+        [
+            exec_req.get("import marimo as mo"),
+            exec_req.get(
+                """
+                def target():
+                    ...
+                """
+            ),
+            er := exec_req.get(
+                """
+                thread = mo.Thread(target=target)
+                thread.start()
+                """
+            ),
+        ]
+    )
+
+    thread = k.globals["thread"]
+    assert not thread.should_exit
+
+    await k.delete_cell(DeleteCellRequest(er.cell_id))
+    assert thread.should_exit


Code review comments:
diff_hunk: @@ -33,17 +35,56 @@ class Thread(threading.Thread):
 
     Writing directly to sys.stdout or sys.stderr, or to file descriptors 1 and
     2, is not yet supported.
+
+    **Thread lifecycle.** When the cell that spawned this thread is invalidated
+    (re-run, deleted, interrupted, or otherwise errored), this thread's
+    `should_exit` property will evaluate to `True`, at which point it
+    is the developer's responsibility to clean up their thread.
+
+    **Example.**
+
+    ```python
+    def target():
+        import time
+        import marimo as mo
+
+        thread = mo.current_thread()
+        while not thread.should_exit:
+            time.sleep(1)
+            print("hello")
+    ```
+
+    ```python
+    import marimo as mo
+
+    mo.Thread(target=target).start()
+    ```
     """
 
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         super().__init__(*args, **kwargs)
 
         self._marimo_ctx: RuntimeContext | None = None
+        exit_event = threading.Event()
+        self._exit_event = exit_event
 
         if not runtime_context_installed():
             return
 
+        class ThreadLifecycle(CellLifecycleItem):
+            def create(self, context: RuntimeContext) -> None:
+                del context
+                pass
path: marimo/_runtime/threads.py
text: @user1:
```suggestion
```