Preserve cell edit state and focus after dragging

#96308

src/vs/workbench/contrib/notebook/browser/view/renderers/cellRenderer.ts

@@ -553,7 +553,7 @@ export class CellDragAndDropController {
 
 			const location = this.getDropInsertLocation(templateData, event);
 			if (location) {
-				this.notebookEditor.moveCell(draggedCell, templateData.currentRenderedCell!, location);
+				this.moveCell(draggedCell, templateData.currentRenderedCell!, location);
 				container.classList.remove(DRAGOVER_TOP_CLASS, DRAGOVER_BOTTOM_CLASS);
 			}
 		}));
@@ -565,6 +565,12 @@ export class CellDragAndDropController {
 		}));
 	}
 
+	private moveCell(draggedCell: ICellViewModel, ontoCell: ICellViewModel, location: 'above' | 'below') {
+		const editState = draggedCell.editState;
+		this.notebookEditor.moveCell(draggedCell, ontoCell, location);
+		this.notebookEditor.focusNotebookCell(draggedCell, editState === CellEditState.Editing);
+	}
+
 	private getDropInsertLocation(templateData: BaseCellRenderTemplate, event: DragEvent): 'above' | 'below' | undefined {
 		if (templateData.currentRenderedCell === this.currentDraggedCell) {
 			return;