Add more tests for CellStateManager.ts

extension/src/__mocks__/TestVsCode.ts

@@ -117,7 +117,7 @@ class NotebookCellOutputItem implements NotebookCellOutputItem {
   }
 }
 
-class NotebookRange implements vscode.NotebookRange {
+export class NotebookRange implements vscode.NotebookRange {
   readonly start: number;
   readonly end: number;
   get isEmpty(): boolean {
@@ -1139,7 +1139,7 @@ export class TestVsCode extends Data.TaggedClass("TestVsCode")<{
     return Ref.get(this.statusBarProviders);
   }
 
-  fireNotebookDocumentChange(event: vscode.NotebookDocumentChangeEvent) {
+  notebookChange(event: vscode.NotebookDocumentChangeEvent) {
     return PubSub.publish(this.documentChangesPubSub, event);
   }
   static make = Effect.fnUntraced(function* (

extension/src/services/VsCode.ts

@@ -271,7 +271,7 @@ export class Workspace extends Effect.Service<Workspace>()("Workspace", {
         return Stream.asyncPush<vscode.NotebookDocumentChangeEvent>((emit) =>
           Effect.acquireRelease(
             Effect.sync(() =>
-              api.onDidChangeNotebookDocument((evt) => emit.single(evt)),
+              api.onDidChangeNotebookDocument((event) => emit.single(event)),
             ),
             (disposable) => Effect.sync(() => disposable.dispose()),
           ),

extension/src/services/__tests__/CellStateManager.test.ts

@@ -1,15 +1,40 @@
 import { describe, expect, it } from "@effect/vitest";
-import { Effect } from "effect";
+import { Effect, Layer, Option, Ref, Stream, TestClock } from "effect";
 import {
   createTestNotebookDocument,
+  NotebookRange,
   TestVsCode,
 } from "../../__mocks__/TestVsCode.ts";
+import type { MarimoCommand } from "../../types.ts";
 import { getNotebookCellId } from "../../utils/notebook.ts";
+import { CellStateManager } from "../CellStateManager.ts";
+import { LanguageClient } from "../LanguageClient.ts";
 import { VsCode } from "../VsCode.ts";
 
 const withTestCtx = Effect.fnUntraced(function* () {
   const vscode = yield* TestVsCode.make();
-  return { vscode, layer: vscode.layer };
+  const executions = yield* Ref.make<ReadonlyArray<MarimoCommand>>([]);
+  const layer = Layer.empty.pipe(
+    Layer.merge(CellStateManager.Default),
+    Layer.provideMerge(vscode.layer),
+    Layer.provide(
+      Layer.succeed(
+        LanguageClient,
+        LanguageClient.make({
+          manage() {
+            return Effect.void;
+          },
+          executeCommand(cmd) {
+            return Ref.update(executions, (arr) => [...arr, cmd]);
+          },
+          streamOf() {
+            return Stream.never;
+          },
+        }),
+      ),
+    ),
+  );
+  return { vscode, layer, executions };
 });
 
 describe("CellStateManager", () => {
@@ -56,4 +81,134 @@ describe("CellStateManager", () => {
       }).pipe(Effect.provide(ctx.layer));
     }),
   );
+
+  it.effect(
+    "deleting cell from notebook sends delete_cell command",
+    Effect.fnUntraced(function* () {
+      const ctx = yield* withTestCtx();
+
+      yield* Effect.gen(function* () {
+        const code = yield* VsCode;
+
+        const notebook = createTestNotebookDocument(
+          "/test/notebook.py",
+          new code.NotebookData([
+            new code.NotebookCellData(
+              code.NotebookCellKind.Code,
+              "x = 1",
+              "python",
+            ),
+            new code.NotebookCellData(
+              code.NotebookCellKind.Code,
+              "y = 2",
+              "python",
+            ),
+            new code.NotebookCellData(
+              code.NotebookCellKind.Code,
+              "z = 3",
+              "python",
+            ),
+          ]),
+        );
+
+        yield* ctx.vscode.addNotebookDocument(notebook);
+        yield* TestClock.adjust("10 millis");
+
+        const cellToDelete = notebook.cellAt(1);
+
+        yield* ctx.vscode.notebookChange({
+          notebook,
+          metadata: undefined,
+          cellChanges: [],
+          contentChanges: [
+            {
+              range: new NotebookRange(1, 2), // Delete cell at index 1
+              removedCells: [cellToDelete],
+              addedCells: [],
+            },
+          ],
+        });
+
+        yield* TestClock.adjust("10 millis");
+
+        const commands = yield* Ref.get(ctx.executions);
+        expect(commands).toMatchInlineSnapshot(`
+          [
+            {
+              "command": "marimo.api",
+              "params": {
+                "method": "delete_cell",
+                "params": {
+                  "inner": {
+                    "cellId": "file:///test/notebook.py#cell-1",
+                  },
+                  "notebookUri": "file:///test/notebook.py",
+                },
+              },
+            },
+          ]
+        `);
+      }).pipe(Effect.provide(ctx.layer));
+    }),
+  );
+
+  it.effect(
+    "updates marimo.notebook.hasStaleCells context when cells become stale",
+    Effect.fnUntraced(function* () {
+      const ctx = yield* withTestCtx();
+
+      yield* Effect.gen(function* () {
+        const code = yield* VsCode;
+
+        const editor = TestVsCode.makeNotebookEditor(
+          "/test/notebook.py",
+          new code.NotebookData([
+            new code.NotebookCellData(
+              code.NotebookCellKind.Code,
+              "x = 1",
+              "python",
+            ),
+            new code.NotebookCellData(
+              code.NotebookCellKind.Code,
+              "y = 2",
+              "python",
+            ),
+          ]),
+        );
+
+        yield* ctx.vscode.addNotebookDocument(editor.notebook);
+        yield* ctx.vscode.setActiveNotebookEditor(Option.some(editor));
+        yield* TestClock.adjust("10 millis");
+
+        // Clear all previous executions
+        yield* Ref.update(ctx.vscode.executions, () => []);
+
+        // Trigger a cell content change to mark it as stale
+        const cell0 = editor.notebook.cellAt(0);
+        yield* ctx.vscode.notebookChange({
+          notebook: editor.notebook,
+          metadata: undefined,
+          cellChanges: [
+            {
+              cell: cell0,
+              document: cell0.document,
+              metadata: undefined,
+              outputs: [],
+              executionSummary: undefined,
+            },
+          ],
+          contentChanges: [],
+        });
+
+        yield* TestClock.adjust("10 millis");
+      }).pipe(Effect.provide(ctx.layer));
+
+      expect(yield* Ref.get(ctx.vscode.executions)).toEqual([
+        {
+          command: "setContext",
+          args: ["marimo.notebook.hasStaleCells", true],
+        },
+      ]);
+    }),
+  );
 });

extension/src/services/__tests__/ExecutionRegistry.test.ts

@@ -28,8 +28,8 @@ const withTestCtx = Effect.fnUntraced(function* (
   const layer = Layer.empty.pipe(
     Layer.merge(ExecutionRegistry.Default),
     Layer.merge(CellStateManager.Default),
-    Layer.provideMerge(vscode.layer),
     Layer.provide(TestLanguageClientLive),
+    Layer.provideMerge(vscode.layer),
   );
   return { vscode, layer };
 });