fix microsoft#202232 (microsoft#205351)

Author: jrieken

src/vs/editor/common/languageSelector.ts

@@ -131,3 +131,14 @@ export function score(selector: LanguageSelector | undefined, candidateUri: URI,
 		return 0;
 	}
 }
+
+
+export function targetsNotebooks(selector: LanguageSelector): boolean {
+	if (typeof selector === 'string') {
+		return false;
+	} else if (Array.isArray(selector)) {
+		return selector.some(targetsNotebooks);
+	} else {
+		return !!(<LanguageFilter>selector).notebookType;
+	}
+}

src/vs/workbench/api/common/extHost.api.impl.ts

@@ -11,7 +11,7 @@ import { Schemas, matchesScheme } from 'vs/base/common/network';
 import Severity from 'vs/base/common/severity';
 import { URI } from 'vs/base/common/uri';
 import { TextEditorCursorStyle } from 'vs/editor/common/config/editorOptions';
-import { score } from 'vs/editor/common/languageSelector';
+import { score, targetsNotebooks } from 'vs/editor/common/languageSelector';
 import * as languageConfiguration from 'vs/editor/common/languages/languageConfiguration';
 import { OverviewRulerLane } from 'vs/editor/common/model';
 import { ExtensionIdentifierSet, IExtensionDescription } from 'vs/platform/extensions/common/extensions';
@@ -533,8 +533,12 @@ export function createApiFactoryAndRegisterActors(accessor: ServicesAccessor): I
 				return extHostLanguages.changeLanguage(document.uri, languageId);
 			},
 			match(selector: vscode.DocumentSelector, document: vscode.TextDocument): number {
-				const notebook = extHostDocuments.getDocumentData(document.uri)?.notebook;
-				return score(typeConverters.LanguageSelector.from(selector), document.uri, document.languageId, true, notebook?.uri, notebook?.notebookType);
+				const interalSelector = typeConverters.LanguageSelector.from(selector);
+				let notebook: vscode.NotebookDocument | undefined;
+				if (targetsNotebooks(interalSelector)) {
+					notebook = extHostNotebook.notebookDocuments.find(value => Boolean(value.getCell(document.uri)))?.apiNotebook;
+				}
+				return score(interalSelector, document.uri, document.languageId, true, notebook?.uri, notebook?.notebookType);
 			},
 			registerCodeActionsProvider(selector: vscode.DocumentSelector, provider: vscode.CodeActionProvider, metadata?: vscode.CodeActionProviderMetadata): vscode.Disposable {
 				return extHostLanguageFeatures.registerCodeActionProvider(extension, checkSelector(selector), provider, metadata);

src/vs/workbench/api/common/extHostDocumentData.ts

@@ -37,7 +37,6 @@ export class ExtHostDocumentData extends MirrorTextModel {
 		uri: URI, lines: string[], eol: string, versionId: number,
 		private _languageId: string,
 		private _isDirty: boolean,
-		public readonly notebook?: vscode.NotebookDocument | undefined
 	) {
 		super(uri, lines, eol, versionId);
 	}

src/vs/workbench/api/common/extHostDocumentsAndEditors.ts

@@ -9,7 +9,7 @@ import { Emitter, Event } from 'vs/base/common/event';
 import { dispose } from 'vs/base/common/lifecycle';
 import { URI } from 'vs/base/common/uri';
 import { createDecorator } from 'vs/platform/instantiation/common/instantiation';
-import { ExtHostDocumentsAndEditorsShape, IDocumentsAndEditorsDelta, IModelAddedData, MainContext } from 'vs/workbench/api/common/extHost.protocol';
+import { ExtHostDocumentsAndEditorsShape, IDocumentsAndEditorsDelta, MainContext } from 'vs/workbench/api/common/extHost.protocol';
 import { ExtHostDocumentData } from 'vs/workbench/api/common/extHostDocumentData';
 import { IExtHostRpcService } from 'vs/workbench/api/common/extHostRpcService';
 import { ExtHostTextEditor } from 'vs/workbench/api/common/extHostTextEditor';
@@ -31,14 +31,6 @@ class Reference<T> {
 	}
 }
 
-export interface IExtHostModelAddedData extends IModelAddedData {
-	notebook?: vscode.NotebookDocument;
-}
-
-export interface IExtHostDocumentsAndEditorsDelta extends IDocumentsAndEditorsDelta {
-	addedDocuments?: IExtHostModelAddedData[];
-}
-
 export class ExtHostDocumentsAndEditors implements ExtHostDocumentsAndEditorsShape {
 
 	readonly _serviceBrand: undefined;
@@ -67,7 +59,7 @@ export class ExtHostDocumentsAndEditors implements ExtHostDocumentsAndEditorsSha
 		this.acceptDocumentsAndEditorsDelta(delta);
 	}
 
