diff --git a/src/kernels/execution/cellExecution.ts b/src/kernels/execution/cellExecution.ts index f1710e1a5e0..a463af4fa03 100644 --- a/src/kernels/execution/cellExecution.ts +++ b/src/kernels/execution/cellExecution.ts @@ -46,7 +46,7 @@ export class CellExecutionFactory { cell: NotebookCell, code: string | undefined, metadata: Readonly, - 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, 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, 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'); diff --git a/src/kernels/kernelExecution.ts b/src/kernels/kernelExecution.ts index fcf9a8c25d0..1965e13c092 100644 --- a/src/kernels/kernelExecution.ts +++ b/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 { traceCellMessage(cell, `NotebookKernelExecution.executeCell (1), ${getDisplayPath(cell.notebook.uri)}`); const stopWatch = new StopWatch(); diff --git a/src/kernels/types.ts b/src/kernels/types.ts index 1d51d627afe..2a05a1ba69c 100644 --- a/src/kernels/types.ts +++ b/src/kernels/types.ts @@ -465,6 +465,12 @@ export interface INotebookKernelExecution { * @param codeOverride Override the code to execute */ executeCell(cell: NotebookCell, codeOverride?: string): Promise; + /** + * 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. */ diff --git a/src/notebooks/notebookCommandListener.ts b/src/notebooks/notebookCommandListener.ts index b95b1475c8c..5e18ba3efbf 100644 --- a/src/notebooks/notebookCommandListener.ts +++ b/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(); + 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(); } } } diff --git a/src/standalone/chat/installPackageTool.node.ts b/src/standalone/chat/installPackageTool.node.ts index 0f9a1b4727c..d2296f6bb8e 100644 --- a/src/standalone/chat/installPackageTool.node.ts +++ b/src/standalone/chat/installPackageTool.node.ts @@ -98,7 +98,9 @@ export class InstallPackagesTool implements vscode.LanguageModelTool