Fix updating outputs of another cell using a NotebookCellExecution

extensions/vscode-api-tests/src/singlefolder-tests/notebook.test.ts

@@ -927,6 +927,80 @@ const apiTestContentProvider: vscode.NotebookContentProvider = {
 
 		assert.strictEqual(executionWasCancelled, true);
 	});
+
+	test('appendOutput to different cell', async function () {
+		const notebook = await openRandomNotebookDocument();
+		const editor = await vscode.window.showNotebookDocument(notebook);
+		const cell0 = editor.notebook.cellAt(0);
+		await vscode.commands.executeCommand('notebook.cell.insertCodeCellBelow');
+		const cell1 = editor.notebook.cellAt(1);
+
+		const nextCellKernel = new class extends Kernel {
+			constructor() {
+				super('nextCellKernel', 'Append to cell kernel');
+			}
+
+			override async _runCell(cell: vscode.NotebookCell) {
+				const task = this.controller.createNotebookCellExecution(cell);
+				task.start();
+				await task.appendOutput([new vscode.NotebookCellOutput([
+					vscode.NotebookCellOutputItem.text('my output')
+				])], cell1);
+				await task.appendOutput([new vscode.NotebookCellOutput([
+					vscode.NotebookCellOutputItem.text('my output 2')
+				])], cell1);
+				task.end(true);
+			}
+		};
+		testDisposables.push(nextCellKernel.controller);
+
+		await withEvent<vscode.NotebookDocumentChangeEvent>(vscode.workspace.onDidChangeNotebookDocument, async (event) => {
+			await assertKernel(nextCellKernel, notebook);
+			await vscode.commands.executeCommand('notebook.cell.execute');
+			await event;
+			assert.strictEqual(cell0.outputs.length, 0, 'should not change cell 0');
+			assert.strictEqual(cell1.outputs.length, 2, 'should update cell 1');
+			assert.strictEqual(cell1.outputs[0].items.length, 1);
+			assert.deepStrictEqual(new TextDecoder().decode(cell1.outputs[0].items[0].data), 'my output');
+		});
+	});
+
+	test('replaceOutput to different cell', async function () {
+		const notebook = await openRandomNotebookDocument();
+		const editor = await vscode.window.showNotebookDocument(notebook);
+		const cell0 = editor.notebook.cellAt(0);
+		await vscode.commands.executeCommand('notebook.cell.insertCodeCellBelow');
+		const cell1 = editor.notebook.cellAt(1);
+
+		const nextCellKernel = new class extends Kernel {
+			constructor() {
+				super('nextCellKernel', 'Replace to cell kernel');
+			}
+
+			override async _runCell(cell: vscode.NotebookCell) {
+				const task = this.controller.createNotebookCellExecution(cell);
+				task.start();
+				await task.replaceOutput([new vscode.NotebookCellOutput([
+					vscode.NotebookCellOutputItem.text('my output')
+				])], cell1);
+				await task.replaceOutput([new vscode.NotebookCellOutput([
+					vscode.NotebookCellOutputItem.text('my output 2')
+				])], cell1);
+				task.end(true);
+			}
+		};
+		testDisposables.push(nextCellKernel.controller);
+
+		await withEvent<vscode.NotebookDocumentChangeEvent>(vscode.workspace.onDidChangeNotebookDocument, async (event) => {
+			await assertKernel(nextCellKernel, notebook);
+			await vscode.commands.executeCommand('notebook.cell.execute');
+			await event;
+			assert.strictEqual(cell0.outputs.length, 0, 'should not change cell 0');
+			assert.strictEqual(cell1.outputs.length, 1, 'should update cell 1');
+			assert.strictEqual(cell1.outputs[0].items.length, 1);
+			assert.deepStrictEqual(new TextDecoder().decode(cell1.outputs[0].items[0].data), 'my output 2');
+		});
+	});
 });
 
 (vscode.env.uiKind === vscode.UIKind.Web ? suite.skip : suite)('statusbar', () => {

src/vs/workbench/api/browser/mainThreadNotebookDto.ts

@@ -96,6 +96,7 @@ export namespace NotebookDto {
 		if (data.editType === CellExecutionUpdateType.Output) {
 			return {
 				editType: data.editType,
+				cellHandle: data.cellHandle,
 				append: data.append,
 				outputs: data.outputs.map(fromNotebookOutputDto)
 			};

src/vs/workbench/api/common/extHost.protocol.ts

@@ -999,6 +999,7 @@ export interface INotebookProxyKernelDto {
 
 export interface ICellExecuteOutputEditDto {
 	editType: CellExecutionUpdateType.Output;
+	cellHandle: number;
 	append?: boolean;
 	outputs: NotebookOutputDto[];
 }

src/vs/workbench/api/common/extHostNotebookKernels.ts

@@ -441,6 +441,17 @@ class NotebookCellExecutionTask extends Disposable {
 		}
 	}
 
+	private cellIndexToHandle(cellOrCellIndex: vscode.NotebookCell | undefined): number {
+		let cell: ExtHostCell | undefined = this._cell;
+		if (cellOrCellIndex) {
+			cell = this._cell.notebook.getCellFromApiCell(cellOrCellIndex);
+		}
+		if (!cell) {
+			throw new Error('INVALID cell');
+		}
+		return cell.handle;
+	}
+
 	private validateAndConvertOutputs(items: vscode.NotebookCellOutput[]): NotebookOutputDto[] {
 		return items.map(output => {
 			const newOutput = NotebookCellOutput.ensureUniqueMimeTypes(output.items, true);
@@ -456,10 +467,12 @@ class NotebookCellExecutionTask extends Disposable {
 	}
 
 	private async updateOutputs(outputs: vscode.NotebookCellOutput | vscode.NotebookCellOutput[], cell: vscode.NotebookCell | undefined, append: boolean): Promise<void> {
+		const handle = this.cellIndexToHandle(cell);
 		const outputDtos = this.validateAndConvertOutputs(asArray(outputs));
 		return this.updateSoon(
 			{
 				editType: CellExecutionUpdateType.Output,
+				cellHandle: handle,
 				append,
 				outputs: outputDtos
 			});

src/vs/workbench/contrib/notebook/browser/notebookEditorWidget.ts

@@ -1459,11 +1459,11 @@ export class NotebookEditorWidget extends Disposable implements INotebookEditorD
 			}
 			hasPendingChangeContentHeight = true;
 
-			DOM.scheduleAtNextAnimationFrame(() => {
+			this._localStore.add(DOM.scheduleAtNextAnimationFrame(() => {
 				hasPendingChangeContentHeight = false;
 				this._updateScrollHeight();
 				this._onDidChangeContentHeight.fire(this._list.getScrollHeight());
-			}, 100);
+			}, 100));
 		}));
 
 		this._localStore.add(this._list.onDidRemoveOutputs(outputs => {

src/vs/workbench/contrib/notebook/browser/notebookExecutionStateServiceImpl.ts

@@ -217,7 +217,7 @@ function updateToEdit(update: ICellExecuteUpdate, cellHandle: number): ICellEdit
 	if (update.editType === CellExecutionUpdateType.Output) {
 		return {
 			editType: CellEditType.Output,
-			handle: cellHandle,
+			handle: update.cellHandle,
 			append: update.append,
 			outputs: update.outputs,
 		};

src/vs/workbench/contrib/notebook/common/notebookExecutionService.ts

@@ -15,6 +15,7 @@ export enum CellExecutionUpdateType {
 
 export interface ICellExecuteOutputEdit {
 	editType: CellExecutionUpdateType.Output;
+	cellHandle: number;
 	append?: boolean;
 	outputs: IOutputDto[];
 }