Preserve previous code cell's language when creating a new code cell from an existing context. (#18646)

src/sql/workbench/contrib/notebook/browser/cellToolbarActions.ts

@@ -75,7 +75,7 @@ export class SplitCellAction extends CellActionBase {
 	doRun(context: CellContext): Promise<void> {
 		let model = context.model;
 		let index = model.cells.findIndex((cell) => cell.id === context.cell.id);
-		context.model?.splitCell(context.cell.cellType, this.notebookService, index);
+		context.model?.splitCell(context.cell.cellType, this.notebookService, index, context.cell.metadata?.language);
 		return Promise.resolve();
 	}
 	public setListener(context: CellContext) {
@@ -246,7 +246,7 @@ export class AddCellFromContextAction extends CellActionBase {
 			if (index !== undefined && this.isAfter) {
 				index += 1;
 			}
-			model.addCell(this.cellType, index);
+			model.addCell(this.cellType, index, context.cell.metadata?.language);
 		} catch (error) {
 			let message = getErrorMessage(error);
 

src/sql/workbench/contrib/notebook/browser/notebookActions.ts

@@ -74,7 +74,7 @@ export class AddCellAction extends Action {
 				}
 			}
 			if (context?.model) {
-				context.model.addCell(this.cellType, index);
+				context.model.addCell(this.cellType, index, context.cell.metadata?.language);
 				context.model.sendNotebookTelemetryActionEvent(TelemetryKeys.NbTelemetryAction.AddCell, { cell_type: this.cellType });
 			}
 		} else {

src/sql/workbench/contrib/notebook/test/electron-browser/notebookModel.test.ts

@@ -1006,6 +1006,33 @@ suite('notebook model', function (): void {
 		assert.strictEqual(model.languageInfo.name, 'fake', 'Notebook language info is not set properly');
 	});
 
+	test('Should add new cell with provided cell language', async function () {
+		let mockContentManager = TypeMoq.Mock.ofType(NotebookEditorContentLoader);
+		mockContentManager.setup(c => c.loadContent()).returns(() => Promise.resolve(expectedNotebookContent));
+		defaultModelOptions.contentLoader = mockContentManager.object;
+
+		let model = new NotebookModel(defaultModelOptions, undefined, logService, undefined, new NullAdsTelemetryService(), queryConnectionService.object, configurationService, undoRedoService);
+		await model.loadContents();
+
+		const newLanguage = 'CustomCellLanguage';
+		let cell = model.addCell(CellTypes.Code, undefined, newLanguage);
+		assert.strictEqual(model.cells.length, expectedNotebookContent.cells.length + 1, 'New cell was not added to list of cells.');
+		assert.strictEqual(cell.language, newLanguage);
+	});
+
+	test('Should add new cell with default language if none is provided', async function () {
+		let mockContentManager = TypeMoq.Mock.ofType(NotebookEditorContentLoader);
+		mockContentManager.setup(c => c.loadContent()).returns(() => Promise.resolve(expectedNotebookContent));
+		defaultModelOptions.contentLoader = mockContentManager.object;
+
+		let model = new NotebookModel(defaultModelOptions, undefined, logService, undefined, new NullAdsTelemetryService(), queryConnectionService.object, configurationService, undoRedoService);
+		await model.loadContents();
+
+		let cell = model.addCell(CellTypes.Code);
+		assert.strictEqual(model.cells.length, expectedNotebookContent.cells.length + 1, 'New cell was not added to list of cells.');
+		assert.strictEqual(cell.language, expectedNotebookContent.metadata.language_info.name);
+	});
+
 	async function loadModelAndStartClientSession(notebookContent: nb.INotebookContents): Promise<NotebookModel> {
 		let mockContentManager = TypeMoq.Mock.ofType(NotebookEditorContentLoader);
 		mockContentManager.setup(c => c.loadContent()).returns(() => Promise.resolve(notebookContent));

src/sql/workbench/contrib/notebook/test/stubs.ts

@@ -122,7 +122,7 @@ export class NotebookModelStub implements INotebookModel {
 	getMetaValue(key: string) {
 		throw new Error('Method not implemented.');
 	}
-	addCell(cellType: CellType, index?: number): void {
+	addCell(cellType: CellType, index?: number, language?: string): void {
 		throw new Error('Method not implemented.');
 	}
 	moveCell(cellModel: ICellModel, direction: MoveDirection): void {

src/sql/workbench/services/notebook/browser/models/modelInterfaces.ts

@@ -389,7 +389,7 @@ export interface INotebookModel {
 	/**
 	 * Adds a cell to the index of the model
 	 */
-	addCell(cellType: CellType, index?: number): void;
+	addCell(cellType: CellType, index?: number, language?: string): void;
 
 	/**
 	 * Moves a cell up/down

src/sql/workbench/services/notebook/browser/models/notebookModel.ts

@@ -543,15 +543,15 @@ export class NotebookModel extends Disposable implements INotebookModel {
 		return this._cells.findIndex(cell => cell.equals(cellModel));
 	}
 
-	public addCell(cellType: CellType, index?: number): ICellModel | undefined {
+	public addCell(cellType: CellType, index?: number, language?: string): ICellModel | undefined {
 		if (this.inErrorState) {
 			return undefined;
 		}
-		let cell = this.createCell(cellType);
+		let cell = this.createCell(cellType, language);
 		return this.insertCell(cell, index, true);
 	}
 
-	public splitCell(cellType: CellType, notebookService: INotebookService, index?: number, addToUndoStack: boolean = true): ICellModel | undefined {
+	public splitCell(cellType: CellType, notebookService: INotebookService, index?: number, language?: string, addToUndoStack: boolean = true): ICellModel | undefined {
 		if (this.inErrorState) {
 			return undefined;
 		}
@@ -626,7 +626,7 @@ export class NotebookModel extends Disposable implements INotebookModel {
 					}
 					//If the selection is not from the start of the cell, create a new cell.
 					if (headContent.length) {
-						newCell = this.createCell(cellType);
+						newCell = this.createCell(cellType, language);
 						newCell.source = newSource;
 						newCellIndex++;
 						this.insertCell(newCell, newCellIndex, false);
@@ -639,7 +639,7 @@ export class NotebookModel extends Disposable implements INotebookModel {
 
 				if (tailCellContent.length) {
 					//tail cell will be of original cell type.
-					tailCell = this.createCell(this._cells[index].cellType);
+					tailCell = this.createCell(this._cells[index].cellType, language);
 					let tailSource = source.slice(tailRange.startLineNumber - 1) as string[];
 					if (selection.endColumn > 1) {
 						partialSource = source.slice(tailRange.startLineNumber - 1, tailRange.startLineNumber)[0].slice(tailRange.startColumn - 1);
@@ -833,13 +833,16 @@ export class NotebookModel extends Disposable implements INotebookModel {
 		}
 	}
 
-	private createCell(cellType: CellType): ICellModel {
+	private createCell(cellType: CellType, language?: string): ICellModel {
 		let singleCell: nb.ICellContents = {
 			cell_type: cellType,
 			source: '',
 			metadata: {},
 			execution_count: undefined
 		};
+		if (language) {
+			singleCell.metadata.language = language;
+		}
 		return this._notebookOptions.factory.createCell(singleCell, { notebook: this, isTrusted: true });
 	}