Preserve code in cells during execution and edits (#10981)

* Preserve code in cells during execution and edits
* Oops

Author: DonJayamanne

news/2 Fixes/10949.md

@@ -0,0 +1 @@
+Ensure user code in cell is preserved between cell execution and cell edits.

src/client/datascience/interactive-common/interactiveBase.ts

@@ -669,7 +669,7 @@ export abstract class InteractiveBase extends WebViewHost<IInteractiveWindowMapp
 
                 case CellState.executing:
                     // Tell the react controls we have an update
-                    this.postMessage(InteractiveWindowMessages.UpdateCell, cell).ignoreErrors();
+                    this.postMessage(InteractiveWindowMessages.UpdateCellWithExecutionResults, cell).ignoreErrors();
                     break;
 
                 case CellState.error:

src/client/datascience/interactive-common/interactiveWindowTypes.ts

@@ -30,7 +30,7 @@ import { BaseReduxActionPayload } from './types';
 export enum InteractiveWindowMessages {
     StartCell = 'start_cell',
     FinishCell = 'finish_cell',
-    UpdateCell = 'update_cell',
+    UpdateCellWithExecutionResults = 'UpdateCellWithExecutionResults',
     GotoCodeCell = 'gotocell_code',
     CopyCodeCell = 'copycell_code',
     NotebookExecutionActivated = 'notebook_execution_activated',
@@ -497,7 +497,7 @@ export class IInteractiveWindowMapping {
     public [IPyWidgetMessages.IPyWidgets_mirror_execute]: { id: string; msg: KernelMessage.IExecuteRequestMsg };
     public [InteractiveWindowMessages.StartCell]: ICell;
     public [InteractiveWindowMessages.FinishCell]: ICell;
-    public [InteractiveWindowMessages.UpdateCell]: ICell;
+    public [InteractiveWindowMessages.UpdateCellWithExecutionResults]: ICell;
     public [InteractiveWindowMessages.GotoCodeCell]: IGotoCode;
     public [InteractiveWindowMessages.CopyCodeCell]: ICopyCode;
     public [InteractiveWindowMessages.NotebookExecutionActivated]: string;

src/client/datascience/interactive-common/synchronization.ts

@@ -166,7 +166,8 @@ const messageWithMessageTypes: MessageMapping<IInteractiveWindowMapping> & Messa
     [InteractiveWindowMessages.Sync]: MessageType.other,
     [InteractiveWindowMessages.Undo]: MessageType.other,
     [InteractiveWindowMessages.UnfocusedCellEditor]: MessageType.syncWithLiveShare,
-    [InteractiveWindowMessages.UpdateCell]: MessageType.syncAcrossSameNotebooks | MessageType.syncWithLiveShare,
+    [InteractiveWindowMessages.UpdateCellWithExecutionResults]:
+        MessageType.syncAcrossSameNotebooks | MessageType.syncWithLiveShare,
     [InteractiveWindowMessages.UpdateModel]: MessageType.syncAcrossSameNotebooks | MessageType.syncWithLiveShare,
     [InteractiveWindowMessages.UpdateKernel]: MessageType.syncAcrossSameNotebooks | MessageType.syncWithLiveShare,
     [InteractiveWindowMessages.UpdateDisplayData]: MessageType.syncWithLiveShare,

src/client/datascience/interactive-ipynb/nativeEditor.ts

@@ -449,6 +449,20 @@ export class NativeEditor extends InteractiveBase implements INotebookEditor {
         const modified = filtered.filter((c) => c.state === CellState.finished || c.state === CellState.error);
         const unmodified = this.model?.cells.filter((c) => modified.find((m) => m.id === c.id));
         if (modified.length > 0 && unmodified && this.model) {
+            // As this point, we're updating the model because of changes to the cell as a result of execution.
+            // The output and execution count change, however we're just going to update everything.
+            // But, we should not update the `source`. The only time source can change is when a request comes from the UI.
+            // Perhaps we need a finer grained update (update only output and execution count along with `source=execution`).
+            // For now retain source from previous model.
+            // E.g. user executes a cell, in the mean time they update the text. Now model contains new value.
+            // However once execution has completed, this code will update the model with results from previous execution (prior to edit).
+            // We now need to give preference to the text in the model, over the old one that was executed.
+            modified.forEach((modifiedCell) => {
+                const originalCell = unmodified.find((unmodifiedCell) => unmodifiedCell.id === modifiedCell.id);
+                if (originalCell) {
+                    modifiedCell.data.source = originalCell.data.source;
+                }
+            });
             this.model.update({
                 source: 'user',
                 kind: 'modify',

src/datascience-ui/history-react/redux/reducers/index.ts

@@ -47,7 +47,7 @@ export const reducerMap: Partial<IInteractiveActionMapping> = {
     [InteractiveWindowMessages.Redo]: Execution.redo,
     [InteractiveWindowMessages.StartCell]: Creation.startCell,
     [InteractiveWindowMessages.FinishCell]: Creation.finishCell,
-    [InteractiveWindowMessages.UpdateCell]: Creation.updateCell,
+    [InteractiveWindowMessages.UpdateCellWithExecutionResults]: Creation.updateCell,
     [InteractiveWindowMessages.Activate]: CommonEffects.activate,
     [InteractiveWindowMessages.RestartKernel]: Kernel.handleRestarted,
     [CssMessages.GetCssResponse]: CommonEffects.handleCss,

src/datascience-ui/interactive-common/redux/reducers/helpers.ts

@@ -83,7 +83,7 @@ export namespace Helpers {
 
             // Prevent updates to the source, as its possible we have recieved a response for a cell execution
             // and the user has updated the cell text since then.
-            const newVM = {
+            const newVM: ICellViewModel = {
                 ...newVMs[index],
                 hasBeenRun: true,
                 cell: {
@@ -95,7 +95,7 @@ export namespace Helpers {
                     }
                 }
             };
-            newVMs[index] = asCellViewModel(newVM);
+            newVMs[index] = newVM;
 
             return {
                 ...arg.prevState,

src/datascience-ui/native-editor/redux/reducers/index.ts

@@ -64,7 +64,7 @@ export const reducerMap: Partial<INativeEditorActionMapping> = {
     // Messages from the webview (some are ignored)
     [InteractiveWindowMessages.StartCell]: Creation.startCell,
     [InteractiveWindowMessages.FinishCell]: Creation.finishCell,
-    [InteractiveWindowMessages.UpdateCell]: Creation.updateCell,
+    [InteractiveWindowMessages.UpdateCellWithExecutionResults]: Creation.updateCell,
     [InteractiveWindowMessages.NotebookDirty]: CommonEffects.notebookDirty,
     [InteractiveWindowMessages.NotebookClean]: CommonEffects.notebookClean,
     [InteractiveWindowMessages.LoadAllCells]: Creation.loadAllCells,