diff --git a/extension/src/__mocks__/TestVsCode.ts b/extension/src/__mocks__/TestVsCode.ts index 7859127a..5dc2361b 100644 --- a/extension/src/__mocks__/TestVsCode.ts +++ b/extension/src/__mocks__/TestVsCode.ts @@ -1474,6 +1474,7 @@ export class TestVsCode extends Data.TaggedClass("TestVsCode")<{ >; readonly documentChangesPubSub: PubSub.PubSub; readonly documentOpenedPubSub: PubSub.PubSub; + readonly documentClosedPubSub: PubSub.PubSub; readonly setActiveNotebookEditor: ( editor: Option.Option, ) => Effect.Effect; @@ -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; @@ -1581,6 +1586,7 @@ export class TestVsCode extends Data.TaggedClass("TestVsCode")<{ yield* PubSub.unbounded(); const documentOpened = yield* PubSub.unbounded(); + const documentClosed = yield* PubSub.unbounded(); const commands = yield* Ref.make( HashSet.empty(), @@ -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* () { diff --git a/extension/src/notebook/NotebookEditorRegistry.ts b/extension/src/notebook/NotebookEditorRegistry.ts index a823c547..40e02327 100644 --- a/extension/src/notebook/NotebookEditorRegistry.ts +++ b/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 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); diff --git a/extension/src/notebook/__tests__/NotebookEditorRegistry.test.ts b/extension/src/notebook/__tests__/NotebookEditorRegistry.test.ts index 2593234f..ecfbec09 100644 --- a/extension/src/notebook/__tests__/NotebookEditorRegistry.test.ts +++ b/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* () {