Fix split cell not working with attachments (#19587)

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

@@ -1079,6 +1079,96 @@ suite('Cell Model', function (): void {
 		assert.deepStrictEqual(serializedCell.attachments, undefined, 'JSON should not include attachments if attachments do not exist');
 	});
 
+	test('Should remove unused attachments name when updating cell source', async function () {
+		const cellAttachment = JSON.parse('{"ads.png":{"image/png":"iVBORw0KGgoAAAANSUhEUgAAAggg=="}}');
+		let notebookModel = new NotebookModelStub({
+			name: '',
+			version: '',
+			mimetype: ''
+		});
+		let contents: nb.ICellContents = {
+			cell_type: CellTypes.Markdown,
+			source: '',
+			attachments: cellAttachment
+		};
+		let model = factory.createCell(contents, { notebook: notebookModel, isTrusted: false });
+		assert.deepStrictEqual(model.attachments, contents.attachments, 'Attachments do not match in cellModel');
+
+		model.source = 'Test';
+
+		assert.notDeepStrictEqual(model.attachments, contents.attachments, 'Unused attachments are not removed after updating cell source');
+	});
+
+	test('Should not remove attachments that are still referenced in cell after updating the cell source', async function () {
+		const cellAttachment = JSON.parse('{"ads.png":{"image/png":"iVBORw0KGgoAAAANSUhEUgAAAggg=="}}');
+		let notebookModel = new NotebookModelStub({
+			name: '',
+			version: '',
+			mimetype: ''
+		});
+		let contents: nb.ICellContents = {
+			cell_type: CellTypes.Markdown,
+			source: '![ads.png](attachment:ads.png)',
+			attachments: cellAttachment
+		};
+		let model = factory.createCell(contents, { notebook: notebookModel, isTrusted: false });
+
+		model.source = '![ads.png](attachment:ads.png) \n Test';
+
+		assert.deepStrictEqual(model.attachments, contents.attachments, 'Attachments referenced in cell were removed');
+	});
+
+	test('Should update cell attachments', async function () {
+		const cellAttachment = JSON.parse('{"ads.png":{"image/png":"iVBORw0KGgoAAAANSUhEUgAAAggg=="}}');
+		let notebookModel = new NotebookModelStub({
+			name: '',
+			version: '',
+			mimetype: ''
+		});
+		let contents: nb.ICellContents = {
+			cell_type: CellTypes.Markdown,
+			source: '![ads.png](attachment:ads.png)'
+		};
+		let model = factory.createCell(contents, { notebook: notebookModel, isTrusted: false });
+		model.updateAttachmentsFromSource('![ads.png](attachment:ads.png)', cellAttachment);
+
+		assert.deepStrictEqual(model.attachments, cellAttachment, 'Cell attachments are not updated correctly');
+	});
+
+	test('Should not update cell attachments if they are not referenced in cell source', async function () {
+		const cellAttachment = JSON.parse('{"ads.png":{"image/png":"iVBORw0KGgoAAAANSUhEUgAAAggg=="}}');
+		let notebookModel = new NotebookModelStub({
+			name: '',
+			version: '',
+			mimetype: ''
+		});
+		let contents: nb.ICellContents = {
+			cell_type: CellTypes.Markdown,
+			source: '![ads.png](attachment:ads.png)'
+		};
+		let model = factory.createCell(contents, { notebook: notebookModel, isTrusted: false });
+		model.updateAttachmentsFromSource('test', cellAttachment);
+
+		assert.notDeepStrictEqual(model.attachments, cellAttachment, 'Cell attachments are updated when they are not referenced in cell source');
+	});
+
+	test('Should not update cell attachments if the attachment reference is not in the correct format', async function () {
+		const cellAttachment = JSON.parse('{"ads.png":{"image/png":"iVBORw0KGgoAAAANSUhEUgAAAggg=="}}');
+		let notebookModel = new NotebookModelStub({
+			name: '',
+			version: '',
+			mimetype: ''
+		});
+		let contents: nb.ICellContents = {
+			cell_type: CellTypes.Markdown,
+			source: '![ads.png](attachment:ads.png)'
+		};
+		let model = factory.createCell(contents, { notebook: notebookModel, isTrusted: false });
+		model.updateAttachmentsFromSource('ads.png', cellAttachment);
+
+		assert.notDeepStrictEqual(model.attachments, cellAttachment, 'Cell attachments are updated when the reference is not in the correct format');
+	});
+
 	test('Should not have cache chart data after new cell created', async function () {
 		let notebookModel = new NotebookModelStub({
 			name: '',

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

@@ -334,7 +334,7 @@ export class CellModel extends Disposable implements ICellModel {
 	}
 
 	public set source(newSource: string | string[]) {
-		this.cleanUnusedAttachments(Array.isArray(newSource) ? newSource.join() : newSource);
+		this.updateAttachmentsFromSource(Array.isArray(newSource) ? newSource.join() : newSource);
 		newSource = this.attachImageFromSource(newSource);
 		newSource = this.getMultilineSource(newSource);
 		if (this._source !== newSource) {
@@ -358,12 +358,8 @@ export class CellModel extends Disposable implements ICellModel {
 		return newSource;
 	}
 
-	/**
-	 * Cleans up the attachments, removing any ones that aren't being currently used in the specified source string.
-	 * @param source The new source string to check for attachments being used
-	 */
-	private cleanUnusedAttachments(source: string): void {
-		const originalAttachments = this._attachments;
+	public updateAttachmentsFromSource(source: string, attachments?: nb.ICellAttachments): void {
+		const originalAttachments = attachments ? attachments : this._attachments;
 		this._attachments = {};
 		// Find existing attachments in the form ![...](attachment:...) so that we can make sure we keep those attachments
 		const attachmentRegex = /!\[.*?\]\(attachment:(.*?)\)/g;

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

@@ -554,6 +554,13 @@ export interface ICellModel {
 	 * Returns the name of the attachment added to metadata.
 	 */
 	addAttachment(mimeType: string, base64Encoding: string, name: string): string;
+	/**
+	 * Updates the current cell attachments with the attachments provided.
+	 * If no attachments are passed in then it cleans up the current cell attachments and removes any ones that aren't being currently used in the specified source string.
+	 * @param source The new source string to check for attachments being used
+	 * @param attachments (Optional) The new attachments for the cell
+	 */
+	updateAttachmentsFromSource(source: string, attachments?: nb.ICellAttachments): void;
 	richTextCursorPosition: ICaretPosition;
 	markdownCursorPosition: IPosition;
 	/**

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

@@ -570,10 +570,11 @@ export class NotebookModel extends Disposable implements INotebookModel {
 				let range = model.getFullModelRange();
 				let selection = editorControl.getSelection();
 				let source = this.cells[index].source;
-				let newCell = undefined, tailCell = undefined, partialSource = undefined;
+				let newCell: ICellModel = undefined, tailCell: ICellModel = undefined, partialSource = undefined;
 				let newCellIndex = index;
 				let tailCellIndex = index;
 				let splitCells: SplitCell[] = [];
+				let attachments = {};
 
 				// Save UI state
 				let showMarkdown = this.cells[index].showMarkdown;
@@ -607,6 +608,9 @@ export class NotebookModel extends Disposable implements INotebookModel {
 						partialSource = source.slice(selection.startLineNumber - 1, selection.startLineNumber)[0].slice(0, selection.startColumn - 1);
 						headsource = headsource.concat(partialSource.toString());
 					}
+					// Save attachments before updating cell contents
+					attachments = this.cells[index].attachments;
+					// No need to update attachments, since unused attachments are removed when updating the cell source
 					this.cells[index].source = headsource;
 					splitCells.push({ cell: this.cells[index], prefix: undefined });
 				}
@@ -628,6 +632,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, language);
+						newCell.updateAttachmentsFromSource(newSource.join(), attachments);
 						newCell.source = newSource;
 						newCellIndex++;
 						this.insertCell(newCell, newCellIndex, false);
@@ -651,6 +656,7 @@ export class NotebookModel extends Disposable implements INotebookModel {
 					if (tailSource[0] === '\r\n' || tailSource[0] === '\n') {
 						newlinesBeforeTailCellContent = tailSource.splice(0, 1)[0];
 					}
+					tailCell.updateAttachmentsFromSource(tailSource.join(), attachments);
 					tailCell.source = tailSource;
 					tailCellIndex = newCellIndex + 1;
 					this.insertCell(tailCell, tailCellIndex, false);
@@ -684,6 +690,7 @@ export class NotebookModel extends Disposable implements INotebookModel {
 		let firstCell = cells[0].cell;
 		// Append the other cell sources to the first cell
 		for (let i = 1; i < cells.length; i++) {
+			firstCell.attachments = { ...firstCell.attachments, ...cells[i].cell.attachments };
 			firstCell.source = cells[i].prefix ? [...firstCell.source, ...cells[i].prefix, ...cells[i].cell.source] : [...firstCell.source, ...cells[i].cell.source];
 		}
 		// Set newly created cell as active cell