Clear execution state when restarting kernel

src/kernels/execution/cellExecution.ts

@@ -46,7 +46,7 @@ export class CellExecutionFactory {
         cell: NotebookCell,
         code: string | undefined,
         metadata: Readonly<KernelConnectionMetadata>,
-        info?: ResumeCellExecutionInformation
+        info?: ResumeCellExecutionInformation | 'preserveOutput'
     ) {
         // eslint-disable-next-line @typescript-eslint/no-use-before-define
         return CellExecution.fromCell(cell, code, metadata, this.controller, this.requestListener, info);
@@ -97,7 +97,7 @@ export class CellExecution implements ICellExecution, IDisposable {
         private readonly kernelConnection: Readonly<KernelConnectionMetadata>,
         private readonly controller: IKernelController,
         private readonly requestListener: CellExecutionMessageHandlerService,
-        private readonly resumeExecution?: ResumeCellExecutionInformation
+        private readonly resumeExecution?: ResumeCellExecutionInformation | 'preserveOutput'
     ) {
         workspace.onDidCloseTextDocument(
             (e) => {
@@ -126,7 +126,7 @@ export class CellExecution implements ICellExecution, IDisposable {
             this.execution = CellExecutionCreator.getOrCreate(
                 cell,
                 this.controller,
-                resumeExecution?.msg_id ? false : true // Do not clear output if we're resuming execution of a cell.
+                resumeExecution === 'preserveOutput' || resumeExecution?.msg_id ? false : true // Do not clear output if we're resuming execution of a cell.
             );
             NotebookCellStateTracker.setCellState(cell, NotebookCellExecutionState.Pending);
         } else {
@@ -139,19 +139,26 @@ export class CellExecution implements ICellExecution, IDisposable {
         }
     }
 
+    public clear() {
+        if (this.canExecuteCell()) {
+            const execution = CellExecutionCreator.getOrCreate(this.cell, this.controller, false);
+            execution.start();
+            execution.end(undefined);
+        }
+    }
     public static fromCell(
         cell: NotebookCell,
         code: string | undefined,
         metadata: Readonly<KernelConnectionMetadata>,
         controller: IKernelController,
         requestListener: CellExecutionMessageHandlerService,
-        info?: ResumeCellExecutionInformation
+        info?: ResumeCellExecutionInformation | 'preserveOutput'
     ) {
         return new CellExecution(cell, code, metadata, controller, requestListener, info);
     }
     public async start(session: IKernelSession) {
         this.session = session;
-        if (this.resumeExecution?.msg_id) {
+        if (this.resumeExecution && typeof this.resumeExecution !== 'string' && this.resumeExecution?.msg_id) {
             return this.resume(session, this.resumeExecution);
         }
         if (this.cancelHandled) {
@@ -399,7 +406,11 @@ export class CellExecution implements ICellExecution, IDisposable {
         // Skip if no code to execute
         if (code.trim().length === 0 || this.cell.document.isClosed) {
             if (code.trim().length === 0) {
-                this.execution?.start(this.resumeExecution?.startTime);
+                this.execution?.start(
+                    this.resumeExecution && typeof this.resumeExecution !== 'string'
+                        ? this.resumeExecution?.startTime
+                        : undefined
+                );
                 this.execution?.clearOutput()?.then(noop, noop);
             }
             traceCellMessage(this.cell, 'Empty cell execution');

src/kernels/kernelExecution.ts

@@ -154,6 +154,20 @@ export class NotebookKernelExecution implements INotebookKernelExecution {
         );
         logger.trace(`Cell ${cell.index} executed ${success ? 'successfully' : 'with an error'}`);
     }
+    public clearState(cell: NotebookCell): void {
+        const execution = this.documentExecutions.get(cell.notebook);
+        const exec = execution?.queue.includes(cell)
+            ? undefined
+            : this.executionFactory.create(cell, undefined, this.kernel.kernelConnectionMetadata, 'preserveOutput');
+        if (exec?.cell.executionSummary?.timing && typeof exec.cell.executionSummary.executionOrder === 'number') {
+            try {
+                exec.clear();
+            } finally {
+                exec.dispose();
+            }
+        }
+    }
+
     public async executeCell(cell: NotebookCell, codeOverride?: string | undefined): Promise<void> {
         traceCellMessage(cell, `NotebookKernelExecution.executeCell (1), ${getDisplayPath(cell.notebook.uri)}`);
         const stopWatch = new StopWatch();

src/kernels/types.ts

@@ -465,6 +465,12 @@ export interface INotebookKernelExecution {
      * @param codeOverride Override the code to execute
      */
     executeCell(cell: NotebookCell, codeOverride?: string): Promise<void>;
+    /**
+     * Clears the execution state of a cell.
+     * This will clear the green check or red cross again cell along with the execution time.
+     * These are populated to indicate whether a cell was executed.
+     */
+    clearState(cell: NotebookCell): void;
     /**
      * Executes 3rd party code against the kernel.
      */

src/notebooks/notebookCommandListener.ts

@@ -31,6 +31,8 @@ import { IDataScienceErrorHandler } from '../kernels/errors/types';
 import { getNotebookMetadata } from '../platform/common/utils';
 import { KernelConnector } from './controllers/kernelConnector';
 import { IControllerRegistration } from './controllers/types';
+import { DisposableStore } from '../platform/common/utils/lifecycle';
+import { createDeferred } from '../platform/common/utils/async';
 
 /**
  * Registers commands specific to the notebook UI
@@ -170,21 +172,43 @@ export class NotebookCommandListener implements IDataScienceCommandListener {
 
         if (kernel) {
             logger.debug(`Restart kernel command handler for ${getDisplayPath(document.uri)}`);
-            if (await this.shouldAskForRestart(document.uri)) {
-                // Ask the user if they want us to restart or not.
-                const message = DataScience.restartKernelMessage;
-                const yes = DataScience.restartKernelMessageYes;
-                const dontAskAgain = DataScience.restartKernelMessageDontAskAgain;
-
-                const response = await window.showInformationMessage(message, { modal: true }, yes, dontAskAgain);
-                if (response === dontAskAgain) {
-                    await this.disableAskForRestart(document.uri);
-                    this.wrapKernelMethod('restart', kernel).catch(noop);
-                } else if (response === yes) {
-                    this.wrapKernelMethod('restart', kernel).catch(noop);
+            const didRestartKernel = createDeferred<void>();
+            const disposable = new DisposableStore();
+            disposable.add(kernel.onRestarted(() => didRestartKernel.resolve()));
+            const resetExecutionOfCells = async () => {
+                await didRestartKernel.promise;
+                const controller = this.controllerRegistration.getSelected(document);
+                const kernel = this.kernelProvider.get(document);
+                const execution = kernel ? this.kernelProvider.getKernelExecution(kernel) : undefined;
+                if (controller && execution) {
+                    const pendingCells = new Set(execution.pendingCells);
+                    document
+                        .getCells()
+                        .filter((c) => c.kind === NotebookCellKind.Code && !pendingCells.has(c))
+                        .forEach((cell) => execution.clearState(cell));
+                }
+            };
+            void resetExecutionOfCells();
+            try {
+                if (await this.shouldAskForRestart(document.uri)) {
+                    // Ask the user if they want us to restart or not.
+                    const message = DataScience.restartKernelMessage;
+                    const yes = DataScience.restartKernelMessageYes;
+                    const dontAskAgain = DataScience.restartKernelMessageDontAskAgain;
+                    const response = await window.showInformationMessage(message, { modal: true }, yes, dontAskAgain);
+                    if (response === dontAskAgain) {
+                        await this.disableAskForRestart(document.uri);
+                        await this.wrapKernelMethod('restart', kernel).catch(noop);
+                    } else if (response === yes) {
+                        await this.wrapKernelMethod('restart', kernel).catch(noop);
+                    }
+                } else {
+                    await this.wrapKernelMethod('restart', kernel).catch(noop);
                 }
-            } else {
-                this.wrapKernelMethod('restart', kernel).catch(noop);
+            } finally {
+                // If we
+                disposable.dispose();
+                didRestartKernel.reject();
             }
         }
     }

src/standalone/chat/installPackageTool.node.ts

@@ -98,7 +98,9 @@ export class InstallPackagesTool implements vscode.LanguageModelTool<IInstallPac
         }
 
         return new vscode.LanguageModelToolResult([
-            new vscode.LanguageModelTextPart('Installation finished successfully.')
+            new vscode.LanguageModelTextPart(
+                `Installation finished successfully. Use the 'copilot_getNotebookSummary' to get the latest summary (state) of the notebook.`
+            )
         ]);
     }