Always clear unused attachments when updating cell source (#17840)

* Always clear unused attachments when updating cell source

* comment

* Fix eslint and improve regex

* Add more tests

src/sql/workbench/contrib/notebook/browser/cellViews/textCell.component.ts

@@ -366,8 +366,6 @@ export class TextCellComponent extends CellView implements OnInit, OnChanges {
 	private updateCellSource(): void {
 		let textOutputElement = <HTMLElement>this.output.nativeElement;
 		let newCellSource: string = this._htmlMarkdownConverter.convert(textOutputElement.innerHTML);
-		// reset cell attachments to remove unused image data since we're going to go through each of them again
-		this.cellModel.attachments = {};
 		this.cellModel.source = newCellSource;
 		this._changeRef.detectChanges();
 	}

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

@@ -1234,6 +1234,97 @@ suite('Cell Model', function (): void {
 		assert.deepStrictEqual(cellModel.attachments, attachments, 'addAttachment should not add duplicate images');
 	});
 
+	test('deleting attachment from cell source should remove it from attachments should add a valid attachment to cell', async function () {
+		let imageFilebase64Value = 'data:application/octet-stream;base64,iVBORw0KGgoAAAANSU';
+		let index = imageFilebase64Value.indexOf('base64,');
+		const testImageAttachment: nb.ICellAttachment = { ['image/png']: imageFilebase64Value.substring(index + 7) };
+		let attachments: nb.ICellAttachments = { 'test.png': testImageAttachment };
+		let notebookModel = new NotebookModelStub({
+			name: '',
+			version: '',
+			mimetype: ''
+		});
+		let contents: nb.ICellContents = {
+			cell_type: CellTypes.Markdown,
+			source: '',
+			metadata: {}
+		};
+		let model = factory.createCell(contents, { notebook: notebookModel, isTrusted: false });
+		model.addAttachment('image/png', imageFilebase64Value, 'test.png');
+		assert.deepStrictEqual(model.attachments, attachments, 'attachment was not added initially');
+		model.source = '';
+		assert.deepStrictEqual(model.attachments, {}, 'attachments should be empty after clearing attachment from cell');
+	});
+
+	test('modifying cell source with existing attachment keeps attachment', async function () {
+		let imageFilebase64Value = 'data:application/octet-stream;base64,iVBORw0KGgoAAAANSU';
+		let index = imageFilebase64Value.indexOf('base64,');
+		const testImageAttachment: nb.ICellAttachment = { ['image/png']: imageFilebase64Value.substring(index + 7) };
+		let attachments: nb.ICellAttachments = { 'test.png': testImageAttachment };
+		let notebookModel = new NotebookModelStub({
+			name: '',
+			version: '',
+			mimetype: ''
+		});
+		let contents: nb.ICellContents = {
+			cell_type: CellTypes.Markdown,
+			source: '![image](attachment:test.png)',
+			metadata: {}
+		};
+		let model = factory.createCell(contents, { notebook: notebookModel, isTrusted: false });
+		model.addAttachment('image/png', imageFilebase64Value, 'test.png');
+		assert.deepStrictEqual(model.attachments, attachments, 'attachment was not added initially');
+		model.source = 'Some new text ' + model.source;
+		assert.deepStrictEqual(model.attachments, attachments, 'attachment should still exist after modifying cell source');
+	});
+
+	test('modifying cell source with multiple attachments keeps all attachments', async function () {
+		let imageFilebase64Value = 'data:application/octet-stream;base64,iVBORw0KGgoAAAANSU';
+		let index = imageFilebase64Value.indexOf('base64,');
+		const testImageAttachment: nb.ICellAttachment = { ['image/png']: imageFilebase64Value.substring(index + 7) };
+		let attachments: nb.ICellAttachments = { 'test.png': testImageAttachment, 'test2.png': testImageAttachment };
+		let notebookModel = new NotebookModelStub({
+			name: '',
+			version: '',
+			mimetype: ''
+		});
+		let contents: nb.ICellContents = {
+			cell_type: CellTypes.Markdown,
+			source: '![image](attachment:test.png)\n![image2](attachment:test2.png)',
+			metadata: {}
+		};
+		let model = factory.createCell(contents, { notebook: notebookModel, isTrusted: false });
+		model.addAttachment('image/png', imageFilebase64Value, 'test.png');
+		model.addAttachment('image/png', imageFilebase64Value, 'test2.png');
+		assert.deepStrictEqual(model.attachments, attachments, 'attachments were not added initially');
+		model.source = 'Some new text ' + model.source;
+		assert.deepStrictEqual(model.attachments, attachments, 'attachments should still exist after modifying cell source');
+	});
+
+	test('removing attachment keeps any others present in cell', async function () {
+		let imageFilebase64Value = 'data:application/octet-stream;base64,iVBORw0KGgoAAAANSU';
+		let index = imageFilebase64Value.indexOf('base64,');
+		const testImageAttachment: nb.ICellAttachment = { ['image/png']: imageFilebase64Value.substring(index + 7) };
+		let attachments: nb.ICellAttachments = { 'test.png': testImageAttachment, 'test2.png': testImageAttachment };
+		let notebookModel = new NotebookModelStub({
+			name: '',
+			version: '',
+			mimetype: ''
+		});
+		let contents: nb.ICellContents = {
+			cell_type: CellTypes.Markdown,
+			source: '![image](attachment:test.png)\n![image2](attachment:test2.png)',
+			metadata: {}
+		};
+		let model = factory.createCell(contents, { notebook: notebookModel, isTrusted: false });
+		model.addAttachment('image/png', imageFilebase64Value, 'test.png');
+		model.addAttachment('image/png', imageFilebase64Value, 'test2.png');
+		assert.deepStrictEqual(model.attachments, attachments, 'attachments were not added initially');
+		model.source = '![image2](attachment:test2.png)';
+		delete attachments['test.png'];
+		assert.deepStrictEqual(model.attachments, attachments, 'we should still have one attachment after removing the other');
+	});
+
 	test('cell should fire onCurrentEditModeChanged on edit', async function () {
 
 		let notebookModel = new NotebookModelStub({

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

@@ -307,6 +307,7 @@ export class CellModel extends Disposable implements ICellModel {
 	}
 
 	public set source(newSource: string | string[]) {
+		this.cleanUnusedAttachments(Array.isArray(newSource) ? newSource.join() : newSource);
 		newSource = this.attachImageFromSource(newSource);
 		newSource = this.getMultilineSource(newSource);
 		if (this._source !== newSource) {
@@ -329,6 +330,22 @@ 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;
+		this._attachments = {};
+		// Find existing attachments in the form ![...](attachment:...) so that we can make sure we keep those attachments
+		const attachmentRegex = /!\[.*?\]\(attachment:(.*?)\)/g;
+		let match;
+		while (match = attachmentRegex.exec(source)) { // eslint-disable-line no-cond-assign
+			this._attachments[match[1]] = originalAttachments[match[1]];
+		}
+	}
+
 	/**
 	 * Gets unique attachment name to add to cell metadata
 	 * @param imgName a string defining name of the image.