Improve notebook cell model lifecycle

packages/monaco/src/browser/monaco-text-model-service.ts

@@ -113,7 +113,7 @@ export class MonacoTextModelService implements ITextModelService {
      * creates a model which is not saved by the model service.
      * this will therefore also not be created on backend side.
      */
-    createUnmangedModel(raw: monaco.Uri | URI): Promise<MonacoEditorModel> {
+    createUnmanagedModel(raw: monaco.Uri | URI): Promise<MonacoEditorModel> {
         return this.loadModel(new URI(raw.toString()));
     }
 

packages/notebook/src/browser/service/notebook-monaco-text-model-service.ts

@@ -31,15 +31,15 @@ export class NotebookMonacoTextModelService {
     protected readonly monacoTextModelService: MonacoTextModelService;
 
     protected readonly cellmodels = new ReferenceCollection<string, MonacoEditorModel>(
-        uri => this.monacoTextModelService.createUnmangedModel(new URI(uri))
+        uri => this.monacoTextModelService.createUnmanagedModel(new URI(uri))
     );
 
     getOrCreateNotebookCellModelReference(uri: URI): Promise<Reference<MonacoEditorModel>> {
         return this.cellmodels.acquire(uri.toString());
     }
 
     async createTextModelsForNotebook(notebook: NotebookModel): Promise<void> {
-        await Promise.all(notebook.cells.map(cell => this.getOrCreateNotebookCellModelReference(cell.uri)));
+        await Promise.all(notebook.cells.map(cell => cell.resolveTextModel()));
     }
 
     get onDidCreateNotebookCellModel(): Event<MonacoEditorModel> {

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

@@ -119,7 +119,7 @@ export class NotebookService implements Disposable {
         this.notebookModels.set(resource.uri.toString(), model);
         // Resolve cell text models right after creating the notebook model
         // This ensures that all text models are available in the plugin host
-        this.textModelService.createTextModelsForNotebook(model);
+        await this.textModelService.createTextModelsForNotebook(model);
         this.didAddNotebookDocumentEmitter.fire(model);
         model.onDidDispose(() => {
             this.notebookModels.delete(resource.uri.toString());

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

@@ -264,7 +264,6 @@ export class NotebookCellModel implements NotebookCell, Disposable {
         this.onDidChangeMetadataEmitter.dispose();
         this.onDidChangeInternalMetadataEmitter.dispose();
         this.onDidChangeLanguageEmitter.dispose();
-        this.textModel?.dispose();
         this.toDispose.dispose();
     }
 
@@ -356,9 +355,10 @@ export class NotebookCellModel implements NotebookCell, Disposable {
 
         const ref = await this.textModelService.getOrCreateNotebookCellModelReference(this.uri);
         this.textModel = ref.object;
-        this.textModel.onDidChangeContent(e => {
+        this.toDispose.push(ref);
+        this.toDispose.push(this.textModel.onDidChangeContent(e => {
             this.props.source = e.model.getText();
-        });
+        }));
         return ref.object;
     }
 

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

@@ -331,7 +331,7 @@ export class NotebookModel implements Saveable, Disposable {
 
     }
 
-    protected async replaceCells(start: number, deleteCount: number, newCells: CellData[], computeUndoRedo: boolean): Promise<void> {
+    protected replaceCells(start: number, deleteCount: number, newCells: CellData[], computeUndoRedo: boolean): void {
         const cells = newCells.map(cell => {
             const handle = this.nextHandle++;
             return this.cellModelFactory({
@@ -362,10 +362,6 @@ export class NotebookModel implements Saveable, Disposable {
                 async () => this.replaceCells(start, deleteCount, newCells, false));
         }
 
-        // Ensure that all text model have been created
-        // Otherwise we run into a race condition once we fire `onDidChangeContent`
-        await Promise.all(cells.map(cell => cell.resolveTextModel()));
-
         this.onDidAddOrRemoveCellEmitter.fire({ rawEvent: { kind: NotebookCellsChangeType.ModelChange, changes }, newCellIds: cells.map(cell => cell.handle) });
         this.onDidChangeContentEmitter.queue({ kind: NotebookCellsChangeType.ModelChange, changes });
         if (cells.length > 0) {

packages/notebook/src/common/notebook-common.ts

@@ -273,7 +273,7 @@ export namespace CellUri {
         const s = handle.toString(_radix);
         const p = s.length < _lengths.length ? _lengths[s.length - 1] : 'z';
 
-        const fragment = `${p}${s}s${Buffer.from(BinaryBuffer.fromString(notebook.scheme).buffer).toString('base64')} `;
+        const fragment = `${p}${s}s${Buffer.from(BinaryBuffer.fromString(notebook.scheme).buffer).toString('base64')}`;
         return notebook.withScheme(cellUriScheme).withFragment(fragment);
     }
 

packages/plugin-ext/src/plugin/editors-and-documents.ts

@@ -46,6 +46,10 @@ export class EditorsAndDocumentsExtImpl implements EditorsAndDocumentsExt {
     private readonly editors = new Map<string, TextEditorExt>();
 
     async $acceptEditorsAndDocumentsDelta(delta: EditorsAndDocumentsDelta): Promise<void> {
+        this.acceptEditorsAndDocumentsDelta(delta);
+    }
+
+    acceptEditorsAndDocumentsDelta(delta: EditorsAndDocumentsDelta): void {
         const removedDocuments = new Array<DocumentDataExt>();
         const addedDocuments = new Array<DocumentDataExt>();
         const removedEditors = new Array<TextEditorExt>();

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

@@ -27,6 +27,7 @@ import * as typeConverters from '../type-converters';
 import { ModelAddedData, NotebookCellDto, NotebookCellsChangedEventDto, NotebookModelAddedData, NotebookOutputDto } from '../../common';
 import { NotebookRange } from '../types-impl';
 import { DocumentsExtImpl } from '../documents';
+import { UriComponents } from '../../common/uri-components';
 
 class RawContentChangeEvent {
 
@@ -345,6 +346,9 @@ export class NotebookDocument implements Disposable {
             return;
         }
 
+        const addedDocuments: ModelAddedData[] = [];
+        const removedDocuments: UriComponents[] = [];
+
         const contentChangeEvents: RawContentChangeEvent[] = [];
 
         splices.reverse().forEach(splice => {
@@ -353,9 +357,7 @@ export class NotebookDocument implements Disposable {
 
                 const extCell = new Cell(this, this.editorsAndDocuments, cell);
                 if (!initialization) {
-                    this.editorsAndDocuments.$acceptEditorsAndDocumentsDelta({
-                        addedDocuments: [Cell.asModelAddData(cell)]
-                    });
+                    addedDocuments.push(Cell.asModelAddData(cell));
                 }
                 return extCell;
             });
@@ -364,10 +366,18 @@ export class NotebookDocument implements Disposable {
             const deletedItems = this.cells.splice(splice.start, splice.deleteCount, ...newCells);
             for (const cell of deletedItems) {
                 changeEvent.deletedItems.push(cell.apiCell);
+                removedDocuments.push(cell.uri.toComponents());
             }
             contentChangeEvents.push(changeEvent);
         });
 
+        if (addedDocuments.length > 0 || removedDocuments.length > 0) {
+            this.editorsAndDocuments.acceptEditorsAndDocumentsDelta({
+                addedDocuments,
+                removedDocuments
+            });
+        }
+
         if (bucket) {
             for (const changeEvent of contentChangeEvents) {
                 bucket.push(changeEvent.asApiEvent());

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, ModelAddedData, NotebookCellStatusBarListDto, NotebookDataDto,
+    CommandRegistryExt, NotebookCellStatusBarListDto, NotebookDataDto,
     NotebookDocumentsAndEditorsDelta, NotebookDocumentShowOptions, NotebookDocumentsMain, NotebookEditorAddData, NotebookEditorsMain, NotebooksExt, NotebooksMain, Plugin,
     PLUGIN_RPC_CONTEXT
 } from '../../common';
@@ -205,7 +205,6 @@ 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);
@@ -226,10 +225,12 @@ export class NotebooksExtImpl implements NotebooksExt {
             }
         }
 
-        // publish all removed cell documents first
-        await this.textDocumentsAndEditors.$acceptEditorsAndDocumentsDelta({
-            removedDocuments: removedCellDocuments
-        });
+        if (removedCellDocuments.length > 0) {
+            // publish all removed cell documents first
+            this.textDocumentsAndEditors.acceptEditorsAndDocumentsDelta({
+                removedDocuments: removedCellDocuments
+            });
+        }
 
         if (delta.addedDocuments) {
             for (const modelData of delta.addedDocuments) {
@@ -250,17 +251,18 @@ export class NotebooksExtImpl implements NotebooksExt {
                 this.documents.get(uri.toString())?.dispose();
                 this.documents.set(uri.toString(), document);
 
-                addedCellDocuments.push(...modelData.cells.map(cell => Cell.asModelAddData(cell)));
+                if (modelData.cells.length > 0) {
+                    // Publish new cell documents before calling the notebook document open event
+                    // During this event, extensions might request the cell document and we want to make sure it is available
+                    this.textDocumentsAndEditors.acceptEditorsAndDocumentsDelta({
+                        addedDocuments: 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)) {