feat: marimo edit --watch (#3437)

Closes #3114
Closes #2675
Closes #1511

This adds support for `marimo edit --watch`. You can do

```python
marimo edit --watch # will watch notebooks as they are opened
# or
marimo edit nb.py --watch # will watch just this notebook
```

This doesn't support a few edge case:
- handle renaming files
- handle naming files (for first time)

docs/guides/editor_features/watching.md

@@ -0,0 +1,29 @@
+# Watching notebooks
+
+marimo's `--watch` flag enables a file watcher that automatically sync your
+notebook file with the marimo editor or running application.
+
+This allows you to edit your notebook file in the editor of your choice, and
+have the changes automatically reflected in the running editor or application.
+
+!!! tip "Install watchdog for better file watching"
+    For better file watching performance, install watchdog with `pip install watchdog`. Without watchdog, marimo will poll for file changes which is less efficient.
+
+## `marimo run --watch`
+
+When you run a notebook with the `--watch` flag, whenever the file watcher
+detects a change to the notebook file, the application will be refreshed.
+The browser will trigger a page refresh to ensure your notebook starts from a fresh state.
+
+## `marimo watch --watch`
+
+When you edit a notebook file with the `--watch` flag, whenever the file watcher
+detects a change to the notebook file, the new cells and code changes will be streamed to
+the browser editor.
+
+This code will not be executed until you run the cell, and instead marked as stale.
+
+## Watching for data changes
+
+!!! note
+    Support for watching data files and automatically refreshing cells that depend on them is coming soon. Follow along at <https://github.com/marimo-team/marimo/issues/3258>

frontend/src/core/cells/__tests__/cells.test.ts

@@ -1305,9 +1305,29 @@ describe("cell reducer", () => {
     actions.setCellIds({ cellIds: newIds });
     expect(state.cellIds.atOrThrow(FIRST_COLUMN).topLevelIds).toEqual(newIds);
 
-    actions.setCellCodes({ codes: newCodes, ids: newIds });
+    // When codeIsStale is false, lastCodeRun should match code
+    actions.setCellCodes({
+      codes: newCodes,
+      ids: newIds,
+      codeIsStale: false,
+    });
     newIds.forEach((id, index) => {
       expect(state.cellData[id].code).toBe(newCodes[index]);
+      expect(state.cellData[id].lastCodeRun).toBe(newCodes[index]);
+      expect(state.cellData[id].edited).toBe(false);
+    });
+
+    // When codeIsStale is true, lastCodeRun should not change
+    const staleCodes = ["stale1", "stale2", "stale3"];
+    actions.setCellCodes({
+      codes: staleCodes,
+      ids: newIds,
+      codeIsStale: true,
+    });
+    newIds.forEach((id, index) => {
+      expect(state.cellData[id].code).toBe(staleCodes[index]);
+      expect(state.cellData[id].lastCodeRun).toBe(newCodes[index]);
+      expect(state.cellData[id].edited).toBe(true);
     });
   });
 

frontend/src/core/cells/cells.ts

@@ -780,7 +780,10 @@ const {
       cellHandles: nextCellHandles,
     };
   },
-  setCellCodes: (state, action: { codes: string[]; ids: CellId[] }) => {
+  setCellCodes: (
+    state,
+    action: { codes: string[]; ids: CellId[]; codeIsStale: boolean },
+  ) => {
     invariant(
       action.codes.length === action.ids.length,
       "Expected codes and ids to have the same length",
@@ -791,11 +794,26 @@ const {
       const code = action.codes[i];
 
       state = updateCellData(state, cellId, (cell) => {
+        // No change
+        if (cell.code.trim() === code.trim()) {
+          return cell;
+        }
+
+        // Update codemirror if mounted
+        const cellHandle = state.cellHandles[cellId].current;
+        if (cellHandle?.editorView) {
+          updateEditorCodeFromPython(cellHandle.editorView, code);
+        }
+
+        // If code is stale, we don't promote it to lastCodeRun
+        const lastCodeRun = action.codeIsStale ? cell.lastCodeRun : code;
+
         return {
           ...cell,
-          code,
-          edited: false,
-          lastCodeRun: code,
+          code: code,
+          // Mark as edited if the code has changed
+          edited: lastCodeRun ? lastCodeRun.trim() !== code.trim() : false,
+          lastCodeRun,
         };
       });
     }

frontend/src/core/websocket/useMarimoWebSocket.tsx

@@ -203,6 +203,7 @@ export function useMarimoWebSocket(opts: {
         setCellCodes({
           codes: msg.data.codes,
           ids: msg.data.cell_ids as CellId[],
+          codeIsStale: msg.data.code_is_stale,
         });
         return;
       case "update-cell-ids":

marimo/_cli/cli.py

@@ -286,6 +286,14 @@ def main(
     help=sandbox_message,
 )
 @click.option("--profile-dir", default=None, type=str, hidden=True)
+@click.option(
+    "--watch",
+    is_flag=True,
+    default=False,
+    show_default=True,
+    type=bool,
+    help="Watch the file for changes and reload the code when saved in another editor.",
+)
 @click.argument("name", required=False, type=click.Path())
 @click.argument("args", nargs=-1, type=click.UNPROCESSED)
 def edit(
@@ -300,6 +308,7 @@ def edit(
     skip_update_check: bool,
     sandbox: bool,
     profile_dir: Optional[str],
+    watch: bool,
     name: Optional[str],
     args: tuple[str, ...],
 ) -> None:
@@ -369,7 +378,7 @@ def edit(
         headless=headless,
         mode=SessionMode.EDIT,
         include_code=True,
-        watch=False,
+        watch=watch,
         cli_args=parse_args(args),
         auth_token=_resolve_token(token, token_password),
         base_url=base_url,

marimo/_messaging/ops.py

@@ -639,6 +639,7 @@ class UpdateCellCodes(Op):
     name: ClassVar[str] = "update-cell-codes"
     cell_ids: List[CellId_t]
     codes: List[str]
+    code_is_stale: bool
 
 
 @dataclass

marimo/_server/api/endpoints/files.py

@@ -104,6 +104,12 @@ async def rename_file(
         from_consumer_id=ConsumerId(app_state.require_current_session_id()),
     )
 
+    if new_path:
+        # Handle rename for watch
+        app_state.session_manager.handle_file_rename_for_watch(
+            app_state.require_current_session_id(), new_path
+        )
+
     return SuccessResponse()
 
 

marimo/_server/api/lifespans.py

@@ -75,15 +75,6 @@ async def lsp(app: Starlette) -> AsyncIterator[None]:
     yield
 
 
-@contextlib.asynccontextmanager
-async def watcher(app: Starlette) -> AsyncIterator[None]:
-    state = AppState.from_app(app)
-    if state.watch:
-        session_mgr = state.session_manager
-        session_mgr.start_file_watcher()
-    yield
-
-
 @contextlib.asynccontextmanager
 async def open_browser(app: Starlette) -> AsyncIterator[None]:
     state = AppState.from_app(app)