diff --git a/extensions/vscode-api-tests/src/singlefolder-tests/notebook.test.ts b/extensions/vscode-api-tests/src/singlefolder-tests/notebook.test.ts
index 33cad6edb97..f863bf014b6 100644
--- a/extensions/vscode-api-tests/src/singlefolder-tests/notebook.test.ts
+++ b/extensions/vscode-api-tests/src/singlefolder-tests/notebook.test.ts
@@ -263,14 +263,25 @@ const apiTestContentProvider: vscode.NotebookContentProvider = {
 		const editor = await vscode.window.showNotebookDocument(notebook);
 
 		const notebookChangeEvent = asPromise<vscode.NotebookDocumentChangeEvent>(vscode.workspace.onDidChangeNotebookDocument);
-		const version = editor.document.version;
-		await editor.edit(editBuilder => {
-			editBuilder.replaceCells(1, 0, [{ kind: vscode.NotebookCellKind.Code, languageId: 'javascript', value: 'test 2', outputs: [], metadata: undefined }]);
-			editBuilder.replaceCellMetadata(0, { inputCollapsed: false });
-		});
-
+		const version = editor.notebook.version;
+		const edit = new vscode.WorkspaceEdit();
+		const cellEdit = vscode.NotebookEdit.replaceCells(new vscode.NotebookRange(1, 0), [{ kind: vscode.NotebookCellKind.Code, languageId: 'javascript', value: 'test 2', outputs: [], metadata: undefined }]);
+		const metdataEdit = vscode.NotebookEdit.updateNotebookMetadata({ ...notebook.metadata, custom: { ...(notebook.metadata.custom || {}), extraNotebookMetadata: true } });
+		edit.set(notebook.uri, [cellEdit, metdataEdit]);
+		await vscode.workspace.applyEdit(edit);
 		await notebookChangeEvent;
-		assert.strictEqual(version + 1, editor.document.version);
+
+		const notebookChangeEvent2 = asPromise<vscode.NotebookDocumentChangeEvent>(vscode.workspace.onDidChangeNotebookDocument);
+		const edit2 = new vscode.WorkspaceEdit();
+		const cellMetadataEdit = vscode.NotebookEdit.updateCellMetadata(0, { extraCellMetadata: true });
+		edit2.set(notebook.uri, [cellMetadataEdit]);
+		await vscode.workspace.applyEdit(edit2);
+		await notebookChangeEvent2;
+
+		assert.strictEqual(version + 2, editor.notebook.version);
+		const cell = editor.notebook.cellAt(0);
+		assert.ok(editor.notebook.metadata.custom.extraNotebookMetadata, `Test metadata not found`);
+		assert.ok(cell.metadata.extraCellMetadata, `Test cell metdata not found`);
 	});
 
 	test('edit API batch edits undo/redo', async function () {
diff --git a/src/vs/workbench/api/common/extHostTypes.ts b/src/vs/workbench/api/common/extHostTypes.ts
index a1c2a3ea107..b5e2cd04ef9 100644
--- a/src/vs/workbench/api/common/extHostTypes.ts
+++ b/src/vs/workbench/api/common/extHostTypes.ts
@@ -807,6 +807,8 @@ export class WorkspaceEdit implements vscode.WorkspaceEdit {
 					if (NotebookEdit.isNotebookCellEdit(edit)) {
 						if (edit.newCellMetadata) {
 							this.replaceNotebookCellMetadata(uri, edit.range.start, edit.newCellMetadata);
+						} else if (edit.newNotebookMetadata) {
+							this.replaceNotebookMetadata(uri, edit.newNotebookMetadata);
 						} else {
 							this.replaceNotebookCells(uri, edit.range, edit.newCells);
 						}