-	acceptDocumentsAndEditorsDelta(delta: IExtHostDocumentsAndEditorsDelta): void {
+	acceptDocumentsAndEditorsDelta(delta: IDocumentsAndEditorsDelta): void {
 
 		const removedDocuments: ExtHostDocumentData[] = [];
 		const addedDocuments: ExtHostDocumentData[] = [];
@@ -105,7 +97,6 @@ export class ExtHostDocumentsAndEditors implements ExtHostDocumentsAndEditorsSha
 						data.versionId,
 						data.languageId,
 						data.isDirty,
-						data.notebook
 					));
 					this._documents.set(resource, ref);
 					addedDocuments.push(ref.value);

src/vs/workbench/api/common/extHostInteractive.ts

@@ -52,7 +52,6 @@ export class ExtHostInteractive implements ExtHostInteractiveShape {
 				uri: uri,
 				isDirty: false,
 				versionId: 1,
-				notebook: this._extHostNotebooks.getNotebookDocument(URI.revive(notebookUri))?.apiNotebook
 			}]
 		});
 	}

src/vs/workbench/api/common/extHostNotebook.ts

@@ -625,7 +625,7 @@ export class ExtHostNotebookController implements ExtHostNotebookShape {
 				);
 
 				// add cell document as vscode.TextDocument
-				addedCellDocuments.push(...modelData.cells.map(cell => ExtHostCell.asModelAddData(document.apiNotebook, cell)));
+				addedCellDocuments.push(...modelData.cells.map(cell => ExtHostCell.asModelAddData(cell)));
 
 				this._documents.get(uri)?.dispose();
 				this._documents.set(uri, document);

src/vs/workbench/api/common/extHostNotebookDocument.ts

@@ -7,7 +7,7 @@ import { Schemas } from 'vs/base/common/network';
 import { URI } from 'vs/base/common/uri';
 import * as extHostProtocol from 'vs/workbench/api/common/extHost.protocol';
 import { ExtHostDocuments } from 'vs/workbench/api/common/extHostDocuments';
-import { ExtHostDocumentsAndEditors, IExtHostModelAddedData } from 'vs/workbench/api/common/extHostDocumentsAndEditors';
+import { ExtHostDocumentsAndEditors } from 'vs/workbench/api/common/extHostDocumentsAndEditors';
 import * as extHostTypeConverters from 'vs/workbench/api/common/extHostTypeConverters';
 import { NotebookRange } from 'vs/workbench/api/common/extHostTypes';
 import * as notebookCommon from 'vs/workbench/contrib/notebook/common/notebookCommon';
@@ -33,15 +33,14 @@ class RawContentChangeEvent {
 
 export class ExtHostCell {
 
-	static asModelAddData(notebook: vscode.NotebookDocument, cell: extHostProtocol.NotebookCellDto): IExtHostModelAddedData {
+	static asModelAddData(cell: extHostProtocol.NotebookCellDto): extHostProtocol.IModelAddedData {
 		return {
 			EOL: cell.eol,
 			lines: cell.source,
 			languageId: cell.language,
 			uri: cell.uri,
 			isDirty: false,
-			versionId: 1,
-			notebook
+			versionId: 1
 		};
 	}
 
@@ -356,7 +355,7 @@ export class ExtHostNotebookDocument {
 		}
 
 		const contentChangeEvents: RawContentChangeEvent[] = [];
-		const addedCellDocuments: IExtHostModelAddedData[] = [];
+		const addedCellDocuments: extHostProtocol.IModelAddedData[] = [];
 		const removedCellDocuments: URI[] = [];
 
 		splices.reverse().forEach(splice => {
@@ -365,7 +364,7 @@ export class ExtHostNotebookDocument {
 
 				const extCell = new ExtHostCell(this, this._textDocumentsAndEditors, cell);
 				if (!initialization) {
-					addedCellDocuments.push(ExtHostCell.asModelAddData(this.apiNotebook, cell));
+					addedCellDocuments.push(ExtHostCell.asModelAddData(cell));
 				}
 				return extCell;
 			});
@@ -443,7 +442,17 @@ export class ExtHostNotebookDocument {
 		return this._cells[index];
 	}
 
-	getCell(cellHandle: number): ExtHostCell | undefined {
+	getCell(cellHandle: number | URI): ExtHostCell | undefined {
+		if (URI.isUri(cellHandle)) {
+			const data = notebookCommon.CellUri.parse(cellHandle);
+			if (!data) {
+				return undefined;
+			}
+			if (data.notebook.toString() !== this.uri.toString()) {
+				return undefined;
+			}
+			cellHandle = data.handle;
+		}
 		return this._cells.find(cell => cell.handle === cellHandle);
 	}