Fix rendering of output of cells added with already existing output

packages/plugin-ext/src/main/browser/notebooks/renderers/cell-output-webview.tsx

@@ -32,12 +32,15 @@ import { WebviewWidget } from '../../webview/webview';
 import { Message, WidgetManager } from '@theia/core/lib/browser';
 import { outputWebviewPreload, PreloadContext } from './output-webview-internal';
 import { WorkspaceTrustService } from '@theia/workspace/lib/browser';
-import { CellsChangedMessage, CellsMoved, CellsSpliced, ChangePreferredMimetypeMessage, FromWebviewMessage, OutputChangedMessage } from './webview-communication';
+import {
+    CellOutputChange, CellsChangedMessage, CellsMoved, CellsSpliced,
+    ChangePreferredMimetypeMessage, FromWebviewMessage, Output, OutputChangedMessage
+} from './webview-communication';
 import { Disposable, DisposableCollection, Emitter, QuickPickService, nls } from '@theia/core';
 import { NotebookModel } from '@theia/notebook/lib/browser/view-model/notebook-model';
 import { NotebookOptionsService, NotebookOutputOptions } from '@theia/notebook/lib/browser/service/notebook-options';
 import { NotebookCellModel } from '@theia/notebook/lib/browser/view-model/notebook-cell-model';
-import { NotebookCellsChangeType } from '@theia/notebook/lib/common';
+import { CellOutput, NotebookCellsChangeType } from '@theia/notebook/lib/common';
 import { NotebookCellOutputModel } from '@theia/notebook/lib/browser/view-model/notebook-cell-output-model';
 
 export const AdditionalNotebookCellOutputCss = Symbol('AdditionalNotebookCellOutputCss');
@@ -363,12 +366,8 @@ export class CellOutputWebviewImpl implements CellOutputWebview, Disposable {
                 .filter(update => visibleCellHandleLookup.has(update.cellHandle))
                 .map(update => ({
                     cellHandle: update.cellHandle,
-                    newOutputs: update.newOutputs.map(output => ({
-                        id: output.outputId,
-                        items: output.outputs.map(item => ({ mime: item.mime, data: item.data.buffer })),
-                        metadata: output.metadata
-                    })),
-                    start: this.toVisibleCellIndex(update.start, visibleCells),
+                    newOutputs: this.mapCellOutputsToWebviewOutput(update.newOutputs),
+                    start: update.start,
                     deleteCount: update.deleteCount
                 }))
         };
@@ -380,6 +379,7 @@ export class CellOutputWebviewImpl implements CellOutputWebview, Disposable {
 
     cellsChanged(cellEvents: NotebookContentChangedEvent[]): void {
         const changes: Array<CellsMoved | CellsSpliced> = [];
+        const outputChanges: CellOutputChange[] = [];
 
         const visibleCells = this.notebook.getVisibleCells();
         const visibleCellLookup = new Set(visibleCells);
@@ -397,16 +397,39 @@ export class CellOutputWebviewImpl implements CellOutputWebview, Disposable {
                     deleteCount: change.deleteCount,
                     newCells: change.newItems.filter(cell => visibleCellLookup.has(cell as NotebookCellModel)).map(cell => cell.handle)
                 } as CellsSpliced)));
+                outputChanges.push(...event.changes
+                    .flatMap(change => change.newItems)
+                    .filter(cell => visibleCellLookup.has(cell as NotebookCellModel) && cell.outputs.length)
+                    .map(newCell => ({
+                        start: 0,
+                        deleteCount: 0,
+                        cellHandle: newCell.handle,
+                        newOutputs: this.mapCellOutputsToWebviewOutput(newCell.outputs)
+                    })));
             }
         }
 
         this.webviewWidget.sendMessage({
             type: 'cellsChanged',
             changes: changes.filter(e => e)
         } as CellsChangedMessage);
+        if (outputChanges.length > 0) {
+            this.webviewWidget.sendMessage({
+                type: 'outputChanged',
+                changes: outputChanges
+            });
+        }
+    }
+
+    protected mapCellOutputsToWebviewOutput(outputs: CellOutput[]): Output[] {
+        return outputs.map(output => ({
+            id: output.outputId,
+            items: output.outputs.map(item => ({ mime: item.mime, data: item.data.buffer })),
+            metadata: output.metadata
+        }));
     }
 
-    toVisibleCellIndex(index: number, visibleCells: Array<NotebookCellModel>): number {
+    protected toVisibleCellIndex(index: number, visibleCells: Array<NotebookCellModel>): number {
         return visibleCells.indexOf(this.notebook.cells[index]);
     }