Fix notebook model/cell disposal

packages/notebook/src/browser/notebook-editor-widget.tsx

@@ -258,6 +258,7 @@ export class NotebookEditorWidget extends ReactWidget implements Navigatable, Sa
         this.onDidReceiveKernelMessageEmitter.dispose();
         this.onPostRendererMessageEmitter.dispose();
         this.viewportService.dispose();
+        this._model?.dispose();
         super.dispose();
     }
 

packages/notebook/src/browser/service/notebook-service.ts

@@ -121,6 +121,10 @@ export class NotebookService implements Disposable {
         // This ensures that all text models are available in the plugin host
         this.textModelService.createTextModelsForNotebook(model);
         this.didAddNotebookDocumentEmitter.fire(model);
+        model.onDidDispose(() => {
+            this.notebookModels.delete(resource.uri.toString());
+            this.didRemoveNotebookDocumentEmitter.fire(model);
+        });
         return model;
     }
 

packages/notebook/src/browser/view-model/notebook-model.ts

@@ -71,6 +71,9 @@ export class NotebookModel implements Saveable, Disposable {
     protected readonly onDidChangeSelectedCellEmitter = new Emitter<NotebookCellModel | undefined>();
     readonly onDidChangeSelectedCell = this.onDidChangeSelectedCellEmitter.event;
 
+    protected readonly onDidDisposeEmitter = new Emitter<void>();
+    readonly onDidDispose = this.onDidDisposeEmitter.event;
+
     get onDidChangeReadOnly(): Event<boolean | MarkdownString> {
         return this.props.resource.onDidChangeReadOnly ?? Event.None;
     }
@@ -150,6 +153,7 @@ export class NotebookModel implements Saveable, Disposable {
         this.onDidChangeContentEmitter.dispose();
         this.onDidChangeSelectedCellEmitter.dispose();
         this.cells.forEach(cell => cell.dispose());
+        this.onDidDisposeEmitter.fire();
     }
 
     async save(options: SaveOptions): Promise<void> {

packages/plugin-ext/src/plugin/notebook/notebooks.ts

@@ -22,7 +22,7 @@ import { CancellationToken, Disposable, DisposableCollection, Emitter, Event, UR
 import { URI as TheiaURI } from '../types-impl';
 import * as theia from '@theia/plugin';
 import {
-    CommandRegistryExt, NotebookCellStatusBarListDto, NotebookDataDto,
+    CommandRegistryExt, ModelAddedData, NotebookCellStatusBarListDto, NotebookDataDto,
     NotebookDocumentsAndEditorsDelta, NotebookDocumentShowOptions, NotebookDocumentsMain, NotebookEditorAddData, NotebookEditorsMain, NotebooksExt, NotebooksMain, Plugin,
     PLUGIN_RPC_CONTEXT
 } from '../../common';
@@ -204,6 +204,8 @@ export class NotebooksExtImpl implements NotebooksExt {
     }
 
     async $acceptDocumentsAndEditorsDelta(delta: NotebookDocumentsAndEditorsDelta): Promise<void> {
+        const removedCellDocuments: UriComponents[] = [];
+        const addedCellDocuments: ModelAddedData[] = [];
         if (delta.removedDocuments) {
             for (const uri of delta.removedDocuments) {
                 const revivedUri = URI.fromComponents(uri);
@@ -213,6 +215,7 @@ export class NotebooksExtImpl implements NotebooksExt {
                     document.dispose();
                     this.documents.delete(revivedUri.toString());
                     this.onDidCloseNotebookDocumentEmitter.fire(document.apiNotebook);
+                    removedCellDocuments.push(...document.apiNotebook.getCells().map(cell => cell.document.uri));
                 }
 
                 for (const editor of this.editors.values()) {
@@ -223,6 +226,11 @@ export class NotebooksExtImpl implements NotebooksExt {
             }
         }
 
+        // publish all removed cell documents first
+        await this.textDocumentsAndEditors.$acceptEditorsAndDocumentsDelta({
+            removedDocuments: removedCellDocuments
+        });
+
         if (delta.addedDocuments) {
             for (const modelData of delta.addedDocuments) {
                 const uri = TheiaURI.from(modelData.uri);
@@ -242,14 +250,17 @@ export class NotebooksExtImpl implements NotebooksExt {
                 this.documents.get(uri.toString())?.dispose();
                 this.documents.set(uri.toString(), document);
 
-                this.textDocumentsAndEditors.$acceptEditorsAndDocumentsDelta({
-                    addedDocuments: modelData.cells.map(cell => Cell.asModelAddData(cell))
-                });
+                addedCellDocuments.push(...modelData.cells.map(cell => Cell.asModelAddData(cell)));
 
                 this.onDidOpenNotebookDocumentEmitter.fire(document.apiNotebook);
             }
         }
 
+        // publish all added cell documents in a separate call
+        await this.textDocumentsAndEditors.$acceptEditorsAndDocumentsDelta({
+            addedDocuments: addedCellDocuments
+        });
+
         if (delta.addedEditors) {
             for (const editorModelData of delta.addedEditors) {
                 if (this.editors.has(editorModelData.id)) {