Fix closed notebook reopening on file auto-reload (#353)

extension/src/__mocks__/TestVsCode.ts

@@ -1474,6 +1474,7 @@ export class TestVsCode extends Data.TaggedClass("TestVsCode")<{
   >;
   readonly documentChangesPubSub: PubSub.PubSub<vscode.NotebookDocumentChangeEvent>;
   readonly documentOpenedPubSub: PubSub.PubSub<vscode.NotebookDocument>;
+  readonly documentClosedPubSub: PubSub.PubSub<vscode.NotebookDocument>;
   readonly setActiveNotebookEditor: (
     editor: Option.Option<vscode.NotebookEditor>,
   ) => Effect.Effect<void>;
@@ -1550,6 +1551,10 @@ export class TestVsCode extends Data.TaggedClass("TestVsCode")<{
     return PubSub.publish(this.documentOpenedPubSub, doc);
   }
 
+  closeNotebook(doc: vscode.NotebookDocument) {
+    return PubSub.publish(this.documentClosedPubSub, doc);
+  }
+
   static make = Effect.fn(function* (
     options: {
       initialDocuments?: Array<vscode.NotebookDocument>;
@@ -1581,6 +1586,7 @@ export class TestVsCode extends Data.TaggedClass("TestVsCode")<{
       yield* PubSub.unbounded<vscode.NotebookDocumentChangeEvent>();
 
     const documentOpened = yield* PubSub.unbounded<vscode.NotebookDocument>();
+    const documentClosed = yield* PubSub.unbounded<vscode.NotebookDocument>();
 
     const commands = yield* Ref.make(
       HashSet.empty<MarimoCommand | DynamicCommand>(),
@@ -1887,7 +1893,7 @@ export class TestVsCode extends Data.TaggedClass("TestVsCode")<{
             return Stream.fromPubSub(documentChanges);
           },
           notebookDocumentClosed() {
-            return Stream.never;
+            return Stream.fromPubSub(documentClosed);
           },
           textDocumentChanges() {
             return Stream.never;
@@ -2219,6 +2225,7 @@ export class TestVsCode extends Data.TaggedClass("TestVsCode")<{
       statusBarProviders,
       documentChangesPubSub: documentChanges,
       documentOpenedPubSub: documentOpened,
+      documentClosedPubSub: documentClosed,
       affinityUpdates,
       setActiveNotebookEditor: (editor) =>
         Effect.gen(function* () {

extension/src/notebook/NotebookEditorRegistry.ts

@@ -1,6 +1,7 @@
 import { Effect, HashMap, Option, Ref, Stream, SubscriptionRef } from "effect";
 import type * as vscode from "vscode";
 
+import { NOTEBOOK_TYPE } from "../constants.ts";
 import { VsCode } from "../platform/VsCode.ts";
 import { MarimoNotebookDocument } from "../schemas/MarimoNotebookDocument.ts";
 import type { NotebookId } from "../schemas/MarimoNotebookDocument.ts";
@@ -69,6 +70,29 @@ export class NotebookEditorRegistry extends Effect.Service<NotebookEditorRegistr
         ),
       );
 
+      // Remove editors when their notebook is closed
+      yield* Effect.forkScoped(
+        code.workspace.notebookDocumentClosed().pipe(
+          Stream.filter((nb) => nb.notebookType === NOTEBOOK_TYPE),
+          Stream.mapEffect(
+            Effect.fn(function* (nb) {
+              const notebook = MarimoNotebookDocument.tryFrom(nb);
+              if (Option.isSome(notebook)) {
+                yield* Ref.update(ref, HashMap.remove(notebook.value.id));
+                yield* Effect.logInfo(
+                  "Removed closed notebook from registry",
+                ).pipe(
+                  Effect.annotateLogs({
+                    notebookUri: notebook.value.id,
+                  }),
+                );
+              }
+            }),
+          ),
+          Stream.runDrain,
+        ),
+      );
+
       return {
         getNotebookEditors() {
           return Effect.map(Ref.get(ref), HashMap.toEntries);

extension/src/notebook/__tests__/NotebookEditorRegistry.test.ts

@@ -94,6 +94,46 @@ it.effect(
   }),
 );
 
+it.effect(
+  "should remove editor when notebook is closed",
+  Effect.fn(function* () {
+    const vscode = yield* TestVsCode.make();
+
+    yield* Effect.provide(
+      Effect.gen(function* () {
+        const code = yield* VsCode;
+        const registry = yield* NotebookEditorRegistry;
+
+        // Create and activate a mock notebook
+        const notebook = createTestNotebookDocument(
+          code.Uri.file("/test/notebook_mo.py"),
+        );
+        const mockEditor = createTestNotebookEditor(notebook);
+
+        yield* vscode.setActiveNotebookEditor(Option.some(mockEditor));
+        yield* TestClock.adjust("10 millis");
+
+        // Verify the editor is tracked
+        const before = yield* registry.getLastNotebookEditor(
+          notebook.uri.toString(),
+        );
+        expect(Option.isSome(before)).toBe(true);
+
+        // Close the notebook
+        yield* vscode.closeNotebook(notebook);
+        yield* TestClock.adjust("10 millis");
+
+        // Verify the editor is removed
+        const after = yield* registry.getLastNotebookEditor(
+          notebook.uri.toString(),
+        );
+        expect(Option.isNone(after)).toBe(true);
+      }),
+      makeRegistryLayer(vscode),
+    );
+  }),
+);
+
 it.effect(
   "should track stream of active notebook editor changes",
   Effect.fn(function* () {